agents 0.0.0-dd6a9e3 → 0.0.0-df716f2

package/dist/mcp/index.d.ts

@@ -0,0 +1,45 @@
+import { DurableObject } from 'cloudflare:workers';
+import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
+import { Connection } from 'partyserver';
+
+interface CORSOptions {
+    origin?: string;
+    methods?: string;
+    headers?: string;
+    maxAge?: number;
+}
+declare abstract class McpAgent<Env = unknown, State = unknown, Props extends Record<string, unknown> = Record<string, unknown>> extends DurableObject<Env> {
+    #private;
+    protected constructor(ctx: DurableObjectState, env: Env);
+    /**
+     * Agents API allowlist
+     */
+    initialState: State;
+    get state(): State;
+    sql<T = Record<string, string | number | boolean | null>>(strings: TemplateStringsArray, ...values: (string | number | boolean | null)[]): T[];
+    setState(state: State): void;
+    onStateUpdate(state: State | undefined, source: Connection | "server"): void;
+    onStart(): Promise<void>;
+    /**
+     * McpAgent API
+     */
+    abstract server: McpServer;
+    props: Props;
+    initRun: boolean;
+    abstract init(): Promise<void>;
+    _init(props: Props): Promise<void>;
+    fetch(request: Request): Promise<Response>;
+    getWebSocket(): WebSocket | null;
+    onMCPMessage(sessionId: string, request: Request): Promise<Response>;
+    webSocketMessage(ws: WebSocket, event: ArrayBuffer | string): Promise<void>;
+    webSocketError(ws: WebSocket, error: unknown): Promise<void>;
+    webSocketClose(ws: WebSocket, code: number, reason: string, wasClean: boolean): Promise<void>;
+    static mount(path: string, { binding, corsOptions, }?: {
+        binding?: string;
+        corsOptions?: CORSOptions;
+    }): {
+        fetch: (request: Request, env: Record<string, DurableObjectNamespace<McpAgent>>, ctx: ExecutionContext) => Promise<Response | undefined>;
+    };
+}
+
+export { McpAgent };

package/dist/react.d.ts

@@ -1,21 +1,18 @@
-import { PartySocket } from "partysocket";
-import { usePartySocket } from "partysocket/react";
-import { StreamOptions } from "./client.js";
+import { PartySocket } from 'partysocket';
+import { usePartySocket } from 'partysocket/react';
+import { StreamOptions } from './client.js';
 
 /**
  * Options for the useAgent hook
  * @template State Type of the Agent's state
  */
-type UseAgentOptions<State = unknown> = Omit<
-  Parameters<typeof usePartySocket>[0],
-  "party" | "room"
-> & {
-  /** Name of the agent to connect to */
-  agent: string;
-  /** Name of the specific Agent instance */
-  name?: string;
-  /** Called when the Agent's state is updated */
-  onStateUpdate?: (state: State, source: "server" | "client") => void;
+type UseAgentOptions<State = unknown> = Omit<Parameters<typeof usePartySocket>[0], "party" | "room"> & {
+    /** Name of the agent to connect to */
+    agent: string;
+    /** Name of the specific Agent instance */
+    name?: string;
+    /** Called when the Agent's state is updated */
+    onStateUpdate?: (state: State, source: "server" | "client") => void;
 };
 /**
  * React hook for connecting to an Agent
@@ -23,17 +20,11 @@ type UseAgentOptions<State = unknown> = Omit<
  * @param options Connection options
  * @returns WebSocket connection with setState and call methods
  */
-declare function useAgent<State = unknown>(
-  options: UseAgentOptions<State>
-): PartySocket & {
-  agent: string;
-  name: string;
-  setState: (state: State) => void;
-  call: <T = unknown>(
-    method: string,
-    args?: unknown[],
-    streamOptions?: StreamOptions
-  ) => Promise<T>;
+declare function useAgent<State = unknown>(options: UseAgentOptions<State>): PartySocket & {
+    agent: string;
+    name: string;
+    setState: (state: State) => void;
+    call: <T = unknown>(method: string, args?: unknown[], streamOptions?: StreamOptions) => Promise<T>;
 };
 
 export { type UseAgentOptions, useAgent };

package/dist/schedule.d.ts

@@ -1,53 +1,43 @@
-import { z } from "zod";
+import { z } from 'zod';
 
 type Schedule = z.infer<typeof unstable_scheduleSchema>;
-declare function unstable_getSchedulePrompt(event: { date: Date }): string;
-declare const unstable_scheduleSchema: z.ZodObject<
-  {
+declare function unstable_getSchedulePrompt(event: {
+    date: Date;
+}): string;
+declare const unstable_scheduleSchema: z.ZodObject<{
     description: z.ZodString;
-    when: z.ZodObject<
-      {
+    when: z.ZodObject<{
         type: z.ZodEnum<["scheduled", "delayed", "cron", "no-schedule"]>;
         date: z.ZodOptional<z.ZodDate>;
         delayInSeconds: z.ZodOptional<z.ZodNumber>;
         cron: z.ZodOptional<z.ZodString>;
-      },
-      "strip",
-      z.ZodTypeAny,
-      {
+    }, "strip", z.ZodTypeAny, {
         type: "scheduled" | "delayed" | "cron" | "no-schedule";
         cron?: string | undefined;
         delayInSeconds?: number | undefined;
         date?: Date | undefined;
-      },
-      {
+    }, {
         type: "scheduled" | "delayed" | "cron" | "no-schedule";
         cron?: string | undefined;
         delayInSeconds?: number | undefined;
         date?: Date | undefined;
-      }
-    >;
-  },
-  "strip",
-  z.ZodTypeAny,
-  {
+    }>;
+}, "strip", z.ZodTypeAny, {
     description: string;
     when: {
-      type: "scheduled" | "delayed" | "cron" | "no-schedule";
-      cron?: string | undefined;
-      delayInSeconds?: number | undefined;
-      date?: Date | undefined;
+        type: "scheduled" | "delayed" | "cron" | "no-schedule";
+        cron?: string | undefined;
+        delayInSeconds?: number | undefined;
+        date?: Date | undefined;
     };
-  },
-  {
+}, {
     description: string;
     when: {
-      type: "scheduled" | "delayed" | "cron" | "no-schedule";
-      cron?: string | undefined;
-      delayInSeconds?: number | undefined;
-      date?: Date | undefined;
+        type: "scheduled" | "delayed" | "cron" | "no-schedule";
+        cron?: string | undefined;
+        delayInSeconds?: number | undefined;
+        date?: Date | undefined;
     };
-  }
->;
+}>;
 
 export { type Schedule, unstable_getSchedulePrompt, unstable_scheduleSchema };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "agents",
-  "version": "0.0.0-dd6a9e3",
+  "version": "0.0.0-df716f2",
   "main": "src/index.ts",
   "type": "module",
   "scripts": {
@@ -44,9 +44,14 @@
       "import": "./dist/schedule.js"
     },
     "./mcp": {
-      "types": "./dist/mcp.d.ts",
-      "require": "./dist/mcp.js",
-      "import": "./dist/mcp.js"
+      "types": "./dist/mcp/index.d.ts",
+      "require": "./dist/mcp/index.js",
+      "import": "./dist/mcp/index.js"
+    },
+    "./mcp/client": {
+      "types": "./dist/mcp/client.d.ts",
+      "require": "./dist/mcp/client.js",
+      "import": "./dist/mcp/client.js"
     }
   },
   "keywords": [],
@@ -65,9 +70,9 @@
     "cron-schedule": "^5.0.4",
     "nanoid": "^5.1.5",
     "partyserver": "^0.0.66",
-    "partysocket": "0.0.0-548c226"
+    "partysocket": "1.1.3"
   },
   "devDependencies": {
-    "@modelcontextprotocol/sdk": "^1.7.0"
+    "@modelcontextprotocol/sdk": "^1.8.0"
   }
 }

package/src/index.ts

@@ -11,6 +11,8 @@ import {
 import { parseCronExpression } from "cron-schedule";
 import { nanoid } from "nanoid";
 
+import { AsyncLocalStorage } from "node:async_hooks";
+
 export type { Connection, WSMessage, ConnectionContext } from "partyserver";
 
 import { WorkflowEntrypoint as CFWorkflowEntrypoint } from "cloudflare:workers";
@@ -168,6 +170,12 @@ const STATE_WAS_CHANGED = "cf_state_was_changed";
 
 const DEFAULT_STATE = {} as unknown;
 
+export const unstable_context = new AsyncLocalStorage<{
+  agent: Agent<unknown>;
+  connection: Connection | undefined;
+  request: Request | undefined;
+}>();
+
 /**
  * Base class for creating Agent implementations
  * @template Env Environment type containing bindings
@@ -292,89 +300,100 @@ export class Agent<Env, State = unknown> extends Server<Env> {
 
     const _onMessage = this.onMessage.bind(this);
     this.onMessage = async (connection: Connection, message: WSMessage) => {
-      if (typeof message !== "string") {
-        return this.#tryCatch(() => _onMessage(connection, message));
-      }
-
-      let parsed: unknown;
-      try {
-        parsed = JSON.parse(message);
-      } catch (e) {
-        // silently fail and let the onMessage handler handle it
-        return this.#tryCatch(() => _onMessage(connection, message));
-      }
-
-      if (isStateUpdateMessage(parsed)) {
-        this.#setStateInternal(parsed.state as State, connection);
-        return;
-      }
-
-      if (isRPCRequest(parsed)) {
-        try {
-          const { id, method, args } = parsed;
-
-          // Check if method exists and is callable
-          const methodFn = this[method as keyof this];
-          if (typeof methodFn !== "function") {
-            throw new Error(`Method ${method} does not exist`);
+      return unstable_context.run(
+        { agent: this, connection, request: undefined },
+        async () => {
+          if (typeof message !== "string") {
+            return this.#tryCatch(() => _onMessage(connection, message));
           }
 
-          if (!this.#isCallable(method)) {
-            throw new Error(`Method ${method} is not callable`);
+          let parsed: unknown;
+          try {
+            parsed = JSON.parse(message);
+          } catch (e) {
+            // silently fail and let the onMessage handler handle it
+            return this.#tryCatch(() => _onMessage(connection, message));
           }
 
-          // biome-ignore lint/complexity/noBannedTypes: <explanation>
-          const metadata = callableMetadata.get(methodFn as Function);
+          if (isStateUpdateMessage(parsed)) {
+            this.#setStateInternal(parsed.state as State, connection);
+            return;
+          }
 
-          // For streaming methods, pass a StreamingResponse object
-          if (metadata?.streaming) {
-            const stream = new StreamingResponse(connection, id);
-            await methodFn.apply(this, [stream, ...args]);
+          if (isRPCRequest(parsed)) {
+            try {
+              const { id, method, args } = parsed;
+
+              // Check if method exists and is callable
+              const methodFn = this[method as keyof this];
+              if (typeof methodFn !== "function") {
+                throw new Error(`Method ${method} does not exist`);
+              }
+
+              if (!this.#isCallable(method)) {
+                throw new Error(`Method ${method} is not callable`);
+              }
+
+              // biome-ignore lint/complexity/noBannedTypes: <explanation>
+              const metadata = callableMetadata.get(methodFn as Function);
+
+              // For streaming methods, pass a StreamingResponse object
+              if (metadata?.streaming) {
+                const stream = new StreamingResponse(connection, id);
+                await methodFn.apply(this, [stream, ...args]);
+                return;
+              }
+
+              // For regular methods, execute and send response
+              const result = await methodFn.apply(this, args);
+              const response: RPCResponse = {
+                type: "rpc",
+                id,
+                success: true,
+                result,
+                done: true,
+              };
+              connection.send(JSON.stringify(response));
+            } catch (e) {
+              // Send error response
+              const response: RPCResponse = {
+                type: "rpc",
+                id: parsed.id,
+                success: false,
+                error:
+                  e instanceof Error ? e.message : "Unknown error occurred",
+              };
+              connection.send(JSON.stringify(response));
+              console.error("RPC error:", e);
+            }
             return;
           }
 
-          // For regular methods, execute and send response
-          const result = await methodFn.apply(this, args);
-          const response: RPCResponse = {
-            type: "rpc",
-            id,
-            success: true,
-            result,
-            done: true,
-          };
-          connection.send(JSON.stringify(response));
-        } catch (e) {
-          // Send error response
-          const response: RPCResponse = {
-            type: "rpc",
-            id: parsed.id,
-            success: false,
-            error: e instanceof Error ? e.message : "Unknown error occurred",
-          };
-          connection.send(JSON.stringify(response));
-          console.error("RPC error:", e);
+          return this.#tryCatch(() => _onMessage(connection, message));
         }
-        return;
-      }
-
-      return this.#tryCatch(() => _onMessage(connection, message));
+      );
     };
 
     const _onConnect = this.onConnect.bind(this);
     this.onConnect = (connection: Connection, ctx: ConnectionContext) => {
       // TODO: This is a hack to ensure the state is sent after the connection is established
       // must fix this
-      setTimeout(() => {
-        if (this.state) {
-          connection.send(
-            JSON.stringify({
-              type: "cf_agent_state",
-              state: this.state,
-            })
-          );
+      return unstable_context.run(
+        { agent: this, connection, request: ctx.request },
+        async () => {
+          setTimeout(() => {
+            if (this.state) {
+              connection.send(
+                JSON.stringify({
+                  type: "cf_agent_state",
+                  state: this.state,
+                })
+              );
+            }
+            return this.#tryCatch(() => _onConnect(connection, ctx));
+          }, 20);
         }
-        return this.#tryCatch(() => _onConnect(connection, ctx));
-      }, 20);
+      );
     };
   }
 
@@ -395,7 +414,15 @@ export class Agent<Env, State = unknown> extends Server<Env> {
       }),
       source !== "server" ? [source.id] : []
     );
-    return this.#tryCatch(() => this.onStateUpdate(state, source));
+    return this.#tryCatch(() => {
+      const { connection, request } = unstable_context.getStore() || {};
+      return unstable_context.run(
+        { agent: this, connection, request },
+        async () => {
+          return this.onStateUpdate(state, source);
+        }
+      );
+    });
   }
 
   /**
@@ -420,7 +447,12 @@ export class Agent<Env, State = unknown> extends Server<Env> {
    * @param email Email message to process
    */
   onEmail(email: ForwardableEmailMessage) {
-    throw new Error("Not implemented");
+    return unstable_context.run(
+      { agent: this, connection: undefined, request: undefined },
+      async () => {
+        console.error("onEmail not implemented");
+      }
+    );
   }
 
   async #tryCatch<T>(fn: () => T | Promise<T>) {
@@ -580,7 +612,6 @@ export class Agent<Env, State = unknown> extends Server<Env> {
    */
   getSchedules<T = string>(
     criteria: {
-      description?: string;
       id?: string;
       type?: "scheduled" | "delayed" | "cron";
       timeRange?: { start?: Date; end?: Date };
@@ -594,11 +625,6 @@ export class Agent<Env, State = unknown> extends Server<Env> {
       params.push(criteria.id);
     }
 
-    if (criteria.description) {
-      query += " AND description = ?";
-      params.push(criteria.description);
-    }
-
     if (criteria.type) {
       query += " AND type = ?";
       params.push(criteria.type);
@@ -671,16 +697,21 @@ export class Agent<Env, State = unknown> extends Server<Env> {
         console.error(`callback ${row.callback} not found`);
         continue;
       }
-      try {
-        await (
-          callback as (
-            payload: unknown,
-            schedule: Schedule<unknown>
-          ) => Promise<void>
-        ).bind(this)(JSON.parse(row.payload as string), row);
-      } catch (e) {
-        console.error(`error executing callback "${row.callback}"`, e);
-      }
+      await unstable_context.run(
+        { agent: this, connection: undefined, request: undefined },
+        async () => {
+          try {
+            await (
+              callback as (
+                payload: unknown,
+                schedule: Schedule<unknown>
+              ) => Promise<void>
+            ).bind(this)(JSON.parse(row.payload as string), row);
+          } catch (e) {
+            console.error(`error executing callback "${row.callback}"`, e);
+          }
+        }
+      );
       if (row.type === "cron") {
         // Update next execution time for cron schedules
         const nextExecutionTime = getNextCronTime(row.cron);

package/dist/ai-chat-agent.js

@@ -1,167 +0,0 @@
-import {
-  Agent
-} from "./chunk-X6BBKLSC.js";
-import {
-  __privateAdd,
-  __privateMethod
-} from "./chunk-HMLY7DHA.js";
-
-// src/ai-chat-agent.ts
-import { appendResponseMessages } from "ai";
-var decoder = new TextDecoder();
-var _AIChatAgent_instances, broadcastChatMessage_fn, tryCatch_fn, persistMessages_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) => {
-      return JSON.parse(row.message);
-    });
-  }
-  async onMessage(connection, message) {
-    if (typeof message === "string") {
-      let data;
-      try {
-        data = JSON.parse(message);
-      } catch (error) {
-        return;
-      }
-      if (data.type === "cf_agent_use_chat_request" && data.init.method === "POST") {
-        const {
-          method,
-          keepalive,
-          headers,
-          body,
-          // we're reading this
-          redirect,
-          integrity,
-          credentials,
-          mode,
-          referrer,
-          referrerPolicy,
-          window
-          // 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 __privateMethod(this, _AIChatAgent_instances, persistMessages_fn).call(this, 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 __privateMethod(this, _AIChatAgent_instances, persistMessages_fn).call(this, finalMessages, [connection.id]);
-          });
-          if (response) {
-            await __privateMethod(this, _AIChatAgent_instances, reply_fn).call(this, data.id, response);
-          }
-        });
-      }
-      if (data.type === "cf_agent_chat_clear") {
-        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 __privateMethod(this, _AIChatAgent_instances, persistMessages_fn).call(this, data.messages, [connection.id]);
-      }
-    }
-  }
-  async onRequest(request) {
-    return __privateMethod(this, _AIChatAgent_instances, tryCatch_fn).call(this, () => {
-      if (request.url.endsWith("/get-messages")) {
-        const messages = (this.sql`select * from cf_ai_chat_agent_messages` || []).map((row) => {
-          return JSON.parse(row.message);
-        });
-        return Response.json(messages);
-      }
-      return super.onRequest(request);
-    });
-  }
-  /**
-   * Handle incoming chat messages and generate a response
-   * @param onFinish Callback to be called when the response is finished
-   * @returns Response to send to the client or undefined
-   */
-  async onChatMessage(onFinish) {
-    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
-   * @param messages Chat messages to save
-   */
-  async saveMessages(messages) {
-    await __privateMethod(this, _AIChatAgent_instances, persistMessages_fn).call(this, messages);
-    const response = await this.onChatMessage(async ({ response: response2 }) => {
-      const finalMessages = appendResponseMessages({
-        messages,
-        responseMessages: response2.messages
-      });
-      await __privateMethod(this, _AIChatAgent_instances, persistMessages_fn).call(this, finalMessages, []);
-    });
-    if (response) {
-      for await (const chunk of response.body) {
-        decoder.decode(chunk);
-      }
-      response.body?.cancel();
-    }
-  }
-};
-_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);
-  }
-};
-persistMessages_fn = async function(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.messages = messages;
-  __privateMethod(this, _AIChatAgent_instances, broadcastChatMessage_fn).call(this, {
-    type: "cf_agent_chat_messages",
-    messages
-  }, excludeBroadcastIds);
-};
-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
-      });
-    }
-    __privateMethod(this, _AIChatAgent_instances, broadcastChatMessage_fn).call(this, {
-      id,
-      type: "cf_agent_use_chat_response",
-      body: "",
-      done: true
-    });
-  });
-};
-export {
-  AIChatAgent
-};
-//# sourceMappingURL=ai-chat-agent.js.map

package/dist/ai-chat-agent.js.map

@@ -1 +0,0 @@
-{"version":3,"sources":["../src/ai-chat-agent.ts"],"sourcesContent":["import { Agent, type AgentContext, type Connection, type WSMessage } from \"./\";\nimport type {\n  Message as ChatMessage,\n  StreamTextOnFinishCallback,\n  ToolSet,\n} from \"ai\";\nimport { appendResponseMessages } from \"ai\";\nimport type { OutgoingMessage, IncomingMessage } from \"./ai-types\";\nconst decoder = new TextDecoder();\n\n/**\n * Extension of Agent with built-in chat capabilities\n * @template Env Environment type containing bindings\n */\nexport class AIChatAgent<Env = unknown, State = unknown> extends Agent<\n  Env,\n  State\n> {\n  /** Array of chat messages for the current conversation */\n  messages: ChatMessage[];\n  constructor(ctx: AgentContext, env: Env) {\n    super(ctx, env);\n    this.sql`create table if not exists cf_ai_chat_agent_messages (\n      id text primary key,\n      message text not null,\n      created_at datetime default current_timestamp\n    )`;\n    this.messages = (\n      this.sql`select * from cf_ai_chat_agent_messages` || []\n    ).map((row) => {\n      return JSON.parse(row.message as string);\n    });\n  }\n\n  #broadcastChatMessage(message: OutgoingMessage, exclude?: string[]) {\n    this.broadcast(JSON.stringify(message), exclude);\n  }\n\n  override async onMessage(connection: Connection, message: WSMessage) {\n    if (typeof message === \"string\") {\n      let data: IncomingMessage;\n      try {\n        data = JSON.parse(message) as IncomingMessage;\n      } catch (error) {\n        // silently ignore invalid messages for now\n        // TODO: log errors with log levels\n        return;\n      }\n      if (\n        data.type === \"cf_agent_use_chat_request\" &&\n        data.init.method === \"POST\"\n      ) {\n        const {\n          method,\n          keepalive,\n          headers,\n          body, // we're reading this\n          redirect,\n          integrity,\n          credentials,\n          mode,\n          referrer,\n          referrerPolicy,\n          window,\n          // dispatcher,\n          // duplex\n        } = data.init;\n        const { messages } = JSON.parse(body as string);\n        this.#broadcastChatMessage(\n          {\n            type: \"cf_agent_chat_messages\",\n            messages,\n          },\n          [connection.id]\n        );\n        await this.#persistMessages(messages, [connection.id]);\n        return this.#tryCatch(async () => {\n          const response = await this.onChatMessage(async ({ response }) => {\n            const finalMessages = appendResponseMessages({\n              messages,\n              responseMessages: response.messages,\n            });\n\n            await this.#persistMessages(finalMessages, [connection.id]);\n          });\n          if (response) {\n            await this.#reply(data.id, response);\n          }\n        });\n      }\n      if (data.type === \"cf_agent_chat_clear\") {\n        this.sql`delete from cf_ai_chat_agent_messages`;\n        this.messages = [];\n        this.#broadcastChatMessage(\n          {\n            type: \"cf_agent_chat_clear\",\n          },\n          [connection.id]\n        );\n      } else if (data.type === \"cf_agent_chat_messages\") {\n        // replace the messages with the new ones\n        await this.#persistMessages(data.messages, [connection.id]);\n      }\n    }\n  }\n\n  override async onRequest(request: Request): Promise<Response> {\n    return this.#tryCatch(() => {\n      if (request.url.endsWith(\"/get-messages\")) {\n        const messages = (\n          this.sql`select * from cf_ai_chat_agent_messages` || []\n        ).map((row) => {\n          return JSON.parse(row.message as string);\n        });\n        return Response.json(messages);\n      }\n      return super.onRequest(request);\n    });\n  }\n\n  async #tryCatch<T>(fn: () => T | Promise<T>) {\n    try {\n      return await fn();\n    } catch (e) {\n      throw this.onError(e);\n    }\n  }\n\n  /**\n   * Handle incoming chat messages and generate a response\n   * @param onFinish Callback to be called when the response is finished\n   * @returns Response to send to the client or undefined\n   */\n  async onChatMessage(\n    onFinish: StreamTextOnFinishCallback<ToolSet>\n  ): Promise<Response | undefined> {\n    throw new Error(\n      \"recieved a chat message, override onChatMessage and return a Response to send to the client\"\n    );\n  }\n\n  /**\n   * Save messages on the server side and trigger AI response\n   * @param messages Chat messages to save\n   */\n  async saveMessages(messages: ChatMessage[]) {\n    await this.#persistMessages(messages);\n    const response = await this.onChatMessage(async ({ response }) => {\n      const finalMessages = appendResponseMessages({\n        messages,\n        responseMessages: response.messages,\n      });\n\n      await this.#persistMessages(finalMessages, []);\n    });\n    if (response) {\n      // we're just going to drain the body\n      // @ts-ignore TODO: fix this type error\n      for await (const chunk of response.body!) {\n        decoder.decode(chunk);\n      }\n      response.body?.cancel();\n    }\n  }\n\n  async #persistMessages(\n    messages: ChatMessage[],\n    excludeBroadcastIds: string[] = []\n  ) {\n    this.sql`delete from cf_ai_chat_agent_messages`;\n    for (const message of messages) {\n      this.sql`insert into cf_ai_chat_agent_messages (id, message) values (${\n        message.id\n      },${JSON.stringify(message)})`;\n    }\n    this.messages = messages;\n    this.#broadcastChatMessage(\n      {\n        type: \"cf_agent_chat_messages\",\n        messages: messages,\n      },\n      excludeBroadcastIds\n    );\n  }\n\n  async #reply(id: string, response: Response) {\n    // now take chunks out from dataStreamResponse and send them to the client\n    return this.#tryCatch(async () => {\n      // @ts-expect-error TODO: fix this type error\n      for await (const chunk of response.body!) {\n        const body = decoder.decode(chunk);\n\n        this.#broadcastChatMessage({\n          id,\n          type: \"cf_agent_use_chat_response\",\n          body,\n          done: false,\n        });\n      }\n\n      this.#broadcastChatMessage({\n        id,\n        type: \"cf_agent_use_chat_response\",\n        body: \"\",\n        done: true,\n      });\n    });\n  }\n}\n"],"mappings":";;;;;;;;;AAMA,SAAS,8BAA8B;AAEvC,IAAM,UAAU,IAAI,YAAY;AARhC;AAcO,IAAM,cAAN,cAA0D,MAG/D;AAAA,EAGA,YAAY,KAAmB,KAAU;AACvC,UAAM,KAAK,GAAG;AAPX;AAQH,SAAK;AAAA;AAAA;AAAA;AAAA;AAKL,SAAK,YACH,KAAK,gDAAgD,CAAC,GACtD,IAAI,CAAC,QAAQ;AACb,aAAO,KAAK,MAAM,IAAI,OAAiB;AAAA,IACzC,CAAC;AAAA,EACH;AAAA,EAMA,MAAe,UAAU,YAAwB,SAAoB;AACnE,QAAI,OAAO,YAAY,UAAU;AAC/B,UAAI;AACJ,UAAI;AACF,eAAO,KAAK,MAAM,OAAO;AAAA,MAC3B,SAAS,OAAO;AAGd;AAAA,MACF;AACA,UACE,KAAK,SAAS,+BACd,KAAK,KAAK,WAAW,QACrB;AACA,cAAM;AAAA,UACJ;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA;AAAA;AAAA,QAGF,IAAI,KAAK;AACT,cAAM,EAAE,SAAS,IAAI,KAAK,MAAM,IAAc;AAC9C,8BAAK,iDAAL,WACE;AAAA,UACE,MAAM;AAAA,UACN;AAAA,QACF,GACA,CAAC,WAAW,EAAE;AAEhB,cAAM,sBAAK,4CAAL,WAAsB,UAAU,CAAC,WAAW,EAAE;AACpD,eAAO,sBAAK,qCAAL,WAAe,YAAY;AAChC,gBAAM,WAAW,MAAM,KAAK,cAAc,OAAO,EAAE,UAAAA,UAAS,MAAM;AAChE,kBAAM,gBAAgB,uBAAuB;AAAA,cAC3C;AAAA,cACA,kBAAkBA,UAAS;AAAA,YAC7B,CAAC;AAED,kBAAM,sBAAK,4CAAL,WAAsB,eAAe,CAAC,WAAW,EAAE;AAAA,UAC3D,CAAC;AACD,cAAI,UAAU;AACZ,kBAAM,sBAAK,kCAAL,WAAY,KAAK,IAAI;AAAA,UAC7B;AAAA,QACF;AAAA,MACF;AACA,UAAI,KAAK,SAAS,uBAAuB;AACvC,aAAK;AACL,aAAK,WAAW,CAAC;AACjB,8BAAK,iDAAL,WACE;AAAA,UACE,MAAM;AAAA,QACR,GACA,CAAC,WAAW,EAAE;AAAA,MAElB,WAAW,KAAK,SAAS,0BAA0B;AAEjD,cAAM,sBAAK,4CAAL,WAAsB,KAAK,UAAU,CAAC,WAAW,EAAE;AAAA,MAC3D;AAAA,IACF;AAAA,EACF;AAAA,EAEA,MAAe,UAAU,SAAqC;AAC5D,WAAO,sBAAK,qCAAL,WAAe,MAAM;AAC1B,UAAI,QAAQ,IAAI,SAAS,eAAe,GAAG;AACzC,cAAM,YACJ,KAAK,gDAAgD,CAAC,GACtD,IAAI,CAAC,QAAQ;AACb,iBAAO,KAAK,MAAM,IAAI,OAAiB;AAAA,QACzC,CAAC;AACD,eAAO,SAAS,KAAK,QAAQ;AAAA,MAC/B;AACA,aAAO,MAAM,UAAU,OAAO;AAAA,IAChC;AAAA,EACF;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAeA,MAAM,cACJ,UAC+B;AAC/B,UAAM,IAAI;AAAA,MACR;AAAA,IACF;AAAA,EACF;AAAA;AAAA;AAAA;AAAA;AAAA,EAMA,MAAM,aAAa,UAAyB;AAC1C,UAAM,sBAAK,4CAAL,WAAsB;AAC5B,UAAM,WAAW,MAAM,KAAK,cAAc,OAAO,EAAE,UAAAA,UAAS,MAAM;AAChE,YAAM,gBAAgB,uBAAuB;AAAA,QAC3C;AAAA,QACA,kBAAkBA,UAAS;AAAA,MAC7B,CAAC;AAED,YAAM,sBAAK,4CAAL,WAAsB,eAAe,CAAC;AAAA,IAC9C,CAAC;AACD,QAAI,UAAU;AAGZ,uBAAiB,SAAS,SAAS,MAAO;AACxC,gBAAQ,OAAO,KAAK;AAAA,MACtB;AACA,eAAS,MAAM,OAAO;AAAA,IACxB;AAAA,EACF;AA6CF;AAlMO;AAoBL,0BAAqB,SAAC,SAA0B,SAAoB;AAClE,OAAK,UAAU,KAAK,UAAU,OAAO,GAAG,OAAO;AACjD;AAoFM,cAAY,eAAC,IAA0B;AAC3C,MAAI;AACF,WAAO,MAAM,GAAG;AAAA,EAClB,SAAS,GAAG;AACV,UAAM,KAAK,QAAQ,CAAC;AAAA,EACtB;AACF;AAuCM,qBAAgB,eACpB,UACA,sBAAgC,CAAC,GACjC;AACA,OAAK;AACL,aAAW,WAAW,UAAU;AAC9B,SAAK,kEACH,QAAQ,EACV,IAAI,KAAK,UAAU,OAAO,CAAC;AAAA,EAC7B;AACA,OAAK,WAAW;AAChB,wBAAK,iDAAL,WACE;AAAA,IACE,MAAM;AAAA,IACN;AAAA,EACF,GACA;AAEJ;AAEM,WAAM,eAAC,IAAY,UAAoB;AAE3C,SAAO,sBAAK,qCAAL,WAAe,YAAY;AAEhC,qBAAiB,SAAS,SAAS,MAAO;AACxC,YAAM,OAAO,QAAQ,OAAO,KAAK;AAEjC,4BAAK,iDAAL,WAA2B;AAAA,QACzB;AAAA,QACA,MAAM;AAAA,QACN;AAAA,QACA,MAAM;AAAA,MACR;AAAA,IACF;AAEA,0BAAK,iDAAL,WAA2B;AAAA,MACzB;AAAA,MACA,MAAM;AAAA,MACN,MAAM;AAAA,MACN,MAAM;AAAA,IACR;AAAA,EACF;AACF;","names":["response"]}