agents 0.0.0-5ebaeb2 → 0.0.0-629836

package/dist/ai-chat-agent.js

@@ -3,13 +3,19 @@ import {
 } from "./chunk-UJVEAURM.js";
 import {
   Agent
-} from "./chunk-IJPBZOSS.js";
-import "./chunk-VYENMKFS.js";
-import "./chunk-LL2AFX7V.js";
+} from "./chunk-XGMKNUJA.js";
+import "./chunk-ZMMHNOMZ.js";
+import "./chunk-C2OEBJZ2.js";
 import "./chunk-QEVM4BVL.js";
-import "./chunk-AVYJQSLW.js";
+import "./chunk-BER7KXUJ.js";
+import "./chunk-PR4QN5HX.js";
 
 // src/ai-chat-agent.ts
+import {
+  getToolName,
+  isToolUIPart,
+  parsePartialJson
+} from "ai";
 var decoder = new TextDecoder();
 var AIChatAgent = class extends Agent {
   constructor(ctx, env) {
@@ -19,15 +25,24 @@ var AIChatAgent = class extends Agent {
       message text not null,
       created_at datetime default current_timestamp
     )`;
-    const rawMessages = (this.sql`select * from cf_ai_chat_agent_messages` || []).map((row) => {
-      return JSON.parse(row.message);
-    });
+    const rawMessages = this._loadMessagesFromDb();
     this.messages = autoTransformMessages(rawMessages);
     this._chatMessageAbortControllers = /* @__PURE__ */ new Map();
   }
   _broadcastChatMessage(message, exclude) {
     this.broadcast(JSON.stringify(message), exclude);
   }
+  _loadMessagesFromDb() {
+    const rows = this.sql`select * from cf_ai_chat_agent_messages order by created_at` || [];
+    return rows.map((row) => {
+      try {
+        return JSON.parse(row.message);
+      } catch (error) {
+        console.error(`Failed to parse message ${row.id}:`, error);
+        return null;
+      }
+    }).filter((msg) => msg !== null);
+  }
   async onMessage(connection, message) {
     if (typeof message === "string") {
       let data;
@@ -127,9 +142,7 @@ var AIChatAgent = class extends Agent {
     return this._tryCatchChat(() => {
       const url = new URL(request.url);
       if (url.pathname.endsWith("/get-messages")) {
-        const messages = (this.sql`select * from cf_ai_chat_agent_messages` || []).map((row) => {
-          return JSON.parse(row.message);
-        });
+        const messages = this._loadMessagesFromDb();
         return Response.json(messages);
       }
       return super.onRequest(request);
@@ -159,13 +172,22 @@ var AIChatAgent = class extends Agent {
    */
   async saveMessages(messages) {
     await this.persistMessages(messages);
+    await this._tryCatchChat(async () => {
+      const response = await this.onChatMessage(() => {
+      });
+      if (response) this._reply(crypto.randomUUID(), response);
+    });
   }
   async persistMessages(messages, excludeBroadcastIds = []) {
-    this.sql`delete from cf_ai_chat_agent_messages`;
     for (const message of messages) {
-      this.sql`insert into cf_ai_chat_agent_messages (id, message) values (${message.id},${JSON.stringify(message)})`;
+      this.sql`
+        insert into cf_ai_chat_agent_messages (id, message)
+        values (${message.id}, ${JSON.stringify(message)})
+        on conflict(id) do update set message = excluded.message
+      `;
     }
-    this.messages = messages;
+    const persisted = this._loadMessagesFromDb();
+    this.messages = autoTransformMessages(persisted);
     this._broadcastChatMessage(
       {
         messages,
@@ -186,8 +208,84 @@ var AIChatAgent = class extends Agent {
         return;
       }
       const reader = response.body.getReader();
-      let fullResponseText = "";
-      const toolCalls = /* @__PURE__ */ new Map();
+      const message = {
+        id: `assistant_${Date.now()}_${Math.random().toString(36).slice(2, 11)}`,
+        // default
+        role: "assistant",
+        parts: []
+      };
+      let activeTextParts = {};
+      let activeReasoningParts = {};
+      const partialToolCalls = {};
+      function updateDynamicToolPart(options) {
+        const part = message.parts.find(
+          (part2) => part2.type === "dynamic-tool" && part2.toolCallId === options.toolCallId
+        );
+        const anyOptions = options;
+        const anyPart = part;
+        if (part != null) {
+          part.state = options.state;
+          anyPart.toolName = options.toolName;
+          anyPart.input = anyOptions.input;
+          anyPart.output = anyOptions.output;
+          anyPart.errorText = anyOptions.errorText;
+          anyPart.rawInput = anyOptions.rawInput ?? anyPart.rawInput;
+          anyPart.preliminary = anyOptions.preliminary;
+          if (anyOptions.providerMetadata != null && part.state === "input-available") {
+            part.callProviderMetadata = anyOptions.providerMetadata;
+          }
+        } else {
+          message.parts.push({
+            type: "dynamic-tool",
+            toolName: options.toolName,
+            toolCallId: options.toolCallId,
+            state: options.state,
+            input: anyOptions.input,
+            output: anyOptions.output,
+            errorText: anyOptions.errorText,
+            preliminary: anyOptions.preliminary,
+            ...anyOptions.providerMetadata != null ? { callProviderMetadata: anyOptions.providerMetadata } : {}
+          });
+        }
+      }
+      function updateToolPart(options) {
+        const part = message.parts.find(
+          (part2) => isToolUIPart(part2) && part2.toolCallId === options.toolCallId
+        );
+        const anyOptions = options;
+        const anyPart = part;
+        if (part != null) {
+          part.state = options.state;
+          anyPart.input = anyOptions.input;
+          anyPart.output = anyOptions.output;
+          anyPart.errorText = anyOptions.errorText;
+          anyPart.rawInput = anyOptions.rawInput;
+          anyPart.preliminary = anyOptions.preliminary;
+          anyPart.providerExecuted = anyOptions.providerExecuted ?? part.providerExecuted;
+          if (anyOptions.providerMetadata != null && part.state === "input-available") {
+            part.callProviderMetadata = anyOptions.providerMetadata;
+          }
+        } else {
+          message.parts.push({
+            type: `tool-${options.toolName}`,
+            toolCallId: options.toolCallId,
+            state: options.state,
+            input: anyOptions.input,
+            output: anyOptions.output,
+            rawInput: anyOptions.rawInput,
+            errorText: anyOptions.errorText,
+            providerExecuted: anyOptions.providerExecuted,
+            preliminary: anyOptions.preliminary,
+            ...anyOptions.providerMetadata != null ? { callProviderMetadata: anyOptions.providerMetadata } : {}
+          });
+        }
+      }
+      async function updateMessageMetadata(metadata) {
+        if (metadata != null) {
+          const mergedMetadata = message.metadata != null ? { ...message.metadata, ...metadata } : metadata;
+          message.metadata = mergedMetadata;
+        }
+      }
       try {
         while (true) {
           const { done, value } = await reader.read();
@@ -210,33 +308,270 @@ var AIChatAgent = class extends Agent {
                 try {
                   const data = JSON.parse(line.slice(6));
                   switch (data.type) {
-                    // SSE event signaling the tool input is ready. We track by
-                    // `toolCallId` so we can persist it as a tool part in the message.
-                    case "tool-input-available": {
-                      const { toolCallId, toolName, input } = data;
-                      toolCalls.set(toolCallId, {
-                        toolCallId,
-                        toolName,
-                        input,
-                        type: toolName ? `tool-${toolName}` : "dynamic-tool",
-                        state: "input-available"
+                    case "text-start": {
+                      const textPart = {
+                        type: "text",
+                        text: "",
+                        providerMetadata: data.providerMetadata,
+                        state: "streaming"
+                      };
+                      activeTextParts[data.id] = textPart;
+                      message.parts.push(textPart);
+                      break;
+                    }
+                    case "text-delta": {
+                      const textPart = activeTextParts[data.id];
+                      textPart.text += data.delta;
+                      textPart.providerMetadata = data.providerMetadata ?? textPart.providerMetadata;
+                      break;
+                    }
+                    case "text-end": {
+                      const textPart = activeTextParts[data.id];
+                      textPart.state = "done";
+                      textPart.providerMetadata = data.providerMetadata ?? textPart.providerMetadata;
+                      delete activeTextParts[data.id];
+                      break;
+                    }
+                    case "reasoning-start": {
+                      const reasoningPart = {
+                        type: "reasoning",
+                        text: "",
+                        providerMetadata: data.providerMetadata,
+                        state: "streaming"
+                      };
+                      activeReasoningParts[data.id] = reasoningPart;
+                      message.parts.push(reasoningPart);
+                      break;
+                    }
+                    case "reasoning-delta": {
+                      const reasoningPart = activeReasoningParts[data.id];
+                      reasoningPart.text += data.delta;
+                      reasoningPart.providerMetadata = data.providerMetadata ?? reasoningPart.providerMetadata;
+                      break;
+                    }
+                    case "reasoning-end": {
+                      const reasoningPart = activeReasoningParts[data.id];
+                      reasoningPart.providerMetadata = data.providerMetadata ?? reasoningPart.providerMetadata;
+                      reasoningPart.state = "done";
+                      delete activeReasoningParts[data.id];
+                      break;
+                    }
+                    case "file": {
+                      message.parts.push({
+                        type: "file",
+                        mediaType: data.mediaType,
+                        url: data.url
+                      });
+                      break;
+                    }
+                    case "source-url": {
+                      message.parts.push({
+                        type: "source-url",
+                        sourceId: data.sourceId,
+                        url: data.url,
+                        title: data.title,
+                        providerMetadata: data.providerMetadata
+                      });
+                      break;
+                    }
+                    case "source-document": {
+                      message.parts.push({
+                        type: "source-document",
+                        sourceId: data.sourceId,
+                        mediaType: data.mediaType,
+                        title: data.title,
+                        filename: data.filename,
+                        providerMetadata: data.providerMetadata
                       });
                       break;
                     }
-                    // SSE event signaling the tool output is ready. We should've
-                    // already received the input in a previous event so an entry
-                    // with `toolCallId` should already be present
+                    case "tool-input-start": {
+                      const toolInvocations = message.parts.filter(isToolUIPart);
+                      partialToolCalls[data.toolCallId] = {
+                        text: "",
+                        toolName: data.toolName,
+                        index: toolInvocations.length,
+                        dynamic: data.dynamic
+                      };
+                      if (data.dynamic) {
+                        updateDynamicToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: data.toolName,
+                          state: "input-streaming",
+                          input: void 0
+                        });
+                      } else {
+                        updateToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: data.toolName,
+                          state: "input-streaming",
+                          input: void 0
+                        });
+                      }
+                      break;
+                    }
+                    case "tool-input-delta": {
+                      const partialToolCall = partialToolCalls[data.toolCallId];
+                      partialToolCall.text += data.inputTextDelta;
+                      const { value: partialArgs } = await parsePartialJson(
+                        partialToolCall.text
+                      );
+                      if (partialToolCall.dynamic) {
+                        updateDynamicToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: partialToolCall.toolName,
+                          state: "input-streaming",
+                          input: partialArgs
+                        });
+                      } else {
+                        updateToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: partialToolCall.toolName,
+                          state: "input-streaming",
+                          input: partialArgs
+                        });
+                      }
+                      break;
+                    }
+                    case "tool-input-available": {
+                      if (data.dynamic) {
+                        updateDynamicToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: data.toolName,
+                          state: "input-available",
+                          input: data.input,
+                          providerMetadata: data.providerMetadata
+                        });
+                      } else {
+                        updateToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: data.toolName,
+                          state: "input-available",
+                          input: data.input,
+                          providerExecuted: data.providerExecuted,
+                          providerMetadata: data.providerMetadata
+                        });
+                      }
+                      break;
+                    }
+                    case "tool-input-error": {
+                      if (data.dynamic) {
+                        updateDynamicToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: data.toolName,
+                          state: "output-error",
+                          input: data.input,
+                          errorText: data.errorText,
+                          providerMetadata: data.providerMetadata
+                        });
+                      } else {
+                        updateToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: data.toolName,
+                          state: "output-error",
+                          input: void 0,
+                          rawInput: data.input,
+                          errorText: data.errorText,
+                          providerExecuted: data.providerExecuted,
+                          providerMetadata: data.providerMetadata
+                        });
+                      }
+                      break;
+                    }
                     case "tool-output-available": {
-                      const { toolCallId, output, isError, errorText } = data;
-                      const toolPart = toolCalls.get(toolCallId);
-                      if (toolPart)
-                        toolCalls.set(toolCallId, {
-                          ...toolPart,
-                          output,
-                          isError,
-                          errorText,
-                          state: "output-available"
+                      if (data.dynamic) {
+                        const toolInvocations = message.parts.filter(
+                          (part) => part.type === "dynamic-tool"
+                        );
+                        const toolInvocation = toolInvocations.find(
+                          (invocation) => invocation.toolCallId === data.toolCallId
+                        );
+                        if (!toolInvocation)
+                          throw new Error("Tool invocation not found");
+                        updateDynamicToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: toolInvocation.toolName,
+                          state: "output-available",
+                          input: toolInvocation.input,
+                          output: data.output,
+                          preliminary: data.preliminary
+                        });
+                      } else {
+                        const toolInvocations = message.parts.filter(isToolUIPart);
+                        const toolInvocation = toolInvocations.find(
+                          (invocation) => invocation.toolCallId === data.toolCallId
+                        );
+                        if (!toolInvocation)
+                          throw new Error("Tool invocation not found");
+                        updateToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: getToolName(toolInvocation),
+                          state: "output-available",
+                          input: toolInvocation.input,
+                          output: data.output,
+                          providerExecuted: data.providerExecuted,
+                          preliminary: data.preliminary
+                        });
+                      }
+                      break;
+                    }
+                    case "tool-output-error": {
+                      if (data.dynamic) {
+                        const toolInvocations = message.parts.filter(
+                          (part) => part.type === "dynamic-tool"
+                        );
+                        const toolInvocation = toolInvocations.find(
+                          (invocation) => invocation.toolCallId === data.toolCallId
+                        );
+                        if (!toolInvocation)
+                          throw new Error("Tool invocation not found");
+                        updateDynamicToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: toolInvocation.toolName,
+                          state: "output-error",
+                          input: toolInvocation.input,
+                          errorText: data.errorText
+                        });
+                      } else {
+                        const toolInvocations = message.parts.filter(isToolUIPart);
+                        const toolInvocation = toolInvocations.find(
+                          (invocation) => invocation.toolCallId === data.toolCallId
+                        );
+                        if (!toolInvocation)
+                          throw new Error("Tool invocation not found");
+                        updateToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: getToolName(toolInvocation),
+                          state: "output-error",
+                          input: toolInvocation.input,
+                          rawInput: "rawInput" in toolInvocation ? toolInvocation.rawInput : void 0,
+                          errorText: data.errorText
                         });
+                      }
+                      break;
+                    }
+                    case "start-step": {
+                      message.parts.push({ type: "step-start" });
+                      break;
+                    }
+                    case "finish-step": {
+                      activeTextParts = {};
+                      activeReasoningParts = {};
+                      break;
+                    }
+                    case "start": {
+                      if (data.messageId != null) {
+                        message.id = data.messageId;
+                      }
+                      await updateMessageMetadata(data.messageMetadata);
+                      break;
+                    }
+                    case "finish": {
+                      await updateMessageMetadata(data.messageMetadata);
+                      break;
+                    }
+                    case "message-metadata": {
+                      await updateMessageMetadata(data.messageMetadata);
                       break;
                     }
                     case "error": {
@@ -247,10 +582,6 @@ var AIChatAgent = class extends Agent {
                         id,
                         type: "cf_agent_use_chat_response" /* CF_AGENT_USE_CHAT_RESPONSE */
                       });
-                      return;
-                    }
-                    case "text-delta": {
-                      if (data.delta) fullResponseText += data.delta;
                       break;
                     }
                   }
@@ -266,7 +597,7 @@ var AIChatAgent = class extends Agent {
             }
           } else {
             if (chunk.length > 0) {
-              fullResponseText += chunk;
+              message.parts.push({ type: "text", text: chunk });
               this._broadcastChatMessage({
                 body: JSON.stringify({ type: "text-delta", delta: chunk }),
                 done: false,
@@ -279,22 +610,8 @@ var AIChatAgent = class extends Agent {
       } finally {
         reader.releaseLock();
       }
-      const messageParts = [];
-      Array.from(toolCalls.values()).forEach((t) => {
-        messageParts.push(t);
-      });
-      if (fullResponseText.trim()) {
-        messageParts.push({ type: "text", text: fullResponseText });
-      }
-      if (messageParts.length > 0) {
-        await this.persistMessages([
-          ...this.messages,
-          {
-            id: `assistant_${Date.now()}_${Math.random().toString(36).slice(2, 11)}`,
-            role: "assistant",
-            parts: messageParts
-          }
-        ]);
+      if (message.parts.length > 0) {
+        await this.persistMessages([...this.messages, message]);
       }
     });
   }