feat: Add usageV2 support

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

@@ -46,6 +46,7 @@
   "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,37 @@
+import type { UsageV2Event } from "src/core/usageV2.js";
+
+/**
+ * Send events to Kafka.
+ * This method may throw. To call this non-blocking:
+ *
+ * ```ts
+ * sendUsageV2Events("production", events).catch(console.error)
+ * ```
+ *
+ * @param environment - The environment the service is running in.
+ * @param events - The events to send.
+ */
+export async function sendUsageV2Events(
+  environment: "development" | "production",
+  events: UsageV2Event[],
+): Promise<void> {
+  const baseUrl =
+    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",
+    },
+    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,110 @@
+import { randomUUID } from "node:crypto";
+import { checkServerIdentity } from "node:tls";
+import { Kafka, type Producer } from "kafkajs";
+import type { UsageV2Event } from "src/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(({ id, created_at, data, ...rest }) => {
+      return {
+        id: id ?? randomUUID(),
+        created_at: created_at ?? new Date(),
+        data: JSON.stringify(data),
+        ...rest,
+      };
+    });
+
+    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;
+    }
+  }
+}