agents 0.0.0-b4d83b9 → 0.0.0-b507395

package/README.md

@@ -549,10 +549,12 @@ export class MyMCP extends McpAgent<Env, State, {}> {
       };
     });
 
-    this.server.tool(
+    this.server.registerTool(
       "add",
-      "Add to the counter, stored in the MCP",
-      { a: z.number() },
+      {
+        description: "Add to the counter, stored in the MCP",
+        inputSchema: { a: z.number() }
+      },
       async ({ a }) => {
         this.setState({ ...this.state, counter: this.state.counter + a });
 

package/dist/ai-chat-agent.d.ts

@@ -1,21 +1,13 @@
-import { UIMessage, StreamTextOnFinishCallback, ToolSet } from "ai";
-import { Agent, AgentContext } from "./index.js";
-import { Connection, WSMessage } from "partyserver";
-import "cloudflare:workers";
-import "@modelcontextprotocol/sdk/client/index.js";
-import "@modelcontextprotocol/sdk/types.js";
-import "./client-DVoPb3-C.js";
-import "zod";
-import "@modelcontextprotocol/sdk/shared/protocol.js";
-import "./mcp-BH1fJeiU.js";
-import "@modelcontextprotocol/sdk/client/sse.js";
-import "@modelcontextprotocol/sdk/client/streamableHttp.js";
-import "./mcp/do-oauth-client-provider.js";
-import "@modelcontextprotocol/sdk/client/auth.js";
-import "@modelcontextprotocol/sdk/shared/auth.js";
-import "./observability/index.js";
-import "./ai-types.js";
+import "./context-_sPQqJWv.js";
+import "./client-DfIOsabL.js";
+import "./mcp-CzbSsLfc.js";
+import "./do-oauth-client-provider-B-ryFIPr.js";
+import "./index-CyDpAVHZ.js";
+import "./ai-types-D_hTbf25.js";
+import { n as AgentContext, t as Agent } from "./index-DPJ32qQn.js";
+import { StreamTextOnFinishCallback, ToolSet, UIMessage } from "ai";
 
+//#region src/ai-chat-agent.d.ts
 /**
  * Extension of Agent with built-in chat capabilities
  * @template Env Environment type containing bindings
@@ -29,11 +21,92 @@ declare class AIChatAgent<Env = unknown, State = unknown> extends Agent<
    * useful to propagate request cancellation signals for any external calls made by the agent
    */
   private _chatMessageAbortControllers;
+  /**
+   * Currently active stream ID for resumable streaming.
+   * Stored in memory for quick access; persisted in stream_metadata table.
+   * @internal Protected for testing purposes.
+   */
+  protected _activeStreamId: string | null;
+  /**
+   * Request ID associated with the active stream.
+   * @internal Protected for testing purposes.
+   */
+  protected _activeRequestId: string | null;
+  /**
+   * Current chunk index for the active stream
+   */
+  private _streamChunkIndex;
+  /**
+   * Buffer for stream chunks pending write to SQLite.
+   * Chunks are batched and flushed when buffer reaches CHUNK_BUFFER_SIZE.
+   */
+  private _chunkBuffer;
+  /**
+   * Lock to prevent concurrent flush operations
+   */
+  private _isFlushingChunks;
+  /**
+   * Timestamp of the last cleanup operation for old streams
+   */
+  private _lastCleanupTime;
   /** Array of chat messages for the current conversation */
   messages: UIMessage[];
   constructor(ctx: AgentContext, env: Env);
+  /**
+   * Restore active stream state if the agent was restarted during streaming.
+   * Called during construction to recover any interrupted streams.
+   * Validates stream freshness to avoid sending stale resume notifications.
+   * @internal Protected for testing purposes.
+   */
+  protected _restoreActiveStream(): void;
+  /**
+   * Notify a connection about an active stream that can be resumed.
+   * The client should respond with CF_AGENT_STREAM_RESUME_ACK to receive chunks.
+   * Uses in-memory state for request ID - no extra DB lookup needed.
+   * @param connection - The WebSocket connection to notify
+   */
+  private _notifyStreamResuming;
+  /**
+   * Send stream chunks to a connection after receiving ACK.
+   * @param connection - The WebSocket connection
+   * @param streamId - The stream to replay
+   * @param requestId - The original request ID
+   */
+  private _sendStreamChunks;
+  /**
+   * Buffer a stream chunk for batch write to SQLite.
+   * @param streamId - The stream this chunk belongs to
+   * @param body - The serialized chunk body
+   * @internal Protected for testing purposes.
+   */
+  protected _storeStreamChunk(streamId: string, body: string): void;
+  /**
+   * Flush buffered chunks to SQLite in a single batch.
+   * Uses a lock to prevent concurrent flush operations.
+   * @internal Protected for testing purposes.
+   */
+  protected _flushChunkBuffer(): void;
+  /**
+   * Start tracking a new stream for resumable streaming.
+   * Creates metadata entry in SQLite and sets up tracking state.
+   * @param requestId - The unique ID of the chat request
+   * @returns The generated stream ID
+   * @internal Protected for testing purposes.
+   */
+  protected _startStream(requestId: string): string;
+  /**
+   * Mark a stream as completed and flush any pending chunks.
+   * @param streamId - The stream to mark as completed
+   * @internal Protected for testing purposes.
+   */
+  protected _completeStream(streamId: string): void;
+  /**
+   * Clean up old completed streams if enough time has passed since last cleanup.
+   * This prevents database growth while avoiding cleanup overhead on every stream completion.
+   */
+  private _maybeCleanupOldStreams;
   private _broadcastChatMessage;
-  onMessage(connection: Connection, message: WSMessage): Promise<void>;
+  private _loadMessagesFromDb;
   onRequest(request: Request): Promise<Response>;
   private _tryCatchChat;
   /**
@@ -58,6 +131,12 @@ declare class AIChatAgent<Env = unknown, State = unknown> extends Agent<
     excludeBroadcastIds?: string[]
   ): Promise<void>;
   private _reply;
+  /**
+   * Mark a stream as errored and clean up state.
+   * @param streamId - The stream to mark as errored
+   * @internal Protected for testing purposes.
+   */
+  protected _markStreamError(streamId: string): void;
   /**
    * For the given message id, look up its associated AbortController
    * If the AbortController does not exist, create and store one in memory
@@ -78,9 +157,10 @@ declare class AIChatAgent<Env = unknown, State = unknown> extends Agent<
    */
   private _destroyAbortControllers;
   /**
-   * When the DO is destroyed, cancel all pending requests
+   * When the DO is destroyed, cancel all pending requests and clean up resources
    */
   destroy(): Promise<void>;
 }
-
+//#endregion
 export { AIChatAgent };
+//# sourceMappingURL=ai-chat-agent.d.ts.map