Refactor ID generation across agents to use nanoid (same as base agent). improve stream manager with a destroy method for proper cleanup during agent destruction

Author: whoiskatrin

packages/agents/src/ai-chat-agent-http.ts

@@ -3,6 +3,7 @@ import type {
   StreamTextOnFinishCallback,
   ToolSet
 } from "ai";
+import { nanoid } from "nanoid";
 import { Agent, type AgentContext } from "./";
 import { autoTransformMessages } from "./ai-chat-v5-migration";
 import { ResumableStreamManager } from "./resumable-stream-manager";
@@ -178,9 +179,7 @@ export class AIHttpChatAgent<
     }
 
     // Generate or reuse stream ID
-    const streamId =
-      requestStreamId ||
-      `stream_${Date.now()}_${Math.random().toString(36).slice(2, 11)}`;
+    const streamId = requestStreamId || `stream_${nanoid()}`;
 
     // Delegate to stream manager
     return this._streamManager.startStream(
@@ -276,4 +275,13 @@ export class AIHttpChatAgent<
   async cleanupOldStreams(maxAgeHours = 24): Promise<void> {
     await this._streamManager.cleanupOldStreams(maxAgeHours);
   }
+
+  /**
+   * Override destroy to properly clean up resumable streaming
+   */
+  override async destroy(): Promise<void> {
+    // Clean up resumable streaming first
+    await this._streamManager.destroy();
+    await super.destroy();
+  }
 }

packages/agents/src/ai-chat-agent.ts

@@ -3,6 +3,7 @@ import type {
   StreamTextOnFinishCallback,
   ToolSet
 } from "ai";
+import { nanoid } from "nanoid";
 import { Agent, type AgentContext, type Connection, type WSMessage } from "./";
 import {
   MessageType,
@@ -400,7 +401,7 @@ export class AIChatAgent<Env = unknown, State = unknown> extends Agent<
         await this.persistMessages([
           ...this.messages,
           {
-            id: `assistant_${Date.now()}_${Math.random().toString(36).slice(2, 11)}`,
+            id: `assistant_${nanoid()}`,
             role: "assistant",
             parts: messageParts
           }

packages/agents/src/client.ts

@@ -3,6 +3,7 @@ import {
   PartySocket,
   type PartySocketOptions
 } from "partysocket";
+import { nanoid } from "nanoid";
 import type { RPCRequest, RPCResponse } from "./";
 import type {
   SerializableReturnValue,
@@ -181,7 +182,7 @@ export class AgentClient<State = unknown> extends PartySocket {
     streamOptions?: StreamOptions
   ): Promise<T> {
     return new Promise<T>((resolve, reject) => {
-      const id = Math.random().toString(36).slice(2);
+      const id = nanoid();
       this._pendingCalls.set(id, {
         reject,
         resolve: (value: unknown) => resolve(value as T),

packages/agents/src/mcp/index.ts

@@ -16,6 +16,7 @@ import type {
   ConnectionContext,
   WSMessage
 } from "../";
+import { nanoid } from "nanoid";
 import { Agent } from "../index";
 import type {
   MaybeConnectionTag,
@@ -303,7 +304,7 @@ export abstract class McpAgent<
     message: string;
     requestedSchema: unknown;
   }): Promise<ElicitResult> {
-    const requestId = `elicit_${Math.random().toString(36).substring(2, 11)}`;
+    const requestId = `elicit_${nanoid()}`;
 
     // Store pending request in durable storage
     await this.ctx.storage.put(`elicitation:${requestId}`, {

packages/agents/src/react.tsx

@@ -1,6 +1,7 @@
 import type { PartySocket } from "partysocket";
 import { usePartySocket } from "partysocket/react";
 import { useCallback, useRef } from "react";
+import { nanoid } from "nanoid";
 import type { Agent, MCPServersState, RPCRequest, RPCResponse } from "./";
 import type { StreamOptions } from "./client";
 import type { Method, RPCMethod } from "./serializable";
@@ -239,7 +240,7 @@ export function useAgent<State>(
       streamOptions?: StreamOptions
     ): Promise<T> => {
       return new Promise((resolve, reject) => {
-        const id = Math.random().toString(36).slice(2);
+        const id = nanoid();
         pendingCallsRef.current.set(id, {
           reject,
           resolve: resolve as (value: unknown) => void,

packages/agents/src/resumable-stream-manager.ts

@@ -3,13 +3,9 @@ import type {
   StreamTextOnFinishCallback,
   ToolSet
 } from "ai";
+import { nanoid } from "nanoid";
 import type { AgentContext } from "./";
 
-type SqlQueryFunction = <T = Record<string, string | number | boolean | null>>(
-  strings: TemplateStringsArray,
-  ...values: (string | number | boolean | null)[]
-) => T[];
-
 interface StreamStateRow {
   stream_id: string;
   seq: number;
@@ -53,9 +49,18 @@ export class ResumableStreamManager<Message extends ChatMessage = ChatMessage> {
   >;
 
   private ctx: AgentContext;
-  private sql: SqlQueryFunction;
-
-  constructor(ctx: AgentContext, sql: SqlQueryFunction) {
+  private sql: <T = Record<string, string | number | boolean | null>>(
+    strings: TemplateStringsArray,
+    ...values: (string | number | boolean | null)[]
+  ) => T[];
+
+  constructor(
+    ctx: AgentContext,
+    sql: <T = Record<string, string | number | boolean | null>>(
+      strings: TemplateStringsArray,
+      ...values: (string | number | boolean | null)[]
+    ) => T[]
+  ) {
     this.ctx = ctx;
     this.sql = sql;
     this._activeStreams = new Map();
@@ -264,6 +269,19 @@ export class ResumableStreamManager<Message extends ChatMessage = ChatMessage> {
     this._activeStreams.clear();
   }
 
+  /**
+   * Destroy all resumable streaming
+   * Should be called during Agent destruction
+   */
+  async destroy(): Promise<void> {
+    // Clear in-memory state first
+    this._activeStreams.clear();
+
+    // Drop all tables
+    this.sql`DROP TABLE IF EXISTS cf_ai_http_chat_streams`;
+    this.sql`DROP TABLE IF EXISTS cf_ai_http_chat_chunks`;
+  }
+
   /**
    * Clean up old completed streams (call periodically)
    */
@@ -369,7 +387,7 @@ export class ResumableStreamManager<Message extends ChatMessage = ChatMessage> {
     if (!streamState) return;
 
     let assistantMessageText = "";
-    const assistantMessageId = `assistant_${Date.now()}_${Math.random().toString(36).slice(2, 11)}`;
+    const assistantMessageId = `assistant_${nanoid()}`;
     let buffer = "";
 
     try {