agents 0.0.0-f5b5854 → 0.0.0-f6b2319

package/dist/ai-chat-agent.js

@@ -1,85 +1,136 @@
 import {
-  Agent
-} from "./chunk-XG52S6YY.js";
+  autoTransformMessages
+} from "./chunk-UJVEAURM.js";
 import {
-  __privateAdd,
-  __privateMethod
-} from "./chunk-HMLY7DHA.js";
+  Agent
+} from "./chunk-254F4GDT.js";
+import "./chunk-3OT2NNEW.js";
+import "./chunk-Z44WASMA.js";
+import "./chunk-QEVM4BVL.js";
+import "./chunk-BER7KXUJ.js";
+import "./chunk-PR4QN5HX.js";
 
 // src/ai-chat-agent.ts
-import { appendResponseMessages } from "ai";
+import {
+  getToolName,
+  isToolUIPart,
+  parsePartialJson
+} from "ai";
 var decoder = new TextDecoder();
-var _AIChatAgent_instances, broadcastChatMessage_fn, tryCatch_fn, reply_fn;
 var AIChatAgent = class extends Agent {
   constructor(ctx, env) {
     super(ctx, env);
-    __privateAdd(this, _AIChatAgent_instances);
     this.sql`create table if not exists cf_ai_chat_agent_messages (
       id text primary key,
       message text not null,
       created_at datetime default current_timestamp
     )`;
-    this.messages = (this.sql`select * from cf_ai_chat_agent_messages` || []).map((row) => {
+    const rawMessages = (this.sql`select * from cf_ai_chat_agent_messages` || []).map((row) => {
       return JSON.parse(row.message);
     });
+    this.messages = autoTransformMessages(rawMessages);
+    this._chatMessageAbortControllers = /* @__PURE__ */ new Map();
+  }
+  _broadcastChatMessage(message, exclude) {
+    this.broadcast(JSON.stringify(message), exclude);
   }
   async onMessage(connection, message) {
     if (typeof message === "string") {
       let data;
       try {
         data = JSON.parse(message);
-      } catch (error) {
+      } catch (_error) {
         return;
       }
-      if (data.type === "cf_agent_use_chat_request" && data.init.method === "POST") {
+      if (data.type === "cf_agent_use_chat_request" /* CF_AGENT_USE_CHAT_REQUEST */ && data.init.method === "POST") {
         const {
-          method,
-          keepalive,
-          headers,
-          body,
+          // method,
+          // keepalive,
+          // headers,
+          body
           // we're reading this
-          redirect,
-          integrity,
-          credentials,
-          mode,
-          referrer,
-          referrerPolicy,
-          window
+          //
+          // // these might not exist?
           // dispatcher,
           // duplex
         } = data.init;
         const { messages } = JSON.parse(body);
-        __privateMethod(this, _AIChatAgent_instances, broadcastChatMessage_fn).call(this, {
-          type: "cf_agent_chat_messages",
-          messages
-        }, [connection.id]);
-        await this.persistMessages(messages, [connection.id]);
-        return __privateMethod(this, _AIChatAgent_instances, tryCatch_fn).call(this, async () => {
-          const response = await this.onChatMessage(async ({ response: response2 }) => {
-            const finalMessages = appendResponseMessages({
-              messages,
-              responseMessages: response2.messages
-            });
-            await this.persistMessages(finalMessages, [connection.id]);
-          });
+        const transformedMessages = autoTransformMessages(messages);
+        this._broadcastChatMessage(
+          {
+            messages: transformedMessages,
+            type: "cf_agent_chat_messages" /* CF_AGENT_CHAT_MESSAGES */
+          },
+          [connection.id]
+        );
+        await this.persistMessages(transformedMessages, [connection.id]);
+        this.observability?.emit(
+          {
+            displayMessage: "Chat message request",
+            id: data.id,
+            payload: {},
+            timestamp: Date.now(),
+            type: "message:request"
+          },
+          this.ctx
+        );
+        const chatMessageId = data.id;
+        const abortSignal = this._getAbortSignal(chatMessageId);
+        return this._tryCatchChat(async () => {
+          const response = await this.onChatMessage(
+            async (_finishResult) => {
+              this._removeAbortController(chatMessageId);
+              this.observability?.emit(
+                {
+                  displayMessage: "Chat message response",
+                  id: data.id,
+                  payload: {},
+                  timestamp: Date.now(),
+                  type: "message:response"
+                },
+                this.ctx
+              );
+            },
+            abortSignal ? { abortSignal } : void 0
+          );
           if (response) {
-            await __privateMethod(this, _AIChatAgent_instances, reply_fn).call(this, data.id, response);
+            await this._reply(data.id, response);
+          } else {
+            console.warn(
+              `[AIChatAgent] onChatMessage returned no response for chatMessageId: ${chatMessageId}`
+            );
+            this._broadcastChatMessage(
+              {
+                body: "No response was generated by the agent.",
+                done: true,
+                id: data.id,
+                type: "cf_agent_use_chat_response" /* CF_AGENT_USE_CHAT_RESPONSE */
+              },
+              [connection.id]
+            );
           }
         });
       }
-      if (data.type === "cf_agent_chat_clear") {
+      if (data.type === "cf_agent_chat_clear" /* CF_AGENT_CHAT_CLEAR */) {
+        this._destroyAbortControllers();
         this.sql`delete from cf_ai_chat_agent_messages`;
         this.messages = [];
-        __privateMethod(this, _AIChatAgent_instances, broadcastChatMessage_fn).call(this, {
-          type: "cf_agent_chat_clear"
-        }, [connection.id]);
-      } else if (data.type === "cf_agent_chat_messages") {
-        await this.persistMessages(data.messages, [connection.id]);
+        this._broadcastChatMessage(
+          {
+            type: "cf_agent_chat_clear" /* CF_AGENT_CHAT_CLEAR */
+          },
+          [connection.id]
+        );
+      } else if (data.type === "cf_agent_chat_messages" /* CF_AGENT_CHAT_MESSAGES */) {
+        const transformedMessages = autoTransformMessages(data.messages);
+        await this.persistMessages(transformedMessages, [connection.id]);
+      } else if (data.type === "cf_agent_chat_request_cancel" /* CF_AGENT_CHAT_REQUEST_CANCEL */) {
+        this._cancelChatRequest(data.id);
       }
     }
   }
   async onRequest(request) {
-    return __privateMethod(this, _AIChatAgent_instances, tryCatch_fn).call(this, () => {
+    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) => {
@@ -90,35 +141,30 @@ var AIChatAgent = class extends Agent {
       return super.onRequest(request);
     });
   }
+  async _tryCatchChat(fn) {
+    try {
+      return await fn();
+    } catch (e) {
+      throw this.onError(e);
+    }
+  }
   /**
    * Handle incoming chat messages and generate a response
    * @param onFinish Callback to be called when the response is finished
+   * @param options.signal A signal to pass to any child requests which can be used to cancel them
    * @returns Response to send to the client or undefined
    */
-  async onChatMessage(onFinish) {
+  async onChatMessage(onFinish, options) {
     throw new Error(
       "recieved a chat message, override onChatMessage and return a Response to send to the client"
     );
   }
   /**
-   * Save messages on the server side and trigger AI response
+   * Save messages on the server side
    * @param messages Chat messages to save
    */
   async saveMessages(messages) {
     await this.persistMessages(messages);
-    const response = await this.onChatMessage(async ({ response: response2 }) => {
-      const finalMessages = appendResponseMessages({
-        messages,
-        responseMessages: response2.messages
-      });
-      await this.persistMessages(finalMessages, []);
-    });
-    if (response) {
-      for await (const chunk of response.body) {
-        decoder.decode(chunk);
-      }
-      response.body?.cancel();
-    }
   }
   async persistMessages(messages, excludeBroadcastIds = []) {
     this.sql`delete from cf_ai_chat_agent_messages`;
@@ -126,41 +172,479 @@ var AIChatAgent = class extends Agent {
       this.sql`insert into cf_ai_chat_agent_messages (id, message) values (${message.id},${JSON.stringify(message)})`;
     }
     this.messages = messages;
-    __privateMethod(this, _AIChatAgent_instances, broadcastChatMessage_fn).call(this, {
-      type: "cf_agent_chat_messages",
-      messages
-    }, excludeBroadcastIds);
+    this._broadcastChatMessage(
+      {
+        messages,
+        type: "cf_agent_chat_messages" /* CF_AGENT_CHAT_MESSAGES */
+      },
+      excludeBroadcastIds
+    );
   }
-};
-_AIChatAgent_instances = new WeakSet();
-broadcastChatMessage_fn = function(message, exclude) {
-  this.broadcast(JSON.stringify(message), exclude);
-};
-tryCatch_fn = async function(fn) {
-  try {
-    return await fn();
-  } catch (e) {
-    throw this.onError(e);
+  async _reply(id, response) {
+    return this._tryCatchChat(async () => {
+      if (!response.body) {
+        this._broadcastChatMessage({
+          body: "",
+          done: true,
+          id,
+          type: "cf_agent_use_chat_response" /* CF_AGENT_USE_CHAT_RESPONSE */
+        });
+        return;
+      }
+      const reader = response.body.getReader();
+      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();
+          if (done) {
+            this._broadcastChatMessage({
+              body: "",
+              done: true,
+              id,
+              type: "cf_agent_use_chat_response" /* CF_AGENT_USE_CHAT_RESPONSE */
+            });
+            break;
+          }
+          const chunk = decoder.decode(value);
+          const contentType = response.headers.get("content-type") || "";
+          const isSSE = contentType.includes("text/event-stream");
+          if (isSSE) {
+            const lines = chunk.split("\n");
+            for (const line of lines) {
+              if (line.startsWith("data: ") && line !== "data: [DONE]") {
+                try {
+                  const data = JSON.parse(line.slice(6));
+                  switch (data.type) {
+                    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;
+                    }
+                    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": {
+                      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": {
+                      this._broadcastChatMessage({
+                        error: true,
+                        body: data.errorText ?? JSON.stringify(data),
+                        done: false,
+                        id,
+                        type: "cf_agent_use_chat_response" /* CF_AGENT_USE_CHAT_RESPONSE */
+                      });
+                      break;
+                    }
+                  }
+                  this._broadcastChatMessage({
+                    body: JSON.stringify(data),
+                    done: false,
+                    id,
+                    type: "cf_agent_use_chat_response" /* CF_AGENT_USE_CHAT_RESPONSE */
+                  });
+                } catch (_e) {
+                }
+              }
+            }
+          } else {
+            if (chunk.length > 0) {
+              message.parts.push({ type: "text", text: chunk });
+              this._broadcastChatMessage({
+                body: JSON.stringify({ type: "text-delta", delta: chunk }),
+                done: false,
+                id,
+                type: "cf_agent_use_chat_response" /* CF_AGENT_USE_CHAT_RESPONSE */
+              });
+            }
+          }
+        }
+      } finally {
+        reader.releaseLock();
+      }
+      if (message.parts.length > 0) {
+        await this.persistMessages([...this.messages, message]);
+      }
+    });
   }
-};
-reply_fn = async function(id, response) {
-  return __privateMethod(this, _AIChatAgent_instances, tryCatch_fn).call(this, async () => {
-    for await (const chunk of response.body) {
-      const body = decoder.decode(chunk);
-      __privateMethod(this, _AIChatAgent_instances, broadcastChatMessage_fn).call(this, {
-        id,
-        type: "cf_agent_use_chat_response",
-        body,
-        done: false
-      });
+  /**
+   * For the given message id, look up its associated AbortController
+   * If the AbortController does not exist, create and store one in memory
+   *
+   * returns the AbortSignal associated with the AbortController
+   */
+  _getAbortSignal(id) {
+    if (typeof id !== "string") {
+      return void 0;
     }
-    __privateMethod(this, _AIChatAgent_instances, broadcastChatMessage_fn).call(this, {
-      id,
-      type: "cf_agent_use_chat_response",
-      body: "",
-      done: true
-    });
-  });
+    if (!this._chatMessageAbortControllers.has(id)) {
+      this._chatMessageAbortControllers.set(id, new AbortController());
+    }
+    return this._chatMessageAbortControllers.get(id)?.signal;
+  }
+  /**
+   * Remove an abort controller from the cache of pending message responses
+   */
+  _removeAbortController(id) {
+    this._chatMessageAbortControllers.delete(id);
+  }
+  /**
+   * Propagate an abort signal for any requests associated with the given message id
+   */
+  _cancelChatRequest(id) {
+    if (this._chatMessageAbortControllers.has(id)) {
+      const abortController = this._chatMessageAbortControllers.get(id);
+      abortController?.abort();
+    }
+  }
+  /**
+   * Abort all pending requests and clear the cache of AbortControllers
+   */
+  _destroyAbortControllers() {
+    for (const controller of this._chatMessageAbortControllers.values()) {
+      controller?.abort();
+    }
+    this._chatMessageAbortControllers.clear();
+  }
+  /**
+   * When the DO is destroyed, cancel all pending requests
+   */
+  async destroy() {
+    this._destroyAbortControllers();
+    await super.destroy();
+  }
 };
 export {
   AIChatAgent