agents 0.0.0-5ebaeb2 → 0.0.0-603b825

package/README.md

@@ -342,16 +342,114 @@ export class DialogueAgent extends AIChatAgent {
     //   messages: this.messages,
     // });
     //
-    // // Optional: you can call onFinish here for custom side effects. Message
-    // // persistence is still handled automatically by AIChatAgent.
-    // await onFinish?.(result);
-    // return new Response(result.text, {
-    //   headers: { 'Content-Type': 'text/plain' }
+    // // For non-streaming with metadata, use toUIMessage:
+    // const message = result.toUIMessage({
+    //   metadata: {
+    //     model: 'gpt-4o',
+    //     totalTokens: result.usage?.totalTokens,
+    //   }
     // });
+    //
+    // return new Response(JSON.stringify(message), {
+    //   headers: { 'Content-Type': 'application/json' }
+    // });
+  }
+}
+```
+
+#### Metadata Support
+
+The AI SDK provides native support for message metadata through the `messageMetadata` callback. This allows you to attach custom information to messages at the message level.
+
+##### AIChatAgent Integration
+
+In the context of `AIChatAgent`, you can use metadata like this:
+
+```typescript
+import { AIChatAgent } from "agents/ai-chat-agent";
+import { streamText } from "ai";
+import { openai } from "@ai-sdk/openai";
+
+export class MyAgent extends AIChatAgent<Env> {
+  async onChatMessage(onFinish) {
+    const startTime = Date.now();
+
+    const result = streamText({
+      model: openai("gpt-4o"),
+      messages: this.messages,
+      onFinish
+    });
+
+    return result.toUIMessageStreamResponse({
+      messageMetadata: ({ part }) => {
+        if (part.type === "start") {
+          return {
+            model: "gpt-4o",
+            createdAt: Date.now(),
+            messageCount: this.messages.length
+          };
+        }
+        if (part.type === "finish") {
+          return {
+            responseTime: Date.now() - startTime,
+            totalTokens: part.totalUsage?.totalTokens
+          };
+        }
+      }
+    });
   }
 }
 ```
 
+##### Accessing Metadata on the Client
+
+Access metadata through the `message.metadata` property:
+
+```typescript
+'use client';
+
+import { useChat } from '@ai-sdk/react';
+import { DefaultChatTransport } from 'ai';
+import type { MyUIMessage } from '@/types';
+
+export default function Chat() {
+  const { messages } = useChat<MyUIMessage>({
+    transport: new DefaultChatTransport({
+      api: '/api/chat',
+    }),
+  });
+
+  return (
+    <div>
+      {messages.map(message => (
+        <div key={message.id}>
+          <div>
+            {message.role === 'user' ? 'User: ' : 'AI: '}
+            {message.metadata?.createdAt && (
+              <span className="text-sm text-gray-500">
+                {new Date(message.metadata.createdAt).toLocaleTimeString()}
+              </span>
+            )}
+          </div>
+          {/* Render message content */}
+          {message.parts.map((part, index) =>
+            part.type === 'text' ? <div key={index}>{part.text}</div> : null,
+          )}
+          {/* Display additional metadata */}
+          {message.metadata?.totalTokens && (
+            <div className="text-xs text-gray-400">
+              {message.metadata.totalTokens} tokens
+            </div>
+          )}
+        </div>
+      ))}
+    </div>
+  );
+}
+```
+
+For more details, see the [AI SDK Message Metadata documentation](https://ai-sdk.dev/docs/ai-sdk-ui/message-metadata).
+
 #### Creating the Interface
 
 Connect with your agent through a React interface:
@@ -451,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,20 +1,12 @@
-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-CcIORE73.js";
-import "zod";
-import "@modelcontextprotocol/sdk/shared/protocol.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 "./client-ctTw3KHG.js";
+import "./mcp-BwPscEiF.js";
+import "./do-oauth-client-provider-UhQDpDb8.js";
+import "./index-BUle9RiP.js";
+import "./ai-types-81H_-Uxh.js";
+import { n as AgentContext, t as Agent } from "./index-D2dnUH0r.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
@@ -28,11 +20,85 @@ 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.
+   */
+  private _activeStreamId;
+  /**
+   * Request ID associated with the active stream.
+   */
+  private _activeRequestId;
+  /**
+   * 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.
+   */
+  private _restoreActiveStream;
+  /**
+   * 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
+   */
+  private _storeStreamChunk;
+  /**
+   * Flush buffered chunks to SQLite in a single batch.
+   * Uses a lock to prevent concurrent flush operations.
+   */
+  private _flushChunkBuffer;
+  /**
+   * 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
+   */
+  private _startStream;
+  /**
+   * Mark a stream as completed and flush any pending chunks.
+   * @param streamId - The stream to mark as completed
+   */
+  private _completeStream;
+  /**
+   * 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;
   /**
@@ -57,6 +123,11 @@ 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
+   */
+  private _markStreamError;
   /**
    * For the given message id, look up its associated AbortController
    * If the AbortController does not exist, create and store one in memory
@@ -77,9 +148,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