[service-utils] feat: Add usageV2 support (#6021)

Author: arcoraven

.changeset/neat-ads-build.md

@@ -0,0 +1,5 @@
+---
+"@thirdweb-dev/service-utils": minor
+---
+
+feat: Add usageV2 support

packages/service-utils/package.json

@@ -26,12 +26,8 @@
   },
   "typesVersions": {
     "*": {
-      "node": [
-        "./dist/types/node/index.d.ts"
-      ],
-      "cf-worker": [
-        "./dist/types/cf-worker/index.d.ts"
-      ]
+      "node": ["./dist/types/node/index.d.ts"],
+      "cf-worker": ["./dist/types/cf-worker/index.d.ts"]
     }
   },
   "repository": "https://github.com/thirdweb-dev/js/tree/main/packages/pay",
@@ -40,12 +36,11 @@
     "url": "https://github.com/thirdweb-dev/js/issues"
   },
   "author": "thirdweb eng <eng@thirdweb.com>",
-  "files": [
-    "dist/"
-  ],
+  "files": ["dist/"],
   "sideEffects": false,
   "dependencies": {
     "aws4fetch": "1.0.20",
+    "kafkajs": "2.2.4",
     "zod": "3.24.1"
   },
   "devDependencies": {

packages/service-utils/src/cf-worker/usageV2.ts

@@ -0,0 +1,45 @@
+import type { UsageV2Event } from "../core/usageV2.js";
+
+/**
+ * Send events to Kafka.
+ * This method may throw. To call this non-blocking:
+ *
+ * ```ts
+ * void sendUsageV2Events(events, {
+ *   environment: "production",
+ *   serviceKey: "..."
+ * }).catch(console.error)
+ * ```
+ *
+ * @param events - The events to send.
+ * @param options.environment - The environment the service is running in.
+ * @param options.serviceKey - The service key required for authentication.
+ */
+export async function sendUsageV2Events(
+  events: UsageV2Event[],
+  options: {
+    environment: "development" | "production";
+    serviceKey: string;
+  },
+): Promise<void> {
+  const baseUrl =
+    options.environment === "production"
+      ? "https://u.thirdweb.com"
+      : "https://u.thirdweb-dev.com";
+
+  const resp = await fetch(`${baseUrl}/usage-v2/raw-events`, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+      "x-service-api-key": options.serviceKey,
+    },
+    body: JSON.stringify({ events }),
+  });
+
+  if (!resp.ok) {
+    throw new Error(
+      `[UsageV2] unexpected response ${resp.status}: ${await resp.text()}`,
+    );
+  }
+  resp.body?.cancel();
+}

packages/service-utils/src/core/usageV2.ts

@@ -0,0 +1,56 @@
+export interface UsageV2Event {
+  /**
+   * A unique identifier for the event. Defaults to a random UUID.
+   * Useful if your service retries sending events.
+   */
+  id?: `${string}-${string}-${string}-${string}-${string}`;
+  /**
+   * The event timestamp. Defaults to now().
+   */
+  created_at?: Date;
+  /**
+   * The source of the event. Example: "storage"
+   */
+  source: string;
+  /**
+   * The action of the event. Example: "upload"
+   */
+  action: string;
+  /**
+   * The team ID.
+   */
+  team_id: string;
+  /**
+   * The client ID, if available.
+   */
+  client_id?: string;
+  /**
+   * The SDK name, if available.
+   */
+  sdk_name?: string;
+  /**
+   * The SDK platform, if available.
+   */
+  sdk_platform?: string;
+  /**
+   * The SDK version, if available.
+   */
+  sdk_version?: string;
+  /**
+   * The SDK OS, if available.
+   */
+  sdk_os?: string;
+  /**
+   * The product name, if available.
+   */
+  product_name?: string;
+  /**
+   * The product version, if available.
+   */
+  product_version?: string;
+  /**
+   * An object of service-specific data. Example: "file_size_bytes"
+   * It is safe to pass any new JSON-serializable data here before updating the usageV2 schema.
+   */
+  data: Record<string, unknown>;
+}

packages/service-utils/src/node/usageV2.ts

@@ -0,0 +1,122 @@
+import { randomUUID } from "node:crypto";
+import { checkServerIdentity } from "node:tls";
+import { Kafka, type Producer } from "kafkajs";
+import type { UsageV2Event } from "../core/usageV2.js";
+
+const TOPIC_USAGE_V2 = "usage_v2.raw_events";
+
+/**
+ * Creates a UsageV2Producer which opens a persistent TCP connection.
+ * This class is thread-safe so your service should re-use one instance.
+ *
+ * Example:
+ * ```ts
+ * usageV2 = new UsageV2Producer(..)
+ * await usageV2.init()
+ * await usageV2.sendEvents(events)
+ * // Non-blocking:
+ * // void usageV2.sendEvents(events).catch(console.error)
+ * ```
+ */
+export class UsageV2Producer {
+  private kafka: Kafka;
+  private producer: Producer | null = null;
+
+  constructor(config: {
+    /**
+     * A descriptive name for your service. Example: "storage-server"
+     */
+    producerName: string;
+    /**
+     * The environment the service is running in.
+     */
+    environment: "development" | "production";
+
+    username: string;
+    password: string;
+  }) {
+    this.kafka = new Kafka({
+      clientId: `${config.producerName}-${config.environment}`,
+      brokers:
+        config.environment === "production"
+          ? ["warpstream.thirdweb.xyz:9092"]
+          : ["warpstream-dev.thirdweb.xyz:9092"],
+      ssl: {
+        checkServerIdentity(hostname, cert) {
+          return checkServerIdentity(hostname.toLowerCase(), cert);
+        },
+      },
+      sasl: {
+        mechanism: "plain",
+        username: config.username,
+        password: config.password,
+      },
+    });
+  }
+
+  /**
+   * Connect the producer.
+   * This must be called before calling `sendEvents()`.
+   */
+  async init() {
+    this.producer = this.kafka.producer({
+      allowAutoTopicCreation: false,
+    });
+    await this.producer.connect();
+  }
+
+  /**
+   * Send usageV2 events.
+   * This method may throw. To call this non-blocking:
+   *
+   * ```ts
+   * usageV2 = new UsageV2Producer(...)
+   * void usageV2.sendEvents(events).catch(console.error)
+   *
+   * @param events - The events to send.
+   */
+  async sendEvents(events: UsageV2Event[]): Promise<void> {
+    if (!this.producer) {
+      throw new Error("Producer not initialized. Call `init()` first.");
+    }
+
+    const parsedEvents = events.map((event) => {
+      return {
+        id: event.id ?? randomUUID(),
+        created_at: event.created_at ?? new Date(),
+        source: event.source,
+        action: event.action,
+        // Remove the "team_" prefix, if any.
+        team_id: event.team_id.startsWith("team_")
+          ? event.team_id.slice(5)
+          : event.team_id,
+        client_id: event.client_id,
+        sdk_name: event.sdk_name,
+        sdk_platform: event.sdk_platform,
+        sdk_version: event.sdk_version,
+        sdk_os: event.sdk_os,
+        product_name: event.product_name,
+        product_version: event.product_version,
+        data: JSON.stringify(event.data),
+      };
+    });
+
+    await this.producer.send({
+      topic: TOPIC_USAGE_V2,
+      messages: parsedEvents.map((event) => ({
+        value: JSON.stringify(event),
+      })),
+    });
+  }
+
+  /**
+   * Disconnects UsageV2Producer.
+   * Useful when shutting down the service to flush in-flight events.
+   */
+  async disconnect() {
+    if (this.producer) {
+      await this.producer.disconnect();
+      this.producer = null;
+    }
+  }
+}