agents 0.0.0-aa5f972 → 0.0.0-ac0e999

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,7 +1,6 @@
1
1
  import { Agent, AgentContext } from "./index.js";
2
2
  import { Message, StreamTextOnFinishCallback, ToolSet } from "ai";
3
3
  import { Connection, WSMessage } from "partyserver";
4
- import "node:async_hooks";
5
4
  import "./mcp/client.js";
6
5
  import "zod";
7
6
  import "@modelcontextprotocol/sdk/types.js";
@@ -20,12 +19,18 @@ declare class AIChatAgent<Env = unknown, State = unknown> extends Agent<
20
19
  Env,
21
20
  State
22
21
  > {
23
- #private;
22
+ /**
23
+ * Map of message `id`s to `AbortController`s
24
+ * useful to propagate request cancellation signals for any external calls made by the agent
25
+ */
26
+ private chatMessageAbortControllers;
24
27
  /** Array of chat messages for the current conversation */
25
28
  messages: Message[];
26
29
  constructor(ctx: AgentContext, env: Env);
30
+ private broadcastChatMessage;
27
31
  onMessage(connection: Connection, message: WSMessage): Promise<void>;
28
32
  onRequest(request: Request): Promise<Response>;
33
+ private tryCatchChat;
29
34
  /**
30
35
  * Handle incoming chat messages and generate a response
31
36
  * @param onFinish Callback to be called when the response is finished
@@ -47,6 +52,26 @@ declare class AIChatAgent<Env = unknown, State = unknown> extends Agent<
47
52
  messages: Message[],
48
53
  excludeBroadcastIds?: string[]
49
54
  ): Promise<void>;
55
+ private reply;
56
+ /**
57
+ * For the given message id, look up its associated AbortController
58
+ * If the AbortController does not exist, create and store one in memory
59
+ *
60
+ * returns the AbortSignal associated with the AbortController
61
+ */
62
+ private getAbortSignal;
63
+ /**
64
+ * Remove an abort controller from the cache of pending message responses
65
+ */
66
+ private removeAbortController;
67
+ /**
68
+ * Propagate an abort signal for any requests associated with the given message id
69
+ */
70
+ private cancelChatRequest;
71
+ /**
72
+ * Abort all pending requests and clear the cache of AbortControllers
73
+ */
74
+ private destroyAbortControllers;
50
75
  /**
51
76
  * When the DO is destroyed, cancel all pending requests
52
77
  */
@@ -1,27 +1,17 @@
1
1
  import {
2
2
  Agent
3
- } from "./chunk-AV3OMRR4.js";
4
- import "./chunk-YZNSS675.js";
5
- import {
6
- __privateAdd,
7
- __privateGet,
8
- __privateMethod,
9
- __privateSet
10
- } from "./chunk-HMLY7DHA.js";
3
+ } from "./chunk-DMJ7L3FI.js";
4
+ import "./chunk-D6UOOELW.js";
5
+ import "./chunk-ZKIVUOTQ.js";
6
+ import "./chunk-25YDMV4H.js";
7
+ import "./chunk-NOUFNU2O.js";
11
8
 
12
9
  // src/ai-chat-agent.ts
13
10
  import { appendResponseMessages } from "ai";
14
11
  var decoder = new TextDecoder();
15
- var _chatMessageAbortControllers, _AIChatAgent_instances, broadcastChatMessage_fn, tryCatch_fn, reply_fn, getAbortSignal_fn, removeAbortController_fn, cancelChatRequest_fn, destroyAbortControllers_fn;
16
12
  var AIChatAgent = class extends Agent {
17
13
  constructor(ctx, env) {
18
14
  super(ctx, env);
19
- __privateAdd(this, _AIChatAgent_instances);
20
- /**
21
- * Map of message `id`s to `AbortController`s
22
- * useful to propagate request cancellation signals for any external calls made by the agent
23
- */
24
- __privateAdd(this, _chatMessageAbortControllers);
25
15
  this.sql`create table if not exists cf_ai_chat_agent_messages (
26
16
  id text primary key,
27
17
  message text not null,
@@ -30,7 +20,10 @@ var AIChatAgent = class extends Agent {
30
20
  this.messages = (this.sql`select * from cf_ai_chat_agent_messages` || []).map((row) => {
31
21
  return JSON.parse(row.message);
32
22
  });
33
- __privateSet(this, _chatMessageAbortControllers, /* @__PURE__ */ new Map());
23
+ this.chatMessageAbortControllers = /* @__PURE__ */ new Map();
24
+ }
25
+ broadcastChatMessage(message, exclude) {
26
+ this.broadcast(JSON.stringify(message), exclude);
34
27
  }
35
28
  async onMessage(connection, message) {
36
29
  if (typeof message === "string") {
@@ -58,14 +51,17 @@ var AIChatAgent = class extends Agent {
58
51
  // duplex
59
52
  } = data.init;
60
53
  const { messages } = JSON.parse(body);
61
- __privateMethod(this, _AIChatAgent_instances, broadcastChatMessage_fn).call(this, {
62
- type: "cf_agent_chat_messages",
63
- messages
64
- }, [connection.id]);
54
+ this.broadcastChatMessage(
55
+ {
56
+ type: "cf_agent_chat_messages",
57
+ messages
58
+ },
59
+ [connection.id]
60
+ );
65
61
  await this.persistMessages(messages, [connection.id]);
66
62
  const chatMessageId = data.id;
67
- const abortSignal = __privateMethod(this, _AIChatAgent_instances, getAbortSignal_fn).call(this, chatMessageId);
68
- return __privateMethod(this, _AIChatAgent_instances, tryCatch_fn).call(this, async () => {
63
+ const abortSignal = this.getAbortSignal(chatMessageId);
64
+ return this.tryCatchChat(async () => {
69
65
  const response = await this.onChatMessage(
70
66
  async ({ response: response2 }) => {
71
67
  const finalMessages = appendResponseMessages({
@@ -73,31 +69,34 @@ var AIChatAgent = class extends Agent {
73
69
  responseMessages: response2.messages
74
70
  });
75
71
  await this.persistMessages(finalMessages, [connection.id]);
76
- __privateMethod(this, _AIChatAgent_instances, removeAbortController_fn).call(this, chatMessageId);
72
+ this.removeAbortController(chatMessageId);
77
73
  },
78
74
  abortSignal ? { abortSignal } : void 0
79
75
  );
80
76
  if (response) {
81
- await __privateMethod(this, _AIChatAgent_instances, reply_fn).call(this, data.id, response);
77
+ await this.reply(data.id, response);
82
78
  }
83
79
  });
84
80
  }
85
81
  if (data.type === "cf_agent_chat_clear") {
86
- __privateMethod(this, _AIChatAgent_instances, destroyAbortControllers_fn).call(this);
82
+ this.destroyAbortControllers();
87
83
  this.sql`delete from cf_ai_chat_agent_messages`;
88
84
  this.messages = [];
89
- __privateMethod(this, _AIChatAgent_instances, broadcastChatMessage_fn).call(this, {
90
- type: "cf_agent_chat_clear"
91
- }, [connection.id]);
85
+ this.broadcastChatMessage(
86
+ {
87
+ type: "cf_agent_chat_clear"
88
+ },
89
+ [connection.id]
90
+ );
92
91
  } else if (data.type === "cf_agent_chat_messages") {
93
92
  await this.persistMessages(data.messages, [connection.id]);
94
93
  } else if (data.type === "cf_agent_chat_request_cancel") {
95
- __privateMethod(this, _AIChatAgent_instances, cancelChatRequest_fn).call(this, data.id);
94
+ this.cancelChatRequest(data.id);
96
95
  }
97
96
  }
98
97
  }
99
98
  async onRequest(request) {
100
- return __privateMethod(this, _AIChatAgent_instances, tryCatch_fn).call(this, () => {
99
+ return this.tryCatchChat(() => {
101
100
  const url = new URL(request.url);
102
101
  if (url.pathname.endsWith("/get-messages")) {
103
102
  const messages = (this.sql`select * from cf_ai_chat_agent_messages` || []).map((row) => {
@@ -108,6 +107,13 @@ var AIChatAgent = class extends Agent {
108
107
  return super.onRequest(request);
109
108
  });
110
109
  }
110
+ async tryCatchChat(fn) {
111
+ try {
112
+ return await fn();
113
+ } catch (e) {
114
+ throw this.onError(e);
115
+ }
116
+ }
111
117
  /**
112
118
  * Handle incoming chat messages and generate a response
113
119
  * @param onFinish Callback to be called when the response is finished
@@ -145,88 +151,79 @@ var AIChatAgent = class extends Agent {
145
151
  this.sql`insert into cf_ai_chat_agent_messages (id, message) values (${message.id},${JSON.stringify(message)})`;
146
152
  }
147
153
  this.messages = messages;
148
- __privateMethod(this, _AIChatAgent_instances, broadcastChatMessage_fn).call(this, {
149
- type: "cf_agent_chat_messages",
150
- messages
151
- }, excludeBroadcastIds);
152
- }
153
- /**
154
- * When the DO is destroyed, cancel all pending requests
155
- */
156
- async destroy() {
157
- __privateMethod(this, _AIChatAgent_instances, destroyAbortControllers_fn).call(this);
158
- await super.destroy();
159
- }
160
- };
161
- _chatMessageAbortControllers = new WeakMap();
162
- _AIChatAgent_instances = new WeakSet();
163
- broadcastChatMessage_fn = function(message, exclude) {
164
- this.broadcast(JSON.stringify(message), exclude);
165
- };
166
- tryCatch_fn = async function(fn) {
167
- try {
168
- return await fn();
169
- } catch (e) {
170
- throw this.onError(e);
154
+ this.broadcastChatMessage(
155
+ {
156
+ type: "cf_agent_chat_messages",
157
+ messages
158
+ },
159
+ excludeBroadcastIds
160
+ );
171
161
  }
172
- };
173
- reply_fn = async function(id, response) {
174
- return __privateMethod(this, _AIChatAgent_instances, tryCatch_fn).call(this, async () => {
175
- for await (const chunk of response.body) {
176
- const body = decoder.decode(chunk);
177
- __privateMethod(this, _AIChatAgent_instances, broadcastChatMessage_fn).call(this, {
162
+ async reply(id, response) {
163
+ return this.tryCatchChat(async () => {
164
+ for await (const chunk of response.body) {
165
+ const body = decoder.decode(chunk);
166
+ this.broadcastChatMessage({
167
+ id,
168
+ type: "cf_agent_use_chat_response",
169
+ body,
170
+ done: false
171
+ });
172
+ }
173
+ this.broadcastChatMessage({
178
174
  id,
179
175
  type: "cf_agent_use_chat_response",
180
- body,
181
- done: false
176
+ body: "",
177
+ done: true
182
178
  });
183
- }
184
- __privateMethod(this, _AIChatAgent_instances, broadcastChatMessage_fn).call(this, {
185
- id,
186
- type: "cf_agent_use_chat_response",
187
- body: "",
188
- done: true
189
179
  });
190
- });
191
- };
192
- /**
193
- * For the given message id, look up its associated AbortController
194
- * If the AbortController does not exist, create and store one in memory
195
- *
196
- * returns the AbortSignal associated with the AbortController
197
- */
198
- getAbortSignal_fn = function(id) {
199
- if (typeof id !== "string") {
200
- return void 0;
201
180
  }
202
- if (!__privateGet(this, _chatMessageAbortControllers).has(id)) {
203
- __privateGet(this, _chatMessageAbortControllers).set(id, new AbortController());
181
+ /**
182
+ * For the given message id, look up its associated AbortController
183
+ * If the AbortController does not exist, create and store one in memory
184
+ *
185
+ * returns the AbortSignal associated with the AbortController
186
+ */
187
+ getAbortSignal(id) {
188
+ if (typeof id !== "string") {
189
+ return void 0;
190
+ }
191
+ if (!this.chatMessageAbortControllers.has(id)) {
192
+ this.chatMessageAbortControllers.set(id, new AbortController());
193
+ }
194
+ return this.chatMessageAbortControllers.get(id)?.signal;
204
195
  }
205
- return __privateGet(this, _chatMessageAbortControllers).get(id)?.signal;
206
- };
207
- /**
208
- * Remove an abort controller from the cache of pending message responses
209
- */
210
- removeAbortController_fn = function(id) {
211
- __privateGet(this, _chatMessageAbortControllers).delete(id);
212
- };
213
- /**
214
- * Propagate an abort signal for any requests associated with the given message id
215
- */
216
- cancelChatRequest_fn = function(id) {
217
- if (__privateGet(this, _chatMessageAbortControllers).has(id)) {
218
- const abortController = __privateGet(this, _chatMessageAbortControllers).get(id);
219
- abortController?.abort();
196
+ /**
197
+ * Remove an abort controller from the cache of pending message responses
198
+ */
199
+ removeAbortController(id) {
200
+ this.chatMessageAbortControllers.delete(id);
220
201
  }
221
- };
222
- /**
223
- * Abort all pending requests and clear the cache of AbortControllers
224
- */
225
- destroyAbortControllers_fn = function() {
226
- for (const controller of __privateGet(this, _chatMessageAbortControllers).values()) {
227
- controller?.abort();
202
+ /**
203
+ * Propagate an abort signal for any requests associated with the given message id
204
+ */
205
+ cancelChatRequest(id) {
206
+ if (this.chatMessageAbortControllers.has(id)) {
207
+ const abortController = this.chatMessageAbortControllers.get(id);
208
+ abortController?.abort();
209
+ }
210
+ }
211
+ /**
212
+ * Abort all pending requests and clear the cache of AbortControllers
213
+ */
214
+ destroyAbortControllers() {
215
+ for (const controller of this.chatMessageAbortControllers.values()) {
216
+ controller?.abort();
217
+ }
218
+ this.chatMessageAbortControllers.clear();
219
+ }
220
+ /**
221
+ * When the DO is destroyed, cancel all pending requests
222
+ */
223
+ async destroy() {
224
+ this.destroyAbortControllers();
225
+ await super.destroy();
228
226
  }
229
- __privateGet(this, _chatMessageAbortControllers).clear();
230
227
  };
231
228
  export {
232
229
  AIChatAgent
@@ -1 +1 @@
1
- {"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\";\n\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 /**\n * Map of message `id`s to `AbortController`s\n * useful to propagate request cancellation signals for any external calls made by the agent\n */\n #chatMessageAbortControllers: Map<string, AbortController>;\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 this.#chatMessageAbortControllers = new Map();\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\n const chatMessageId = data.id;\n const abortSignal = this.#getAbortSignal(chatMessageId);\n\n return this.#tryCatch(async () => {\n const response = await this.onChatMessage(\n async ({ response }) => {\n const finalMessages = appendResponseMessages({\n messages,\n responseMessages: response.messages,\n });\n\n await this.persistMessages(finalMessages, [connection.id]);\n this.#removeAbortController(chatMessageId);\n },\n abortSignal ? { abortSignal } : undefined\n );\n\n if (response) {\n await this.#reply(data.id, response);\n }\n });\n }\n if (data.type === \"cf_agent_chat_clear\") {\n this.#destroyAbortControllers();\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 } else if (data.type === \"cf_agent_chat_request_cancel\") {\n // propagate an abort signal for the associated request\n this.#cancelChatRequest(data.id);\n }\n }\n }\n\n override async onRequest(request: Request): Promise<Response> {\n return this.#tryCatch(() => {\n const url = new URL(request.url);\n if (url.pathname.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 * @param options.signal A signal to pass to any child requests which can be used to cancel them\n * @returns Response to send to the client or undefined\n */\n async onChatMessage(\n onFinish: StreamTextOnFinishCallback<ToolSet>,\n options?: { abortSignal: AbortSignal | undefined }\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 /**\n * For the given message id, look up its associated AbortController\n * If the AbortController does not exist, create and store one in memory\n *\n * returns the AbortSignal associated with the AbortController\n */\n #getAbortSignal(id: string): AbortSignal | undefined {\n // Defensive check, since we're coercing message types at the moment\n if (typeof id !== \"string\") {\n return undefined;\n }\n\n if (!this.#chatMessageAbortControllers.has(id)) {\n this.#chatMessageAbortControllers.set(id, new AbortController());\n }\n\n return this.#chatMessageAbortControllers.get(id)?.signal;\n }\n\n /**\n * Remove an abort controller from the cache of pending message responses\n */\n #removeAbortController(id: string) {\n this.#chatMessageAbortControllers.delete(id);\n }\n\n /**\n * Propagate an abort signal for any requests associated with the given message id\n */\n #cancelChatRequest(id: string) {\n if (this.#chatMessageAbortControllers.has(id)) {\n const abortController = this.#chatMessageAbortControllers.get(id);\n abortController?.abort();\n }\n }\n\n /**\n * Abort all pending requests and clear the cache of AbortControllers\n */\n #destroyAbortControllers() {\n for (const controller of this.#chatMessageAbortControllers.values()) {\n controller?.abort();\n }\n this.#chatMessageAbortControllers.clear();\n }\n\n /**\n * When the DO is destroyed, cancel all pending requests\n */\n async destroy() {\n this.#destroyAbortControllers();\n await super.destroy();\n }\n}\n"],"mappings":";;;;;;;;;;;;AAMA,SAAS,8BAA8B;AAGvC,IAAM,UAAU,IAAI,YAAY;AAThC;AAeO,IAAM,cAAN,cAA0D,MAG/D;AAAA,EAQA,YAAY,KAAmB,KAAU;AACvC,UAAM,KAAK,GAAG;AAZX;AAQL;AAAA;AAAA;AAAA;AAAA;AAKE,SAAK;AAAA;AAAA;AAAA;AAAA;AAKL,SAAK,YACH,KAAK,gDAAgD,CAAC,GACtD,IAAI,CAAC,QAAQ;AACb,aAAO,KAAK,MAAM,IAAI,OAAiB;AAAA,IACzC,CAAC;AAED,uBAAK,8BAA+B,oBAAI,IAAI;AAAA,EAC9C;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,KAAK,gBAAgB,UAAU,CAAC,WAAW,EAAE,CAAC;AAEpD,cAAM,gBAAgB,KAAK;AAC3B,cAAM,cAAc,sBAAK,2CAAL,WAAqB;AAEzC,eAAO,sBAAK,qCAAL,WAAe,YAAY;AAChC,gBAAM,WAAW,MAAM,KAAK;AAAA,YAC1B,OAAO,EAAE,UAAAA,UAAS,MAAM;AACtB,oBAAM,gBAAgB,uBAAuB;AAAA,gBAC3C;AAAA,gBACA,kBAAkBA,UAAS;AAAA,cAC7B,CAAC;AAED,oBAAM,KAAK,gBAAgB,eAAe,CAAC,WAAW,EAAE,CAAC;AACzD,oCAAK,kDAAL,WAA4B;AAAA,YAC9B;AAAA,YACA,cAAc,EAAE,YAAY,IAAI;AAAA,UAClC;AAEA,cAAI,UAAU;AACZ,kBAAM,sBAAK,kCAAL,WAAY,KAAK,IAAI;AAAA,UAC7B;AAAA,QACF;AAAA,MACF;AACA,UAAI,KAAK,SAAS,uBAAuB;AACvC,8BAAK,oDAAL;AACA,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,KAAK,gBAAgB,KAAK,UAAU,CAAC,WAAW,EAAE,CAAC;AAAA,MAC3D,WAAW,KAAK,SAAS,gCAAgC;AAEvD,8BAAK,8CAAL,WAAwB,KAAK;AAAA,MAC/B;AAAA,IACF;AAAA,EACF;AAAA,EAEA,MAAe,UAAU,SAAqC;AAC5D,WAAO,sBAAK,qCAAL,WAAe,MAAM;AAC1B,YAAM,MAAM,IAAI,IAAI,QAAQ,GAAG;AAC/B,UAAI,IAAI,SAAS,SAAS,eAAe,GAAG;AAC1C,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;AAAA,EAgBA,MAAM,cACJ,UACA,SAC+B;AAC/B,UAAM,IAAI;AAAA,MACR;AAAA,IACF;AAAA,EACF;AAAA;AAAA;AAAA;AAAA;AAAA,EAMA,MAAM,aAAa,UAAyB;AAC1C,UAAM,KAAK,gBAAgB,QAAQ;AACnC,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,KAAK,gBAAgB,eAAe,CAAC,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;AAAA,EAEA,MAAM,gBACJ,UACA,sBAAgC,CAAC,GACjC;AACA,SAAK;AACL,eAAW,WAAW,UAAU;AAC9B,WAAK,kEACH,QAAQ,EACV,IAAI,KAAK,UAAU,OAAO,CAAC;AAAA,IAC7B;AACA,SAAK,WAAW;AAChB,0BAAK,iDAAL,WACE;AAAA,MACE,MAAM;AAAA,MACN;AAAA,IACF,GACA;AAAA,EAEJ;AAAA;AAAA;AAAA;AAAA,EA2EA,MAAM,UAAU;AACd,0BAAK,oDAAL;AACA,UAAM,MAAM,QAAQ;AAAA,EACtB;AACF;AAvQE;AARK;AA2BL,0BAAqB,SAAC,SAA0B,SAAoB;AAClE,OAAK,UAAU,KAAK,UAAU,OAAO,GAAG,OAAO;AACjD;AAkGM,cAAY,eAAC,IAA0B;AAC3C,MAAI;AACF,WAAO,MAAM,GAAG;AAAA,EAClB,SAAS,GAAG;AACV,UAAM,KAAK,QAAQ,CAAC;AAAA,EACtB;AACF;AA6DM,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;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAQA,oBAAe,SAAC,IAAqC;AAEnD,MAAI,OAAO,OAAO,UAAU;AAC1B,WAAO;AAAA,EACT;AAEA,MAAI,CAAC,mBAAK,8BAA6B,IAAI,EAAE,GAAG;AAC9C,uBAAK,8BAA6B,IAAI,IAAI,IAAI,gBAAgB,CAAC;AAAA,EACjE;AAEA,SAAO,mBAAK,8BAA6B,IAAI,EAAE,GAAG;AACpD;AAAA;AAAA;AAAA;AAKA,2BAAsB,SAAC,IAAY;AACjC,qBAAK,8BAA6B,OAAO,EAAE;AAC7C;AAAA;AAAA;AAAA;AAKA,uBAAkB,SAAC,IAAY;AAC7B,MAAI,mBAAK,8BAA6B,IAAI,EAAE,GAAG;AAC7C,UAAM,kBAAkB,mBAAK,8BAA6B,IAAI,EAAE;AAChE,qBAAiB,MAAM;AAAA,EACzB;AACF;AAAA;AAAA;AAAA;AAKA,6BAAwB,WAAG;AACzB,aAAW,cAAc,mBAAK,8BAA6B,OAAO,GAAG;AACnE,gBAAY,MAAM;AAAA,EACpB;AACA,qBAAK,8BAA6B,MAAM;AAC1C;","names":["response"]}
1
+ {"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\";\n\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 /**\n * Map of message `id`s to `AbortController`s\n * useful to propagate request cancellation signals for any external calls made by the agent\n */\n private chatMessageAbortControllers: Map<string, AbortController>;\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 this.chatMessageAbortControllers = new Map();\n }\n\n private 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\n const chatMessageId = data.id;\n const abortSignal = this.getAbortSignal(chatMessageId);\n\n return this.tryCatchChat(async () => {\n const response = await this.onChatMessage(\n async ({ response }) => {\n const finalMessages = appendResponseMessages({\n messages,\n responseMessages: response.messages,\n });\n\n await this.persistMessages(finalMessages, [connection.id]);\n this.removeAbortController(chatMessageId);\n },\n abortSignal ? { abortSignal } : undefined\n );\n\n if (response) {\n await this.reply(data.id, response);\n }\n });\n }\n if (data.type === \"cf_agent_chat_clear\") {\n this.destroyAbortControllers();\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 } else if (data.type === \"cf_agent_chat_request_cancel\") {\n // propagate an abort signal for the associated request\n this.cancelChatRequest(data.id);\n }\n }\n }\n\n override async onRequest(request: Request): Promise<Response> {\n return this.tryCatchChat(() => {\n const url = new URL(request.url);\n if (url.pathname.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 private async tryCatchChat<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 * @param options.signal A signal to pass to any child requests which can be used to cancel them\n * @returns Response to send to the client or undefined\n */\n async onChatMessage(\n onFinish: StreamTextOnFinishCallback<ToolSet>,\n options?: { abortSignal: AbortSignal | undefined }\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 private async reply(id: string, response: Response) {\n // now take chunks out from dataStreamResponse and send them to the client\n return this.tryCatchChat(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 /**\n * For the given message id, look up its associated AbortController\n * If the AbortController does not exist, create and store one in memory\n *\n * returns the AbortSignal associated with the AbortController\n */\n private getAbortSignal(id: string): AbortSignal | undefined {\n // Defensive check, since we're coercing message types at the moment\n if (typeof id !== \"string\") {\n return undefined;\n }\n\n if (!this.chatMessageAbortControllers.has(id)) {\n this.chatMessageAbortControllers.set(id, new AbortController());\n }\n\n return this.chatMessageAbortControllers.get(id)?.signal;\n }\n\n /**\n * Remove an abort controller from the cache of pending message responses\n */\n private removeAbortController(id: string) {\n this.chatMessageAbortControllers.delete(id);\n }\n\n /**\n * Propagate an abort signal for any requests associated with the given message id\n */\n private cancelChatRequest(id: string) {\n if (this.chatMessageAbortControllers.has(id)) {\n const abortController = this.chatMessageAbortControllers.get(id);\n abortController?.abort();\n }\n }\n\n /**\n * Abort all pending requests and clear the cache of AbortControllers\n */\n private destroyAbortControllers() {\n for (const controller of this.chatMessageAbortControllers.values()) {\n controller?.abort();\n }\n this.chatMessageAbortControllers.clear();\n }\n\n /**\n * When the DO is destroyed, cancel all pending requests\n */\n async destroy() {\n this.destroyAbortControllers();\n await super.destroy();\n }\n}\n"],"mappings":";;;;;;;;;AAMA,SAAS,8BAA8B;AAGvC,IAAM,UAAU,IAAI,YAAY;AAMzB,IAAM,cAAN,cAA0D,MAG/D;AAAA,EAQA,YAAY,KAAmB,KAAU;AACvC,UAAM,KAAK,GAAG;AACd,SAAK;AAAA;AAAA;AAAA;AAAA;AAKL,SAAK,YACH,KAAK,gDAAgD,CAAC,GACtD,IAAI,CAAC,QAAQ;AACb,aAAO,KAAK,MAAM,IAAI,OAAiB;AAAA,IACzC,CAAC;AAED,SAAK,8BAA8B,oBAAI,IAAI;AAAA,EAC7C;AAAA,EAEQ,qBAAqB,SAA0B,SAAoB;AACzE,SAAK,UAAU,KAAK,UAAU,OAAO,GAAG,OAAO;AAAA,EACjD;AAAA,EAEA,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,aAAK;AAAA,UACH;AAAA,YACE,MAAM;AAAA,YACN;AAAA,UACF;AAAA,UACA,CAAC,WAAW,EAAE;AAAA,QAChB;AACA,cAAM,KAAK,gBAAgB,UAAU,CAAC,WAAW,EAAE,CAAC;AAEpD,cAAM,gBAAgB,KAAK;AAC3B,cAAM,cAAc,KAAK,eAAe,aAAa;AAErD,eAAO,KAAK,aAAa,YAAY;AACnC,gBAAM,WAAW,MAAM,KAAK;AAAA,YAC1B,OAAO,EAAE,UAAAA,UAAS,MAAM;AACtB,oBAAM,gBAAgB,uBAAuB;AAAA,gBAC3C;AAAA,gBACA,kBAAkBA,UAAS;AAAA,cAC7B,CAAC;AAED,oBAAM,KAAK,gBAAgB,eAAe,CAAC,WAAW,EAAE,CAAC;AACzD,mBAAK,sBAAsB,aAAa;AAAA,YAC1C;AAAA,YACA,cAAc,EAAE,YAAY,IAAI;AAAA,UAClC;AAEA,cAAI,UAAU;AACZ,kBAAM,KAAK,MAAM,KAAK,IAAI,QAAQ;AAAA,UACpC;AAAA,QACF,CAAC;AAAA,MACH;AACA,UAAI,KAAK,SAAS,uBAAuB;AACvC,aAAK,wBAAwB;AAC7B,aAAK;AACL,aAAK,WAAW,CAAC;AACjB,aAAK;AAAA,UACH;AAAA,YACE,MAAM;AAAA,UACR;AAAA,UACA,CAAC,WAAW,EAAE;AAAA,QAChB;AAAA,MACF,WAAW,KAAK,SAAS,0BAA0B;AAEjD,cAAM,KAAK,gBAAgB,KAAK,UAAU,CAAC,WAAW,EAAE,CAAC;AAAA,MAC3D,WAAW,KAAK,SAAS,gCAAgC;AAEvD,aAAK,kBAAkB,KAAK,EAAE;AAAA,MAChC;AAAA,IACF;AAAA,EACF;AAAA,EAEA,MAAe,UAAU,SAAqC;AAC5D,WAAO,KAAK,aAAa,MAAM;AAC7B,YAAM,MAAM,IAAI,IAAI,QAAQ,GAAG;AAC/B,UAAI,IAAI,SAAS,SAAS,eAAe,GAAG;AAC1C,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,CAAC;AAAA,EACH;AAAA,EAEA,MAAc,aAAgB,IAA0B;AACtD,QAAI;AACF,aAAO,MAAM,GAAG;AAAA,IAClB,SAAS,GAAG;AACV,YAAM,KAAK,QAAQ,CAAC;AAAA,IACtB;AAAA,EACF;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAQA,MAAM,cACJ,UACA,SAC+B;AAC/B,UAAM,IAAI;AAAA,MACR;AAAA,IACF;AAAA,EACF;AAAA;AAAA;AAAA;AAAA;AAAA,EAMA,MAAM,aAAa,UAAyB;AAC1C,UAAM,KAAK,gBAAgB,QAAQ;AACnC,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,KAAK,gBAAgB,eAAe,CAAC,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;AAAA,EAEA,MAAM,gBACJ,UACA,sBAAgC,CAAC,GACjC;AACA,SAAK;AACL,eAAW,WAAW,UAAU;AAC9B,WAAK,kEACH,QAAQ,EACV,IAAI,KAAK,UAAU,OAAO,CAAC;AAAA,IAC7B;AACA,SAAK,WAAW;AAChB,SAAK;AAAA,MACH;AAAA,QACE,MAAM;AAAA,QACN;AAAA,MACF;AAAA,MACA;AAAA,IACF;AAAA,EACF;AAAA,EAEA,MAAc,MAAM,IAAY,UAAoB;AAElD,WAAO,KAAK,aAAa,YAAY;AAEnC,uBAAiB,SAAS,SAAS,MAAO;AACxC,cAAM,OAAO,QAAQ,OAAO,KAAK;AAEjC,aAAK,qBAAqB;AAAA,UACxB;AAAA,UACA,MAAM;AAAA,UACN;AAAA,UACA,MAAM;AAAA,QACR,CAAC;AAAA,MACH;AAEA,WAAK,qBAAqB;AAAA,QACxB;AAAA,QACA,MAAM;AAAA,QACN,MAAM;AAAA,QACN,MAAM;AAAA,MACR,CAAC;AAAA,IACH,CAAC;AAAA,EACH;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAQQ,eAAe,IAAqC;AAE1D,QAAI,OAAO,OAAO,UAAU;AAC1B,aAAO;AAAA,IACT;AAEA,QAAI,CAAC,KAAK,4BAA4B,IAAI,EAAE,GAAG;AAC7C,WAAK,4BAA4B,IAAI,IAAI,IAAI,gBAAgB,CAAC;AAAA,IAChE;AAEA,WAAO,KAAK,4BAA4B,IAAI,EAAE,GAAG;AAAA,EACnD;AAAA;AAAA;AAAA;AAAA,EAKQ,sBAAsB,IAAY;AACxC,SAAK,4BAA4B,OAAO,EAAE;AAAA,EAC5C;AAAA;AAAA;AAAA;AAAA,EAKQ,kBAAkB,IAAY;AACpC,QAAI,KAAK,4BAA4B,IAAI,EAAE,GAAG;AAC5C,YAAM,kBAAkB,KAAK,4BAA4B,IAAI,EAAE;AAC/D,uBAAiB,MAAM;AAAA,IACzB;AAAA,EACF;AAAA;AAAA;AAAA;AAAA,EAKQ,0BAA0B;AAChC,eAAW,cAAc,KAAK,4BAA4B,OAAO,GAAG;AAClE,kBAAY,MAAM;AAAA,IACpB;AACA,SAAK,4BAA4B,MAAM;AAAA,EACzC;AAAA;AAAA;AAAA;AAAA,EAKA,MAAM,UAAU;AACd,SAAK,wBAAwB;AAC7B,UAAM,MAAM,QAAQ;AAAA,EACtB;AACF;","names":["response"]}
@@ -4,6 +4,17 @@ import { useChat } from "@ai-sdk/react";
4
4
  import { useAgent } from "./react.js";
5
5
  import "partysocket";
6
6
  import "partysocket/react";
7
+ import "./index.js";
8
+ import "partyserver";
9
+ import "./mcp/client.js";
10
+ import "zod";
11
+ import "@modelcontextprotocol/sdk/types.js";
12
+ import "@modelcontextprotocol/sdk/client/index.js";
13
+ import "@modelcontextprotocol/sdk/client/sse.js";
14
+ import "./mcp/do-oauth-client-provider.js";
15
+ import "@modelcontextprotocol/sdk/client/auth.js";
16
+ import "@modelcontextprotocol/sdk/shared/auth.js";
17
+ import "@modelcontextprotocol/sdk/shared/protocol.js";
7
18
  import "./client.js";
8
19
 
9
20
  type GetInitialMessagesOptions = {
@@ -52,6 +63,7 @@ declare function useAgentChat<State = unknown>(
52
63
  chatRequestOptions?: ai.ChatRequestOptions
53
64
  ) => Promise<string | null | undefined>;
54
65
  stop: () => void;
66
+ experimental_resume: () => void;
55
67
  input: string;
56
68
  setInput: React.Dispatch<React.SetStateAction<string>>;
57
69
  handleInputChange: (
package/dist/ai-react.js CHANGED
@@ -1,8 +1,9 @@
1
- import "./chunk-HMLY7DHA.js";
1
+ import "./chunk-NOUFNU2O.js";
2
2
 
3
3
  // src/ai-react.tsx
4
4
  import { useChat } from "@ai-sdk/react";
5
5
  import { use, useEffect } from "react";
6
+ import { nanoid } from "nanoid";
6
7
  var requestCache = /* @__PURE__ */ new Map();
7
8
  function useAgentChat(options) {
8
9
  const { agent, getInitialMessages, ...rest } = options;
@@ -65,7 +66,7 @@ function useAgentChat(options) {
65
66
  window
66
67
  // dispatcher, duplex
67
68
  } = options2;
68
- const id = crypto.randomUUID();
69
+ const id = nanoid(8);
69
70
  const abortController = new AbortController();
70
71
  signal?.addEventListener("abort", () => {
71
72
  agent.send(
@@ -75,6 +76,7 @@ function useAgentChat(options) {
75
76
  })
76
77
  );
77
78
  abortController.abort();
79
+ controller.close();
78
80
  });
79
81
  agent.addEventListener(
80
82
  "message",
@@ -1 +1 @@
1
- {"version":3,"sources":["../src/ai-react.tsx"],"sourcesContent":["import { useChat } from \"@ai-sdk/react\";\nimport type { Message } from \"ai\";\nimport { use, useEffect } from \"react\";\nimport type { OutgoingMessage } from \"./ai-types\";\nimport type { useAgent } from \"./react\";\n\ntype GetInitialMessagesOptions = {\n agent: string;\n name: string;\n url: string;\n};\n\n/**\n * Options for the useAgentChat hook\n */\ntype UseAgentChatOptions<State> = Omit<\n Parameters<typeof useChat>[0] & {\n /** Agent connection from useAgent */\n agent: ReturnType<typeof useAgent<State>>;\n getInitialMessages?:\n | undefined\n | null\n // | (() => Message[])\n | ((options: GetInitialMessagesOptions) => Promise<Message[]>);\n },\n \"fetch\"\n>;\n\nconst requestCache = new Map<string, Promise<Message[]>>();\n\n/**\n * React hook for building AI chat interfaces using an Agent\n * @param options Chat options including the agent connection\n * @returns Chat interface controls and state with added clearHistory method\n */\nexport function useAgentChat<State = unknown>(\n options: UseAgentChatOptions<State>\n) {\n const { agent, getInitialMessages, ...rest } = options;\n\n const agentUrl = new URL(\n `${// @ts-expect-error we're using a protected _url property that includes query params\n ((agent._url as string | null) || agent._pkurl)\n ?.replace(\"ws://\", \"http://\")\n .replace(\"wss://\", \"https://\")}`\n );\n\n // delete the _pk query param\n agentUrl.searchParams.delete(\"_pk\");\n const agentUrlString = agentUrl.toString();\n\n async function defaultGetInitialMessagesFetch({\n url,\n }: GetInitialMessagesOptions) {\n const getMessagesUrl = new URL(url);\n getMessagesUrl.pathname += \"/get-messages\";\n const response = await fetch(getMessagesUrl.toString(), {\n headers: options.headers,\n credentials: options.credentials,\n });\n return response.json<Message[]>();\n }\n\n const getInitialMessagesFetch =\n getInitialMessages || defaultGetInitialMessagesFetch;\n\n function doGetInitialMessages(\n getInitialMessagesOptions: GetInitialMessagesOptions\n ) {\n if (requestCache.has(agentUrlString)) {\n return requestCache.get(agentUrlString)!;\n }\n const promise = getInitialMessagesFetch(getInitialMessagesOptions);\n // immediately cache the promise so that we don't\n // create multiple requests for the same agent during multiple\n // concurrent renders\n requestCache.set(agentUrlString, promise);\n return promise;\n }\n\n const initialMessagesPromise =\n getInitialMessages === null\n ? null\n : doGetInitialMessages({\n agent: agent.agent,\n name: agent.name,\n url: agentUrlString,\n });\n const initialMessages = initialMessagesPromise\n ? use(initialMessagesPromise)\n : (rest.initialMessages ?? []);\n\n // manages adding and removing the promise from the cache\n useEffect(() => {\n if (!initialMessagesPromise) {\n return;\n }\n // this effect is responsible for removing the promise from the cache\n // when the component unmounts or the promise changes,\n // but that means it also must add the promise to the cache\n // so that multiple arbitrary effect runs produce the expected state\n // when resolved.\n requestCache.set(agentUrlString, initialMessagesPromise!);\n return () => {\n if (requestCache.get(agentUrlString) === initialMessagesPromise) {\n requestCache.delete(agentUrlString);\n }\n };\n }, [agentUrlString, initialMessagesPromise]);\n\n async function aiFetch(\n request: RequestInfo | URL,\n options: RequestInit = {}\n ) {\n // we're going to use a websocket to do the actual \"fetching\"\n // but still satisfy the type signature of the fetch function\n // so we'll return a promise that resolves to a response\n\n const {\n method,\n keepalive,\n headers,\n body,\n redirect,\n integrity,\n signal,\n credentials,\n mode,\n referrer,\n referrerPolicy,\n window,\n // dispatcher, duplex\n } = options;\n const id = crypto.randomUUID();\n const abortController = new AbortController();\n\n signal?.addEventListener(\"abort\", () => {\n // Propagate request cancellation to the Agent\n // We need to communciate cancellation as a websocket message, instead of a request signal\n agent.send(\n JSON.stringify({\n type: \"cf_agent_chat_request_cancel\",\n id,\n })\n );\n\n // NOTE - If we wanted to, we could preserve the \"interrupted\" message here, with the code below\n // However, I think it might be the responsibility of the library user to implement that behavior manually?\n // Reasoning: This code could be subject to collisions, as it \"force saves\" the messages we have locally\n //\n // agent.send(JSON.stringify({\n // type: \"cf_agent_chat_messages\",\n // messages: ... /* some way of getting current messages ref? */\n // }))\n\n abortController.abort();\n });\n\n agent.addEventListener(\n \"message\",\n (event) => {\n let data: OutgoingMessage;\n try {\n data = JSON.parse(event.data) as OutgoingMessage;\n } catch (error) {\n // silently ignore invalid messages for now\n // TODO: log errors with log levels\n return;\n }\n if (data.type === \"cf_agent_use_chat_response\") {\n if (data.id === id) {\n controller.enqueue(new TextEncoder().encode(data.body));\n if (data.done) {\n controller.close();\n abortController.abort();\n }\n }\n }\n },\n { signal: abortController.signal }\n );\n\n let controller: ReadableStreamDefaultController;\n\n const stream = new ReadableStream({\n start(c) {\n controller = c;\n },\n });\n\n agent.send(\n JSON.stringify({\n type: \"cf_agent_use_chat_request\",\n id,\n url: request.toString(),\n init: {\n method,\n keepalive,\n headers,\n body,\n redirect,\n integrity,\n credentials,\n mode,\n referrer,\n referrerPolicy,\n window,\n // dispatcher,\n // duplex\n },\n })\n );\n\n return new Response(stream);\n }\n const useChatHelpers = useChat({\n initialMessages,\n sendExtraMessageFields: true,\n fetch: aiFetch,\n ...rest,\n });\n\n useEffect(() => {\n function onClearHistory(event: MessageEvent) {\n if (typeof event.data !== \"string\") {\n return;\n }\n let data: OutgoingMessage;\n try {\n data = JSON.parse(event.data) as OutgoingMessage;\n } catch (error) {\n // silently ignore invalid messages for now\n // TODO: log errors with log levels\n return;\n }\n if (data.type === \"cf_agent_chat_clear\") {\n useChatHelpers.setMessages([]);\n }\n }\n\n function onMessages(event: MessageEvent) {\n if (typeof event.data !== \"string\") {\n return;\n }\n let data: OutgoingMessage;\n try {\n data = JSON.parse(event.data) as OutgoingMessage;\n } catch (error) {\n // silently ignore invalid messages for now\n // TODO: log errors with log levels\n return;\n }\n if (data.type === \"cf_agent_chat_messages\") {\n useChatHelpers.setMessages(data.messages);\n }\n }\n\n agent.addEventListener(\"message\", onClearHistory);\n agent.addEventListener(\"message\", onMessages);\n\n return () => {\n agent.removeEventListener(\"message\", onClearHistory);\n agent.removeEventListener(\"message\", onMessages);\n };\n }, [agent, useChatHelpers.setMessages]);\n\n return {\n ...useChatHelpers,\n /**\n * Set the chat messages and synchronize with the Agent\n * @param messages New messages to set\n */\n setMessages: (messages: Message[]) => {\n useChatHelpers.setMessages(messages);\n agent.send(\n JSON.stringify({\n type: \"cf_agent_chat_messages\",\n messages,\n })\n );\n },\n /**\n * Clear chat history on both client and Agent\n */\n clearHistory: () => {\n useChatHelpers.setMessages([]);\n agent.send(\n JSON.stringify({\n type: \"cf_agent_chat_clear\",\n })\n );\n },\n };\n}\n"],"mappings":";;;AAAA,SAAS,eAAe;AAExB,SAAS,KAAK,iBAAiB;AA0B/B,IAAM,eAAe,oBAAI,IAAgC;AAOlD,SAAS,aACd,SACA;AACA,QAAM,EAAE,OAAO,oBAAoB,GAAG,KAAK,IAAI;AAE/C,QAAM,WAAW,IAAI;AAAA,IACnB;AAAA,KACE,MAAM,QAA0B,MAAM,SACpC,QAAQ,SAAS,SAAS,EAC3B,QAAQ,UAAU,UAAU,CAAC;AAAA,EAClC;AAGA,WAAS,aAAa,OAAO,KAAK;AAClC,QAAM,iBAAiB,SAAS,SAAS;AAEzC,iBAAe,+BAA+B;AAAA,IAC5C;AAAA,EACF,GAA8B;AAC5B,UAAM,iBAAiB,IAAI,IAAI,GAAG;AAClC,mBAAe,YAAY;AAC3B,UAAM,WAAW,MAAM,MAAM,eAAe,SAAS,GAAG;AAAA,MACtD,SAAS,QAAQ;AAAA,MACjB,aAAa,QAAQ;AAAA,IACvB,CAAC;AACD,WAAO,SAAS,KAAgB;AAAA,EAClC;AAEA,QAAM,0BACJ,sBAAsB;AAExB,WAAS,qBACP,2BACA;AACA,QAAI,aAAa,IAAI,cAAc,GAAG;AACpC,aAAO,aAAa,IAAI,cAAc;AAAA,IACxC;AACA,UAAM,UAAU,wBAAwB,yBAAyB;AAIjE,iBAAa,IAAI,gBAAgB,OAAO;AACxC,WAAO;AAAA,EACT;AAEA,QAAM,yBACJ,uBAAuB,OACnB,OACA,qBAAqB;AAAA,IACnB,OAAO,MAAM;AAAA,IACb,MAAM,MAAM;AAAA,IACZ,KAAK;AAAA,EACP,CAAC;AACP,QAAM,kBAAkB,yBACpB,IAAI,sBAAsB,IACzB,KAAK,mBAAmB,CAAC;AAG9B,YAAU,MAAM;AACd,QAAI,CAAC,wBAAwB;AAC3B;AAAA,IACF;AAMA,iBAAa,IAAI,gBAAgB,sBAAuB;AACxD,WAAO,MAAM;AACX,UAAI,aAAa,IAAI,cAAc,MAAM,wBAAwB;AAC/D,qBAAa,OAAO,cAAc;AAAA,MACpC;AAAA,IACF;AAAA,EACF,GAAG,CAAC,gBAAgB,sBAAsB,CAAC;AAE3C,iBAAe,QACb,SACAA,WAAuB,CAAC,GACxB;AAKA,UAAM;AAAA,MACJ;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA;AAAA,IAEF,IAAIA;AACJ,UAAM,KAAK,OAAO,WAAW;AAC7B,UAAM,kBAAkB,IAAI,gBAAgB;AAE5C,YAAQ,iBAAiB,SAAS,MAAM;AAGtC,YAAM;AAAA,QACJ,KAAK,UAAU;AAAA,UACb,MAAM;AAAA,UACN;AAAA,QACF,CAAC;AAAA,MACH;AAWA,sBAAgB,MAAM;AAAA,IACxB,CAAC;AAED,UAAM;AAAA,MACJ;AAAA,MACA,CAAC,UAAU;AACT,YAAI;AACJ,YAAI;AACF,iBAAO,KAAK,MAAM,MAAM,IAAI;AAAA,QAC9B,SAAS,OAAO;AAGd;AAAA,QACF;AACA,YAAI,KAAK,SAAS,8BAA8B;AAC9C,cAAI,KAAK,OAAO,IAAI;AAClB,uBAAW,QAAQ,IAAI,YAAY,EAAE,OAAO,KAAK,IAAI,CAAC;AACtD,gBAAI,KAAK,MAAM;AACb,yBAAW,MAAM;AACjB,8BAAgB,MAAM;AAAA,YACxB;AAAA,UACF;AAAA,QACF;AAAA,MACF;AAAA,MACA,EAAE,QAAQ,gBAAgB,OAAO;AAAA,IACnC;AAEA,QAAI;AAEJ,UAAM,SAAS,IAAI,eAAe;AAAA,MAChC,MAAM,GAAG;AACP,qBAAa;AAAA,MACf;AAAA,IACF,CAAC;AAED,UAAM;AAAA,MACJ,KAAK,UAAU;AAAA,QACb,MAAM;AAAA,QACN;AAAA,QACA,KAAK,QAAQ,SAAS;AAAA,QACtB,MAAM;AAAA,UACJ;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA;AAAA;AAAA,QAGF;AAAA,MACF,CAAC;AAAA,IACH;AAEA,WAAO,IAAI,SAAS,MAAM;AAAA,EAC5B;AACA,QAAM,iBAAiB,QAAQ;AAAA,IAC7B;AAAA,IACA,wBAAwB;AAAA,IACxB,OAAO;AAAA,IACP,GAAG;AAAA,EACL,CAAC;AAED,YAAU,MAAM;AACd,aAAS,eAAe,OAAqB;AAC3C,UAAI,OAAO,MAAM,SAAS,UAAU;AAClC;AAAA,MACF;AACA,UAAI;AACJ,UAAI;AACF,eAAO,KAAK,MAAM,MAAM,IAAI;AAAA,MAC9B,SAAS,OAAO;AAGd;AAAA,MACF;AACA,UAAI,KAAK,SAAS,uBAAuB;AACvC,uBAAe,YAAY,CAAC,CAAC;AAAA,MAC/B;AAAA,IACF;AAEA,aAAS,WAAW,OAAqB;AACvC,UAAI,OAAO,MAAM,SAAS,UAAU;AAClC;AAAA,MACF;AACA,UAAI;AACJ,UAAI;AACF,eAAO,KAAK,MAAM,MAAM,IAAI;AAAA,MAC9B,SAAS,OAAO;AAGd;AAAA,MACF;AACA,UAAI,KAAK,SAAS,0BAA0B;AAC1C,uBAAe,YAAY,KAAK,QAAQ;AAAA,MAC1C;AAAA,IACF;AAEA,UAAM,iBAAiB,WAAW,cAAc;AAChD,UAAM,iBAAiB,WAAW,UAAU;AAE5C,WAAO,MAAM;AACX,YAAM,oBAAoB,WAAW,cAAc;AACnD,YAAM,oBAAoB,WAAW,UAAU;AAAA,IACjD;AAAA,EACF,GAAG,CAAC,OAAO,eAAe,WAAW,CAAC;AAEtC,SAAO;AAAA,IACL,GAAG;AAAA;AAAA;AAAA;AAAA;AAAA,IAKH,aAAa,CAAC,aAAwB;AACpC,qBAAe,YAAY,QAAQ;AACnC,YAAM;AAAA,QACJ,KAAK,UAAU;AAAA,UACb,MAAM;AAAA,UACN;AAAA,QACF,CAAC;AAAA,MACH;AAAA,IACF;AAAA;AAAA;AAAA;AAAA,IAIA,cAAc,MAAM;AAClB,qBAAe,YAAY,CAAC,CAAC;AAC7B,YAAM;AAAA,QACJ,KAAK,UAAU;AAAA,UACb,MAAM;AAAA,QACR,CAAC;AAAA,MACH;AAAA,IACF;AAAA,EACF;AACF;","names":["options"]}
1
+ {"version":3,"sources":["../src/ai-react.tsx"],"sourcesContent":["import { useChat } from \"@ai-sdk/react\";\nimport type { Message } from \"ai\";\nimport { use, useEffect } from \"react\";\nimport type { OutgoingMessage } from \"./ai-types\";\nimport type { useAgent } from \"./react\";\nimport { nanoid } from \"nanoid\";\n\ntype GetInitialMessagesOptions = {\n agent: string;\n name: string;\n url: string;\n};\n\n/**\n * Options for the useAgentChat hook\n */\ntype UseAgentChatOptions<State> = Omit<\n Parameters<typeof useChat>[0] & {\n /** Agent connection from useAgent */\n agent: ReturnType<typeof useAgent<State>>;\n getInitialMessages?:\n | undefined\n | null\n // | (() => Message[])\n | ((options: GetInitialMessagesOptions) => Promise<Message[]>);\n },\n \"fetch\"\n>;\n\nconst requestCache = new Map<string, Promise<Message[]>>();\n\n/**\n * React hook for building AI chat interfaces using an Agent\n * @param options Chat options including the agent connection\n * @returns Chat interface controls and state with added clearHistory method\n */\nexport function useAgentChat<State = unknown>(\n options: UseAgentChatOptions<State>\n) {\n const { agent, getInitialMessages, ...rest } = options;\n\n const agentUrl = new URL(\n `${// @ts-expect-error we're using a protected _url property that includes query params\n ((agent._url as string | null) || agent._pkurl)\n ?.replace(\"ws://\", \"http://\")\n .replace(\"wss://\", \"https://\")}`\n );\n\n // delete the _pk query param\n agentUrl.searchParams.delete(\"_pk\");\n const agentUrlString = agentUrl.toString();\n\n async function defaultGetInitialMessagesFetch({\n url,\n }: GetInitialMessagesOptions) {\n const getMessagesUrl = new URL(url);\n getMessagesUrl.pathname += \"/get-messages\";\n const response = await fetch(getMessagesUrl.toString(), {\n headers: options.headers,\n credentials: options.credentials,\n });\n return response.json<Message[]>();\n }\n\n const getInitialMessagesFetch =\n getInitialMessages || defaultGetInitialMessagesFetch;\n\n function doGetInitialMessages(\n getInitialMessagesOptions: GetInitialMessagesOptions\n ) {\n if (requestCache.has(agentUrlString)) {\n return requestCache.get(agentUrlString)!;\n }\n const promise = getInitialMessagesFetch(getInitialMessagesOptions);\n // immediately cache the promise so that we don't\n // create multiple requests for the same agent during multiple\n // concurrent renders\n requestCache.set(agentUrlString, promise);\n return promise;\n }\n\n const initialMessagesPromise =\n getInitialMessages === null\n ? null\n : doGetInitialMessages({\n agent: agent.agent,\n name: agent.name,\n url: agentUrlString,\n });\n const initialMessages = initialMessagesPromise\n ? use(initialMessagesPromise)\n : (rest.initialMessages ?? []);\n\n // manages adding and removing the promise from the cache\n useEffect(() => {\n if (!initialMessagesPromise) {\n return;\n }\n // this effect is responsible for removing the promise from the cache\n // when the component unmounts or the promise changes,\n // but that means it also must add the promise to the cache\n // so that multiple arbitrary effect runs produce the expected state\n // when resolved.\n requestCache.set(agentUrlString, initialMessagesPromise!);\n return () => {\n if (requestCache.get(agentUrlString) === initialMessagesPromise) {\n requestCache.delete(agentUrlString);\n }\n };\n }, [agentUrlString, initialMessagesPromise]);\n\n async function aiFetch(\n request: RequestInfo | URL,\n options: RequestInit = {}\n ) {\n // we're going to use a websocket to do the actual \"fetching\"\n // but still satisfy the type signature of the fetch function\n // so we'll return a promise that resolves to a response\n\n const {\n method,\n keepalive,\n headers,\n body,\n redirect,\n integrity,\n signal,\n credentials,\n mode,\n referrer,\n referrerPolicy,\n window,\n // dispatcher, duplex\n } = options;\n const id = nanoid(8);\n const abortController = new AbortController();\n\n signal?.addEventListener(\"abort\", () => {\n // Propagate request cancellation to the Agent\n // We need to communciate cancellation as a websocket message, instead of a request signal\n agent.send(\n JSON.stringify({\n type: \"cf_agent_chat_request_cancel\",\n id,\n })\n );\n\n // NOTE - If we wanted to, we could preserve the \"interrupted\" message here, with the code below\n // However, I think it might be the responsibility of the library user to implement that behavior manually?\n // Reasoning: This code could be subject to collisions, as it \"force saves\" the messages we have locally\n //\n // agent.send(JSON.stringify({\n // type: \"cf_agent_chat_messages\",\n // messages: ... /* some way of getting current messages ref? */\n // }))\n\n abortController.abort();\n // Make sure to also close the stream (cf. https://github.com/cloudflare/agents-starter/issues/69)\n controller.close();\n });\n\n agent.addEventListener(\n \"message\",\n (event) => {\n let data: OutgoingMessage;\n try {\n data = JSON.parse(event.data) as OutgoingMessage;\n } catch (error) {\n // silently ignore invalid messages for now\n // TODO: log errors with log levels\n return;\n }\n if (data.type === \"cf_agent_use_chat_response\") {\n if (data.id === id) {\n controller.enqueue(new TextEncoder().encode(data.body));\n if (data.done) {\n controller.close();\n abortController.abort();\n }\n }\n }\n },\n { signal: abortController.signal }\n );\n\n let controller: ReadableStreamDefaultController;\n\n const stream = new ReadableStream({\n start(c) {\n controller = c;\n },\n });\n\n agent.send(\n JSON.stringify({\n type: \"cf_agent_use_chat_request\",\n id,\n url: request.toString(),\n init: {\n method,\n keepalive,\n headers,\n body,\n redirect,\n integrity,\n credentials,\n mode,\n referrer,\n referrerPolicy,\n window,\n // dispatcher,\n // duplex\n },\n })\n );\n\n return new Response(stream);\n }\n const useChatHelpers = useChat({\n initialMessages,\n sendExtraMessageFields: true,\n fetch: aiFetch,\n ...rest,\n });\n\n useEffect(() => {\n function onClearHistory(event: MessageEvent) {\n if (typeof event.data !== \"string\") {\n return;\n }\n let data: OutgoingMessage;\n try {\n data = JSON.parse(event.data) as OutgoingMessage;\n } catch (error) {\n // silently ignore invalid messages for now\n // TODO: log errors with log levels\n return;\n }\n if (data.type === \"cf_agent_chat_clear\") {\n useChatHelpers.setMessages([]);\n }\n }\n\n function onMessages(event: MessageEvent) {\n if (typeof event.data !== \"string\") {\n return;\n }\n let data: OutgoingMessage;\n try {\n data = JSON.parse(event.data) as OutgoingMessage;\n } catch (error) {\n // silently ignore invalid messages for now\n // TODO: log errors with log levels\n return;\n }\n if (data.type === \"cf_agent_chat_messages\") {\n useChatHelpers.setMessages(data.messages);\n }\n }\n\n agent.addEventListener(\"message\", onClearHistory);\n agent.addEventListener(\"message\", onMessages);\n\n return () => {\n agent.removeEventListener(\"message\", onClearHistory);\n agent.removeEventListener(\"message\", onMessages);\n };\n }, [agent, useChatHelpers.setMessages]);\n\n return {\n ...useChatHelpers,\n /**\n * Set the chat messages and synchronize with the Agent\n * @param messages New messages to set\n */\n setMessages: (messages: Message[]) => {\n useChatHelpers.setMessages(messages);\n agent.send(\n JSON.stringify({\n type: \"cf_agent_chat_messages\",\n messages,\n })\n );\n },\n /**\n * Clear chat history on both client and Agent\n */\n clearHistory: () => {\n useChatHelpers.setMessages([]);\n agent.send(\n JSON.stringify({\n type: \"cf_agent_chat_clear\",\n })\n );\n },\n };\n}\n"],"mappings":";;;AAAA,SAAS,eAAe;AAExB,SAAS,KAAK,iBAAiB;AAG/B,SAAS,cAAc;AAwBvB,IAAM,eAAe,oBAAI,IAAgC;AAOlD,SAAS,aACd,SACA;AACA,QAAM,EAAE,OAAO,oBAAoB,GAAG,KAAK,IAAI;AAE/C,QAAM,WAAW,IAAI;AAAA,IACnB;AAAA,KACE,MAAM,QAA0B,MAAM,SACpC,QAAQ,SAAS,SAAS,EAC3B,QAAQ,UAAU,UAAU,CAAC;AAAA,EAClC;AAGA,WAAS,aAAa,OAAO,KAAK;AAClC,QAAM,iBAAiB,SAAS,SAAS;AAEzC,iBAAe,+BAA+B;AAAA,IAC5C;AAAA,EACF,GAA8B;AAC5B,UAAM,iBAAiB,IAAI,IAAI,GAAG;AAClC,mBAAe,YAAY;AAC3B,UAAM,WAAW,MAAM,MAAM,eAAe,SAAS,GAAG;AAAA,MACtD,SAAS,QAAQ;AAAA,MACjB,aAAa,QAAQ;AAAA,IACvB,CAAC;AACD,WAAO,SAAS,KAAgB;AAAA,EAClC;AAEA,QAAM,0BACJ,sBAAsB;AAExB,WAAS,qBACP,2BACA;AACA,QAAI,aAAa,IAAI,cAAc,GAAG;AACpC,aAAO,aAAa,IAAI,cAAc;AAAA,IACxC;AACA,UAAM,UAAU,wBAAwB,yBAAyB;AAIjE,iBAAa,IAAI,gBAAgB,OAAO;AACxC,WAAO;AAAA,EACT;AAEA,QAAM,yBACJ,uBAAuB,OACnB,OACA,qBAAqB;AAAA,IACnB,OAAO,MAAM;AAAA,IACb,MAAM,MAAM;AAAA,IACZ,KAAK;AAAA,EACP,CAAC;AACP,QAAM,kBAAkB,yBACpB,IAAI,sBAAsB,IACzB,KAAK,mBAAmB,CAAC;AAG9B,YAAU,MAAM;AACd,QAAI,CAAC,wBAAwB;AAC3B;AAAA,IACF;AAMA,iBAAa,IAAI,gBAAgB,sBAAuB;AACxD,WAAO,MAAM;AACX,UAAI,aAAa,IAAI,cAAc,MAAM,wBAAwB;AAC/D,qBAAa,OAAO,cAAc;AAAA,MACpC;AAAA,IACF;AAAA,EACF,GAAG,CAAC,gBAAgB,sBAAsB,CAAC;AAE3C,iBAAe,QACb,SACAA,WAAuB,CAAC,GACxB;AAKA,UAAM;AAAA,MACJ;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA,MACA;AAAA;AAAA,IAEF,IAAIA;AACJ,UAAM,KAAK,OAAO,CAAC;AACnB,UAAM,kBAAkB,IAAI,gBAAgB;AAE5C,YAAQ,iBAAiB,SAAS,MAAM;AAGtC,YAAM;AAAA,QACJ,KAAK,UAAU;AAAA,UACb,MAAM;AAAA,UACN;AAAA,QACF,CAAC;AAAA,MACH;AAWA,sBAAgB,MAAM;AAEtB,iBAAW,MAAM;AAAA,IACnB,CAAC;AAED,UAAM;AAAA,MACJ;AAAA,MACA,CAAC,UAAU;AACT,YAAI;AACJ,YAAI;AACF,iBAAO,KAAK,MAAM,MAAM,IAAI;AAAA,QAC9B,SAAS,OAAO;AAGd;AAAA,QACF;AACA,YAAI,KAAK,SAAS,8BAA8B;AAC9C,cAAI,KAAK,OAAO,IAAI;AAClB,uBAAW,QAAQ,IAAI,YAAY,EAAE,OAAO,KAAK,IAAI,CAAC;AACtD,gBAAI,KAAK,MAAM;AACb,yBAAW,MAAM;AACjB,8BAAgB,MAAM;AAAA,YACxB;AAAA,UACF;AAAA,QACF;AAAA,MACF;AAAA,MACA,EAAE,QAAQ,gBAAgB,OAAO;AAAA,IACnC;AAEA,QAAI;AAEJ,UAAM,SAAS,IAAI,eAAe;AAAA,MAChC,MAAM,GAAG;AACP,qBAAa;AAAA,MACf;AAAA,IACF,CAAC;AAED,UAAM;AAAA,MACJ,KAAK,UAAU;AAAA,QACb,MAAM;AAAA,QACN;AAAA,QACA,KAAK,QAAQ,SAAS;AAAA,QACtB,MAAM;AAAA,UACJ;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA;AAAA;AAAA,QAGF;AAAA,MACF,CAAC;AAAA,IACH;AAEA,WAAO,IAAI,SAAS,MAAM;AAAA,EAC5B;AACA,QAAM,iBAAiB,QAAQ;AAAA,IAC7B;AAAA,IACA,wBAAwB;AAAA,IACxB,OAAO;AAAA,IACP,GAAG;AAAA,EACL,CAAC;AAED,YAAU,MAAM;AACd,aAAS,eAAe,OAAqB;AAC3C,UAAI,OAAO,MAAM,SAAS,UAAU;AAClC;AAAA,MACF;AACA,UAAI;AACJ,UAAI;AACF,eAAO,KAAK,MAAM,MAAM,IAAI;AAAA,MAC9B,SAAS,OAAO;AAGd;AAAA,MACF;AACA,UAAI,KAAK,SAAS,uBAAuB;AACvC,uBAAe,YAAY,CAAC,CAAC;AAAA,MAC/B;AAAA,IACF;AAEA,aAAS,WAAW,OAAqB;AACvC,UAAI,OAAO,MAAM,SAAS,UAAU;AAClC;AAAA,MACF;AACA,UAAI;AACJ,UAAI;AACF,eAAO,KAAK,MAAM,MAAM,IAAI;AAAA,MAC9B,SAAS,OAAO;AAGd;AAAA,MACF;AACA,UAAI,KAAK,SAAS,0BAA0B;AAC1C,uBAAe,YAAY,KAAK,QAAQ;AAAA,MAC1C;AAAA,IACF;AAEA,UAAM,iBAAiB,WAAW,cAAc;AAChD,UAAM,iBAAiB,WAAW,UAAU;AAE5C,WAAO,MAAM;AACX,YAAM,oBAAoB,WAAW,cAAc;AACnD,YAAM,oBAAoB,WAAW,UAAU;AAAA,IACjD;AAAA,EACF,GAAG,CAAC,OAAO,eAAe,WAAW,CAAC;AAEtC,SAAO;AAAA,IACL,GAAG;AAAA;AAAA;AAAA;AAAA;AAAA,IAKH,aAAa,CAAC,aAAwB;AACpC,qBAAe,YAAY,QAAQ;AACnC,YAAM;AAAA,QACJ,KAAK,UAAU;AAAA,UACb,MAAM;AAAA,UACN;AAAA,QACF,CAAC;AAAA,MACH;AAAA,IACF;AAAA;AAAA;AAAA;AAAA,IAIA,cAAc,MAAM;AAClB,qBAAe,YAAY,CAAC,CAAC;AAC7B,YAAM;AAAA,QACJ,KAAK,UAAU;AAAA,UACb,MAAM;AAAA,QACR,CAAC;AAAA,MACH;AAAA,IACF;AAAA,EACF;AACF;","names":["options"]}
@@ -12,16 +12,22 @@ var SSEEdgeClientTransport = class extends SSEClientTransport {
12
12
  const workerOptions = {
13
13
  ...fetchInit,
14
14
  headers: {
15
+ ...options.requestInit?.headers,
15
16
  ...fetchInit?.headers,
16
17
  ...headers
17
18
  }
18
19
  };
19
20
  delete workerOptions.mode;
20
- return fetch(fetchUrl, workerOptions);
21
+ return options.eventSourceInit?.fetch?.(
22
+ fetchUrl,
23
+ // @ts-expect-error Expects FetchLikeInit from EventSource but is compatible with RequestInit
24
+ workerOptions
25
+ ) || fetch(fetchUrl, workerOptions);
21
26
  };
22
27
  super(url, {
23
28
  ...options,
24
29
  eventSourceInit: {
30
+ ...options.eventSourceInit,
25
31
  fetch: fetchOverride
26
32
  }
27
33
  });
@@ -47,7 +53,7 @@ import {
47
53
  } from "@modelcontextprotocol/sdk/types.js";
48
54
  import { Client } from "@modelcontextprotocol/sdk/client/index.js";
49
55
  var MCPClientConnection = class {
50
- constructor(url, info, options = { transport: {}, client: {}, capabilities: {} }) {
56
+ constructor(url, info, options = { transport: {}, client: {} }) {
51
57
  this.url = url;
52
58
  this.options = options;
53
59
  this.connectionState = "connecting";
@@ -56,7 +62,6 @@ var MCPClientConnection = class {
56
62
  this.resources = [];
57
63
  this.resourceTemplates = [];
58
64
  this.client = new Client(info, options.client);
59
- this.client.registerCapabilities(options.capabilities);
60
65
  }
61
66
  /**
62
67
  * Initialize a client connection
@@ -64,7 +69,7 @@ var MCPClientConnection = class {
64
69
  * @param code Optional OAuth code to initialize the connection with if auth hasn't been initialized
65
70
  * @returns
66
71
  */
67
- async init(code, clientId) {
72
+ async init(code) {
68
73
  try {
69
74
  const transport = new SSEEdgeClientTransport(
70
75
  this.url,
@@ -218,6 +223,7 @@ function capabilityErrorHandler(empty, method) {
218
223
 
219
224
  // src/mcp/client.ts
220
225
  import { jsonSchema } from "ai";
226
+ import { nanoid } from "nanoid";
221
227
  var MCPClientManager = class {
222
228
  /**
223
229
  * @param name Name of the MCP client
@@ -238,13 +244,16 @@ var MCPClientManager = class {
238
244
  * @param capabilities Client capabilities (i.e. if the client supports roots/sampling)
239
245
  */
240
246
  async connect(url, options = {}) {
241
- const id = options.reconnect?.id ?? crypto.randomUUID();
247
+ const id = options.reconnect?.id ?? nanoid(8);
242
248
  if (!options.transport?.authProvider) {
243
249
  console.warn(
244
250
  "No authProvider provided in the transport options. This client will only support unauthenticated remote MCP Servers"
245
251
  );
246
252
  } else {
247
253
  options.transport.authProvider.serverId = id;
254
+ if (options.reconnect?.oauthClientId) {
255
+ options.transport.authProvider.clientId = options.reconnect?.oauthClientId;
256
+ }
248
257
  }
249
258
  this.mcpConnections[id] = new MCPClientConnection(
250
259
  new URL(url),
@@ -254,23 +263,23 @@ var MCPClientManager = class {
254
263
  },
255
264
  {
256
265
  transport: options.transport ?? {},
257
- client: options.client ?? {},
258
- capabilities: options.client ?? {}
266
+ client: options.client ?? {}
259
267
  }
260
268
  );
261
- await this.mcpConnections[id].init(
262
- options.reconnect?.oauthCode,
263
- options.reconnect?.oauthClientId
264
- );
269
+ await this.mcpConnections[id].init(options.reconnect?.oauthCode);
265
270
  const authUrl = options.transport?.authProvider?.authUrl;
266
271
  if (authUrl && options.transport?.authProvider?.redirectUrl) {
267
272
  this.callbackUrls.push(
268
273
  options.transport.authProvider.redirectUrl.toString()
269
274
  );
275
+ return {
276
+ id,
277
+ authUrl,
278
+ clientId: options.transport?.authProvider?.clientId
279
+ };
270
280
  }
271
281
  return {
272
- id,
273
- authUrl
282
+ id
274
283
  };
275
284
  }
276
285
  isCallbackRequest(req) {
@@ -341,7 +350,7 @@ var MCPClientManager = class {
341
350
  return Object.fromEntries(
342
351
  getNamespacedData(this.mcpConnections, "tools").map((tool) => {
343
352
  return [
344
- tool.name,
353
+ `${tool.serverId}_${tool.name}`,
345
354
  {
346
355
  parameters: jsonSchema(tool.inputSchema),
347
356
  description: tool.description,
@@ -361,6 +370,26 @@ var MCPClientManager = class {
361
370
  })
362
371
  );
363
372
  }
373
+ /**
374
+ * Closes all connections to MCP servers
375
+ */
376
+ async closeAllConnections() {
377
+ return Promise.all(
378
+ Object.values(this.mcpConnections).map(async (connection) => {
379
+ await connection.client.close();
380
+ })
381
+ );
382
+ }
383
+ /**
384
+ * Closes a connection to an MCP server
385
+ * @param id The id of the connection to close
386
+ */
387
+ async closeConnection(id) {
388
+ if (!this.mcpConnections[id]) {
389
+ throw new Error(`Connection with id "${id}" does not exist.`);
390
+ }
391
+ await this.mcpConnections[id].client.close();
392
+ }
364
393
  /**
365
394
  * @returns namespaced list of prompts
366
395
  */
@@ -432,4 +461,4 @@ export {
432
461
  MCPClientManager,
433
462
  getNamespacedData
434
463
  };
435
- //# sourceMappingURL=chunk-YZNSS675.js.map
464
+ //# sourceMappingURL=chunk-25YDMV4H.js.map