agents 0.0.0-cf3b3d7 → 0.0.0-d08612f

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