agents 0.2.23 → 0.2.24

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.
Files changed (58) hide show
  1. package/README.md +5 -3
  2. package/dist/ai-chat-agent.d.ts +86 -13
  3. package/dist/ai-chat-agent.js +340 -74
  4. package/dist/ai-chat-agent.js.map +1 -1
  5. package/dist/{ai-chat-v5-migration-DBHGW4Hv.js → ai-chat-v5-migration-BSiGZmYU.js} +1 -1
  6. package/dist/{ai-chat-v5-migration-DBHGW4Hv.js.map → ai-chat-v5-migration-BSiGZmYU.js.map} +1 -1
  7. package/dist/ai-chat-v5-migration.js +1 -1
  8. package/dist/ai-react.d.ts +14 -9
  9. package/dist/ai-react.js +164 -29
  10. package/dist/ai-react.js.map +1 -1
  11. package/dist/{ai-types-D5YoPrBZ.d.ts → ai-types-81H_-Uxh.d.ts} +15 -7
  12. package/dist/{ai-types-B3aQaFv3.js → ai-types-CrMqkwc_.js} +5 -1
  13. package/dist/ai-types-CrMqkwc_.js.map +1 -0
  14. package/dist/ai-types.d.ts +1 -1
  15. package/dist/ai-types.js +1 -1
  16. package/dist/{client-BfiZ3HQd.js → client-B3SR12TQ.js} +2 -2
  17. package/dist/{client-BfiZ3HQd.js.map → client-B3SR12TQ.js.map} +1 -1
  18. package/dist/{client-CbWe9FBd.d.ts → client-BAQA84dr.d.ts} +2 -2
  19. package/dist/client-BZq9qau2.js +1093 -0
  20. package/dist/client-BZq9qau2.js.map +1 -0
  21. package/dist/client-CsaP9Irq.d.ts +1528 -0
  22. package/dist/client.d.ts +2 -2
  23. package/dist/client.js +2 -2
  24. package/dist/codemode/ai.js +5 -5
  25. package/dist/{do-oauth-client-provider-DGc5pP0l.d.ts → do-oauth-client-provider-C2CHH5x-.d.ts} +1 -1
  26. package/dist/{do-oauth-client-provider-CswoD5Lu.js → do-oauth-client-provider-CwqK5SXm.js} +2 -1
  27. package/dist/do-oauth-client-provider-CwqK5SXm.js.map +1 -0
  28. package/dist/{index-DhJCaDWd.d.ts → index-BUle9RiP.d.ts} +2 -2
  29. package/dist/{index-DEZFuVoU.d.ts → index-Bx5KK3VJ.d.ts} +44 -23
  30. package/dist/index.d.ts +6 -6
  31. package/dist/index.js +5 -5
  32. package/dist/mcp/client.d.ts +4 -4
  33. package/dist/mcp/client.js +2 -1
  34. package/dist/mcp/do-oauth-client-provider.d.ts +1 -1
  35. package/dist/mcp/do-oauth-client-provider.js +1 -1
  36. package/dist/mcp/index.d.ts +11 -11
  37. package/dist/mcp/index.js +21 -14
  38. package/dist/mcp/index.js.map +1 -1
  39. package/dist/mcp/x402.js +10 -6
  40. package/dist/mcp/x402.js.map +1 -1
  41. package/dist/{mcp-Dw5vDrY8.d.ts → mcp-BwPscEiF.d.ts} +1 -1
  42. package/dist/observability/index.d.ts +2 -2
  43. package/dist/observability/index.js +5 -5
  44. package/dist/{react-CrV9uP64.d.ts → react-CbwD4fBf.d.ts} +4 -4
  45. package/dist/react.d.ts +9 -9
  46. package/dist/react.js +2 -2
  47. package/dist/react.js.map +1 -1
  48. package/dist/{serializable-CymX8ovI.d.ts → serializable-faDkMCai.d.ts} +1 -1
  49. package/dist/serializable.d.ts +1 -1
  50. package/dist/{src-C8K3lu37.js → src-D_KKH_4c.js} +99 -146
  51. package/dist/src-D_KKH_4c.js.map +1 -0
  52. package/package.json +7 -7
  53. package/dist/ai-types-B3aQaFv3.js.map +0 -1
  54. package/dist/client-DZhjV_XA.js +0 -719
  55. package/dist/client-DZhjV_XA.js.map +0 -1
  56. package/dist/client-GfgZTqrS.d.ts +0 -5283
  57. package/dist/do-oauth-client-provider-CswoD5Lu.js.map +0 -1
  58. package/dist/src-C8K3lu37.js.map +0 -1
package/README.md CHANGED
@@ -549,10 +549,12 @@ export class MyMCP extends McpAgent<Env, State, {}> {
549
549
  };
550
550
  });
551
551
 
552
- this.server.tool(
552
+ this.server.registerTool(
553
553
  "add",
554
- "Add to the counter, stored in the MCP",
555
- { a: z.number() },
554
+ {
555
+ description: "Add to the counter, stored in the MCP",
556
+ inputSchema: { a: z.number() }
557
+ },
556
558
  async ({ a }) => {
557
559
  this.setState({ ...this.state, counter: this.state.counter + a });
558
560
 
@@ -1,14 +1,9 @@
1
- import "./client-GfgZTqrS.js";
2
- import "./mcp-Dw5vDrY8.js";
3
- import "./do-oauth-client-provider-DGc5pP0l.js";
4
- import "./index-DhJCaDWd.js";
5
- import "./ai-types-D5YoPrBZ.js";
6
- import {
7
- n as AgentContext,
8
- s as Connection,
9
- t as Agent,
10
- x as WSMessage
11
- } from "./index-DEZFuVoU.js";
1
+ import "./client-CsaP9Irq.js";
2
+ import "./mcp-BwPscEiF.js";
3
+ import "./do-oauth-client-provider-C2CHH5x-.js";
4
+ import "./index-BUle9RiP.js";
5
+ import "./ai-types-81H_-Uxh.js";
6
+ import { n as AgentContext, t as Agent } from "./index-Bx5KK3VJ.js";
12
7
  import { StreamTextOnFinishCallback, ToolSet, UIMessage } from "ai";
13
8
 
14
9
  //#region src/ai-chat-agent.d.ts
@@ -25,12 +20,85 @@ declare class AIChatAgent<Env = unknown, State = unknown> extends Agent<
25
20
  * useful to propagate request cancellation signals for any external calls made by the agent
26
21
  */
27
22
  private _chatMessageAbortControllers;
23
+ /**
24
+ * Currently active stream ID for resumable streaming.
25
+ * Stored in memory for quick access; persisted in stream_metadata table.
26
+ */
27
+ private _activeStreamId;
28
+ /**
29
+ * Request ID associated with the active stream.
30
+ */
31
+ private _activeRequestId;
32
+ /**
33
+ * Current chunk index for the active stream
34
+ */
35
+ private _streamChunkIndex;
36
+ /**
37
+ * Buffer for stream chunks pending write to SQLite.
38
+ * Chunks are batched and flushed when buffer reaches CHUNK_BUFFER_SIZE.
39
+ */
40
+ private _chunkBuffer;
41
+ /**
42
+ * Lock to prevent concurrent flush operations
43
+ */
44
+ private _isFlushingChunks;
45
+ /**
46
+ * Timestamp of the last cleanup operation for old streams
47
+ */
48
+ private _lastCleanupTime;
28
49
  /** Array of chat messages for the current conversation */
29
50
  messages: UIMessage[];
30
51
  constructor(ctx: AgentContext, env: Env);
52
+ /**
53
+ * Restore active stream state if the agent was restarted during streaming.
54
+ * Called during construction to recover any interrupted streams.
55
+ * Validates stream freshness to avoid sending stale resume notifications.
56
+ */
57
+ private _restoreActiveStream;
58
+ /**
59
+ * Notify a connection about an active stream that can be resumed.
60
+ * The client should respond with CF_AGENT_STREAM_RESUME_ACK to receive chunks.
61
+ * Uses in-memory state for request ID - no extra DB lookup needed.
62
+ * @param connection - The WebSocket connection to notify
63
+ */
64
+ private _notifyStreamResuming;
65
+ /**
66
+ * Send stream chunks to a connection after receiving ACK.
67
+ * @param connection - The WebSocket connection
68
+ * @param streamId - The stream to replay
69
+ * @param requestId - The original request ID
70
+ */
71
+ private _sendStreamChunks;
72
+ /**
73
+ * Buffer a stream chunk for batch write to SQLite.
74
+ * @param streamId - The stream this chunk belongs to
75
+ * @param body - The serialized chunk body
76
+ */
77
+ private _storeStreamChunk;
78
+ /**
79
+ * Flush buffered chunks to SQLite in a single batch.
80
+ * Uses a lock to prevent concurrent flush operations.
81
+ */
82
+ private _flushChunkBuffer;
83
+ /**
84
+ * Start tracking a new stream for resumable streaming.
85
+ * Creates metadata entry in SQLite and sets up tracking state.
86
+ * @param requestId - The unique ID of the chat request
87
+ * @returns The generated stream ID
88
+ */
89
+ private _startStream;
90
+ /**
91
+ * Mark a stream as completed and flush any pending chunks.
92
+ * @param streamId - The stream to mark as completed
93
+ */
94
+ private _completeStream;
95
+ /**
96
+ * Clean up old completed streams if enough time has passed since last cleanup.
97
+ * This prevents database growth while avoiding cleanup overhead on every stream completion.
98
+ */
99
+ private _maybeCleanupOldStreams;
31
100
  private _broadcastChatMessage;
32
101
  private _loadMessagesFromDb;
33
- onMessage(connection: Connection, message: WSMessage): Promise<void>;
34
102
  onRequest(request: Request): Promise<Response>;
35
103
  private _tryCatchChat;
36
104
  /**
@@ -55,6 +123,11 @@ declare class AIChatAgent<Env = unknown, State = unknown> extends Agent<
55
123
  excludeBroadcastIds?: string[]
56
124
  ): Promise<void>;
57
125
  private _reply;
126
+ /**
127
+ * Mark a stream as errored and clean up state.
128
+ * @param streamId - The stream to mark as errored
129
+ */
130
+ private _markStreamError;
58
131
  /**
59
132
  * For the given message id, look up its associated AbortController
60
133
  * If the AbortController does not exist, create and store one in memory
@@ -75,7 +148,7 @@ declare class AIChatAgent<Env = unknown, State = unknown> extends Agent<
75
148
  */
76
149
  private _destroyAbortControllers;
77
150
  /**
78
- * When the DO is destroyed, cancel all pending requests
151
+ * When the DO is destroyed, cancel all pending requests and clean up resources
79
152
  */
80
153
  destroy(): Promise<void>;
81
154
  }
@@ -1,11 +1,22 @@
1
- import { t as MessageType } from "./ai-types-B3aQaFv3.js";
2
- import "./client-BfiZ3HQd.js";
3
- import "./client-DZhjV_XA.js";
4
- import "./do-oauth-client-provider-CswoD5Lu.js";
5
- import { t as Agent } from "./src-C8K3lu37.js";
6
- import { r as autoTransformMessages } from "./ai-chat-v5-migration-DBHGW4Hv.js";
1
+ import { t as MessageType } from "./ai-types-CrMqkwc_.js";
2
+ import "./client-B3SR12TQ.js";
3
+ import "./client-BZq9qau2.js";
4
+ import "./do-oauth-client-provider-CwqK5SXm.js";
5
+ import { t as Agent } from "./src-D_KKH_4c.js";
6
+ import { r as autoTransformMessages } from "./ai-chat-v5-migration-BSiGZmYU.js";
7
+ import { nanoid } from "nanoid";
7
8
 
8
9
  //#region src/ai-chat-agent.ts
10
+ /** Number of chunks to buffer before flushing to SQLite */
11
+ const CHUNK_BUFFER_SIZE = 10;
12
+ /** Maximum buffer size to prevent memory issues on rapid reconnections */
13
+ const CHUNK_BUFFER_MAX_SIZE = 100;
14
+ /** Maximum age for a "streaming" stream before considering it stale (ms) - 5 minutes */
15
+ const STREAM_STALE_THRESHOLD_MS = 300 * 1e3;
16
+ /** Default cleanup interval for old streams (ms) - every 10 minutes */
17
+ const CLEANUP_INTERVAL_MS = 600 * 1e3;
18
+ /** Default age threshold for cleaning up completed streams (ms) - 24 hours */
19
+ const CLEANUP_AGE_THRESHOLD_MS = 1440 * 60 * 1e3;
9
20
  const decoder = new TextDecoder();
10
21
  /**
11
22
  * Extension of Agent with built-in chat capabilities
@@ -14,13 +25,279 @@ const decoder = new TextDecoder();
14
25
  var AIChatAgent = class extends Agent {
15
26
  constructor(ctx, env) {
16
27
  super(ctx, env);
28
+ this._activeStreamId = null;
29
+ this._activeRequestId = null;
30
+ this._streamChunkIndex = 0;
31
+ this._chunkBuffer = [];
32
+ this._isFlushingChunks = false;
33
+ this._lastCleanupTime = 0;
17
34
  this.sql`create table if not exists cf_ai_chat_agent_messages (
18
35
  id text primary key,
19
36
  message text not null,
20
37
  created_at datetime default current_timestamp
21
38
  )`;
39
+ this.sql`create table if not exists cf_ai_chat_stream_chunks (
40
+ id text primary key,
41
+ stream_id text not null,
42
+ body text not null,
43
+ chunk_index integer not null,
44
+ created_at integer not null
45
+ )`;
46
+ this.sql`create table if not exists cf_ai_chat_stream_metadata (
47
+ id text primary key,
48
+ request_id text not null,
49
+ status text not null,
50
+ created_at integer not null,
51
+ completed_at integer
52
+ )`;
53
+ this.sql`create index if not exists idx_stream_chunks_stream_id
54
+ on cf_ai_chat_stream_chunks(stream_id, chunk_index)`;
22
55
  this.messages = autoTransformMessages(this._loadMessagesFromDb());
23
56
  this._chatMessageAbortControllers = /* @__PURE__ */ new Map();
57
+ this._restoreActiveStream();
58
+ const _onConnect = this.onConnect.bind(this);
59
+ this.onConnect = async (connection, ctx$1) => {
60
+ if (this._activeStreamId) this._notifyStreamResuming(connection);
61
+ return _onConnect(connection, ctx$1);
62
+ };
63
+ const _onMessage = this.onMessage.bind(this);
64
+ this.onMessage = async (connection, message) => {
65
+ if (typeof message === "string") {
66
+ let data;
67
+ try {
68
+ data = JSON.parse(message);
69
+ } catch (_error) {
70
+ return _onMessage(connection, message);
71
+ }
72
+ if (data.type === MessageType.CF_AGENT_USE_CHAT_REQUEST && data.init.method === "POST") {
73
+ const { body } = data.init;
74
+ const { messages } = JSON.parse(body);
75
+ const transformedMessages = autoTransformMessages(messages);
76
+ this._broadcastChatMessage({
77
+ messages: transformedMessages,
78
+ type: MessageType.CF_AGENT_CHAT_MESSAGES
79
+ }, [connection.id]);
80
+ await this.persistMessages(transformedMessages, [connection.id]);
81
+ this.observability?.emit({
82
+ displayMessage: "Chat message request",
83
+ id: data.id,
84
+ payload: {},
85
+ timestamp: Date.now(),
86
+ type: "message:request"
87
+ }, this.ctx);
88
+ const chatMessageId = data.id;
89
+ const abortSignal = this._getAbortSignal(chatMessageId);
90
+ return this._tryCatchChat(async () => {
91
+ const response = await this.onChatMessage(async (_finishResult) => {
92
+ this._removeAbortController(chatMessageId);
93
+ this.observability?.emit({
94
+ displayMessage: "Chat message response",
95
+ id: data.id,
96
+ payload: {},
97
+ timestamp: Date.now(),
98
+ type: "message:response"
99
+ }, this.ctx);
100
+ }, abortSignal ? { abortSignal } : void 0);
101
+ if (response) await this._reply(data.id, response);
102
+ else {
103
+ console.warn(`[AIChatAgent] onChatMessage returned no response for chatMessageId: ${chatMessageId}`);
104
+ this._broadcastChatMessage({
105
+ body: "No response was generated by the agent.",
106
+ done: true,
107
+ id: data.id,
108
+ type: MessageType.CF_AGENT_USE_CHAT_RESPONSE
109
+ }, [connection.id]);
110
+ }
111
+ });
112
+ }
113
+ if (data.type === MessageType.CF_AGENT_CHAT_CLEAR) {
114
+ this._destroyAbortControllers();
115
+ this.sql`delete from cf_ai_chat_agent_messages`;
116
+ this.sql`delete from cf_ai_chat_stream_chunks`;
117
+ this.sql`delete from cf_ai_chat_stream_metadata`;
118
+ this._activeStreamId = null;
119
+ this._activeRequestId = null;
120
+ this._streamChunkIndex = 0;
121
+ this.messages = [];
122
+ this._broadcastChatMessage({ type: MessageType.CF_AGENT_CHAT_CLEAR }, [connection.id]);
123
+ return;
124
+ }
125
+ if (data.type === MessageType.CF_AGENT_CHAT_MESSAGES) {
126
+ const transformedMessages = autoTransformMessages(data.messages);
127
+ await this.persistMessages(transformedMessages, [connection.id]);
128
+ return;
129
+ }
130
+ if (data.type === MessageType.CF_AGENT_CHAT_REQUEST_CANCEL) {
131
+ this._cancelChatRequest(data.id);
132
+ return;
133
+ }
134
+ if (data.type === MessageType.CF_AGENT_STREAM_RESUME_ACK) {
135
+ if (this._activeStreamId && this._activeRequestId && this._activeRequestId === data.id) this._sendStreamChunks(connection, this._activeStreamId, this._activeRequestId);
136
+ return;
137
+ }
138
+ }
139
+ return _onMessage(connection, message);
140
+ };
141
+ }
142
+ /**
143
+ * Restore active stream state if the agent was restarted during streaming.
144
+ * Called during construction to recover any interrupted streams.
145
+ * Validates stream freshness to avoid sending stale resume notifications.
146
+ */
147
+ _restoreActiveStream() {
148
+ const activeStreams = this.sql`
149
+ select * from cf_ai_chat_stream_metadata
150
+ where status = 'streaming'
151
+ order by created_at desc
152
+ limit 1
153
+ `;
154
+ if (activeStreams && activeStreams.length > 0) {
155
+ const stream = activeStreams[0];
156
+ const streamAge = Date.now() - stream.created_at;
157
+ if (streamAge > STREAM_STALE_THRESHOLD_MS) {
158
+ this.sql`delete from cf_ai_chat_stream_chunks where stream_id = ${stream.id}`;
159
+ this.sql`delete from cf_ai_chat_stream_metadata where id = ${stream.id}`;
160
+ console.warn(`[AIChatAgent] Deleted stale stream ${stream.id} (age: ${Math.round(streamAge / 1e3)}s)`);
161
+ return;
162
+ }
163
+ this._activeStreamId = stream.id;
164
+ this._activeRequestId = stream.request_id;
165
+ const lastChunk = this.sql`
166
+ select max(chunk_index) as max_index
167
+ from cf_ai_chat_stream_chunks
168
+ where stream_id = ${this._activeStreamId}
169
+ `;
170
+ this._streamChunkIndex = lastChunk && lastChunk[0]?.max_index != null ? lastChunk[0].max_index + 1 : 0;
171
+ }
172
+ }
173
+ /**
174
+ * Notify a connection about an active stream that can be resumed.
175
+ * The client should respond with CF_AGENT_STREAM_RESUME_ACK to receive chunks.
176
+ * Uses in-memory state for request ID - no extra DB lookup needed.
177
+ * @param connection - The WebSocket connection to notify
178
+ */
179
+ _notifyStreamResuming(connection) {
180
+ if (!this._activeStreamId || !this._activeRequestId) return;
181
+ connection.send(JSON.stringify({
182
+ type: MessageType.CF_AGENT_STREAM_RESUMING,
183
+ id: this._activeRequestId
184
+ }));
185
+ }
186
+ /**
187
+ * Send stream chunks to a connection after receiving ACK.
188
+ * @param connection - The WebSocket connection
189
+ * @param streamId - The stream to replay
190
+ * @param requestId - The original request ID
191
+ */
192
+ _sendStreamChunks(connection, streamId, requestId) {
193
+ this._flushChunkBuffer();
194
+ const chunks = this.sql`
195
+ select * from cf_ai_chat_stream_chunks
196
+ where stream_id = ${streamId}
197
+ order by chunk_index asc
198
+ `;
199
+ for (const chunk of chunks || []) connection.send(JSON.stringify({
200
+ body: chunk.body,
201
+ done: false,
202
+ id: requestId,
203
+ type: MessageType.CF_AGENT_USE_CHAT_RESPONSE
204
+ }));
205
+ if (this._activeStreamId !== streamId) connection.send(JSON.stringify({
206
+ body: "",
207
+ done: true,
208
+ id: requestId,
209
+ type: MessageType.CF_AGENT_USE_CHAT_RESPONSE
210
+ }));
211
+ }
212
+ /**
213
+ * Buffer a stream chunk for batch write to SQLite.
214
+ * @param streamId - The stream this chunk belongs to
215
+ * @param body - The serialized chunk body
216
+ */
217
+ _storeStreamChunk(streamId, body) {
218
+ if (this._chunkBuffer.length >= CHUNK_BUFFER_MAX_SIZE) this._flushChunkBuffer();
219
+ this._chunkBuffer.push({
220
+ id: nanoid(),
221
+ streamId,
222
+ body,
223
+ index: this._streamChunkIndex
224
+ });
225
+ this._streamChunkIndex++;
226
+ if (this._chunkBuffer.length >= CHUNK_BUFFER_SIZE) this._flushChunkBuffer();
227
+ }
228
+ /**
229
+ * Flush buffered chunks to SQLite in a single batch.
230
+ * Uses a lock to prevent concurrent flush operations.
231
+ */
232
+ _flushChunkBuffer() {
233
+ if (this._isFlushingChunks || this._chunkBuffer.length === 0) return;
234
+ this._isFlushingChunks = true;
235
+ try {
236
+ const chunks = this._chunkBuffer;
237
+ this._chunkBuffer = [];
238
+ const now = Date.now();
239
+ for (const chunk of chunks) this.sql`
240
+ insert into cf_ai_chat_stream_chunks (id, stream_id, body, chunk_index, created_at)
241
+ values (${chunk.id}, ${chunk.streamId}, ${chunk.body}, ${chunk.index}, ${now})
242
+ `;
243
+ } finally {
244
+ this._isFlushingChunks = false;
245
+ }
246
+ }
247
+ /**
248
+ * Start tracking a new stream for resumable streaming.
249
+ * Creates metadata entry in SQLite and sets up tracking state.
250
+ * @param requestId - The unique ID of the chat request
251
+ * @returns The generated stream ID
252
+ */
253
+ _startStream(requestId) {
254
+ this._flushChunkBuffer();
255
+ const streamId = nanoid();
256
+ this._activeStreamId = streamId;
257
+ this._activeRequestId = requestId;
258
+ this._streamChunkIndex = 0;
259
+ this.sql`
260
+ insert into cf_ai_chat_stream_metadata (id, request_id, status, created_at)
261
+ values (${streamId}, ${requestId}, 'streaming', ${Date.now()})
262
+ `;
263
+ return streamId;
264
+ }
265
+ /**
266
+ * Mark a stream as completed and flush any pending chunks.
267
+ * @param streamId - The stream to mark as completed
268
+ */
269
+ _completeStream(streamId) {
270
+ this._flushChunkBuffer();
271
+ this.sql`
272
+ update cf_ai_chat_stream_metadata
273
+ set status = 'completed', completed_at = ${Date.now()}
274
+ where id = ${streamId}
275
+ `;
276
+ this._activeStreamId = null;
277
+ this._activeRequestId = null;
278
+ this._streamChunkIndex = 0;
279
+ this._maybeCleanupOldStreams();
280
+ }
281
+ /**
282
+ * Clean up old completed streams if enough time has passed since last cleanup.
283
+ * This prevents database growth while avoiding cleanup overhead on every stream completion.
284
+ */
285
+ _maybeCleanupOldStreams() {
286
+ const now = Date.now();
287
+ if (now - this._lastCleanupTime < CLEANUP_INTERVAL_MS) return;
288
+ this._lastCleanupTime = now;
289
+ const cutoff = now - CLEANUP_AGE_THRESHOLD_MS;
290
+ this.sql`
291
+ delete from cf_ai_chat_stream_chunks
292
+ where stream_id in (
293
+ select id from cf_ai_chat_stream_metadata
294
+ where status = 'completed' and completed_at < ${cutoff}
295
+ )
296
+ `;
297
+ this.sql`
298
+ delete from cf_ai_chat_stream_metadata
299
+ where status = 'completed' and completed_at < ${cutoff}
300
+ `;
24
301
  }
25
302
  _broadcastChatMessage(message, exclude) {
26
303
  this.broadcast(JSON.stringify(message), exclude);
@@ -35,68 +312,8 @@ var AIChatAgent = class extends Agent {
35
312
  }
36
313
  }).filter((msg) => msg !== null);
37
314
  }
38
- async onMessage(connection, message) {
39
- if (typeof message === "string") {
40
- let data;
41
- try {
42
- data = JSON.parse(message);
43
- } catch (_error) {
44
- return;
45
- }
46
- if (data.type === MessageType.CF_AGENT_USE_CHAT_REQUEST && data.init.method === "POST") {
47
- const { body } = data.init;
48
- const { messages } = JSON.parse(body);
49
- const transformedMessages = autoTransformMessages(messages);
50
- this._broadcastChatMessage({
51
- messages: transformedMessages,
52
- type: MessageType.CF_AGENT_CHAT_MESSAGES
53
- }, [connection.id]);
54
- await this.persistMessages(transformedMessages, [connection.id]);
55
- this.observability?.emit({
56
- displayMessage: "Chat message request",
57
- id: data.id,
58
- payload: {},
59
- timestamp: Date.now(),
60
- type: "message:request"
61
- }, this.ctx);
62
- const chatMessageId = data.id;
63
- const abortSignal = this._getAbortSignal(chatMessageId);
64
- return this._tryCatchChat(async () => {
65
- const response = await this.onChatMessage(async (_finishResult) => {
66
- this._removeAbortController(chatMessageId);
67
- this.observability?.emit({
68
- displayMessage: "Chat message response",
69
- id: data.id,
70
- payload: {},
71
- timestamp: Date.now(),
72
- type: "message:response"
73
- }, this.ctx);
74
- }, abortSignal ? { abortSignal } : void 0);
75
- if (response) await this._reply(data.id, response);
76
- else {
77
- console.warn(`[AIChatAgent] onChatMessage returned no response for chatMessageId: ${chatMessageId}`);
78
- this._broadcastChatMessage({
79
- body: "No response was generated by the agent.",
80
- done: true,
81
- id: data.id,
82
- type: MessageType.CF_AGENT_USE_CHAT_RESPONSE
83
- }, [connection.id]);
84
- }
85
- });
86
- }
87
- if (data.type === MessageType.CF_AGENT_CHAT_CLEAR) {
88
- this._destroyAbortControllers();
89
- this.sql`delete from cf_ai_chat_agent_messages`;
90
- this.messages = [];
91
- this._broadcastChatMessage({ type: MessageType.CF_AGENT_CHAT_CLEAR }, [connection.id]);
92
- } else if (data.type === MessageType.CF_AGENT_CHAT_MESSAGES) {
93
- const transformedMessages = autoTransformMessages(data.messages);
94
- await this.persistMessages(transformedMessages, [connection.id]);
95
- } else if (data.type === MessageType.CF_AGENT_CHAT_REQUEST_CANCEL) this._cancelChatRequest(data.id);
96
- }
97
- }
98
315
  async onRequest(request) {
99
- return this._tryCatchChat(() => {
316
+ return this._tryCatchChat(async () => {
100
317
  if (new URL(request.url).pathname.endsWith("/get-messages")) {
101
318
  const messages = this._loadMessagesFromDb();
102
319
  return Response.json(messages);
@@ -154,6 +371,7 @@ var AIChatAgent = class extends Agent {
154
371
  });
155
372
  return;
156
373
  }
374
+ const streamId = this._startStream(id);
157
375
  const { getToolName, isToolUIPart, parsePartialJson } = await import("ai");
158
376
  const reader = response.body.getReader();
159
377
  const message = {
@@ -221,10 +439,13 @@ var AIChatAgent = class extends Agent {
221
439
  ...metadata
222
440
  } : metadata;
223
441
  }
442
+ let streamCompleted = false;
224
443
  try {
225
444
  while (true) {
226
445
  const { done, value } = await reader.read();
227
446
  if (done) {
447
+ this._completeStream(streamId);
448
+ streamCompleted = true;
228
449
  this._broadcastChatMessage({
229
450
  body: "",
230
451
  done: true,
@@ -467,29 +688,54 @@ var AIChatAgent = class extends Agent {
467
688
  });
468
689
  break;
469
690
  }
691
+ let eventToSend = data;
692
+ if (data.type === "finish" && "finishReason" in data) {
693
+ const { finishReason, ...rest } = data;
694
+ eventToSend = {
695
+ ...rest,
696
+ type: "finish",
697
+ messageMetadata: { finishReason }
698
+ };
699
+ }
700
+ const chunkBody = JSON.stringify(eventToSend);
701
+ this._storeStreamChunk(streamId, chunkBody);
470
702
  this._broadcastChatMessage({
471
- body: JSON.stringify(data),
703
+ body: chunkBody,
472
704
  done: false,
473
705
  id,
474
706
  type: MessageType.CF_AGENT_USE_CHAT_RESPONSE
475
707
  });
476
- } catch (_e) {}
708
+ } catch (_error) {}
477
709
  } else if (chunk.length > 0) {
478
710
  message.parts.push({
479
711
  type: "text",
480
712
  text: chunk
481
713
  });
714
+ const chunkBody = JSON.stringify({
715
+ type: "text-delta",
716
+ delta: chunk
717
+ });
718
+ this._storeStreamChunk(streamId, chunkBody);
482
719
  this._broadcastChatMessage({
483
- body: JSON.stringify({
484
- type: "text-delta",
485
- delta: chunk
486
- }),
720
+ body: chunkBody,
487
721
  done: false,
488
722
  id,
489
723
  type: MessageType.CF_AGENT_USE_CHAT_RESPONSE
490
724
  });
491
725
  }
492
726
  }
727
+ } catch (error) {
728
+ if (!streamCompleted) {
729
+ this._markStreamError(streamId);
730
+ this._broadcastChatMessage({
731
+ body: error instanceof Error ? error.message : "Stream error",
732
+ done: true,
733
+ error: true,
734
+ id,
735
+ type: MessageType.CF_AGENT_USE_CHAT_RESPONSE
736
+ });
737
+ }
738
+ throw error;
493
739
  } finally {
494
740
  reader.releaseLock();
495
741
  }
@@ -497,6 +743,21 @@ var AIChatAgent = class extends Agent {
497
743
  });
498
744
  }
499
745
  /**
746
+ * Mark a stream as errored and clean up state.
747
+ * @param streamId - The stream to mark as errored
748
+ */
749
+ _markStreamError(streamId) {
750
+ this._flushChunkBuffer();
751
+ this.sql`
752
+ update cf_ai_chat_stream_metadata
753
+ set status = 'error', completed_at = ${Date.now()}
754
+ where id = ${streamId}
755
+ `;
756
+ this._activeStreamId = null;
757
+ this._activeRequestId = null;
758
+ this._streamChunkIndex = 0;
759
+ }
760
+ /**
500
761
  * For the given message id, look up its associated AbortController
501
762
  * If the AbortController does not exist, create and store one in memory
502
763
  *
@@ -527,10 +788,15 @@ var AIChatAgent = class extends Agent {
527
788
  this._chatMessageAbortControllers.clear();
528
789
  }
529
790
  /**
530
- * When the DO is destroyed, cancel all pending requests
791
+ * When the DO is destroyed, cancel all pending requests and clean up resources
531
792
  */
532
793
  async destroy() {
533
794
  this._destroyAbortControllers();
795
+ this._flushChunkBuffer();
796
+ this.sql`drop table if exists cf_ai_chat_stream_chunks`;
797
+ this.sql`drop table if exists cf_ai_chat_stream_metadata`;
798
+ this._activeStreamId = null;
799
+ this._activeRequestId = null;
534
800
  await super.destroy();
535
801
  }
536
802
  };