agents 0.2.23 → 0.2.25
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.
- package/README.md +5 -3
- package/dist/ai-chat-agent.d.ts +86 -13
- package/dist/ai-chat-agent.js +340 -74
- package/dist/ai-chat-agent.js.map +1 -1
- package/dist/{ai-chat-v5-migration-DBHGW4Hv.js → ai-chat-v5-migration-BSiGZmYU.js} +1 -1
- package/dist/{ai-chat-v5-migration-DBHGW4Hv.js.map → ai-chat-v5-migration-BSiGZmYU.js.map} +1 -1
- package/dist/ai-chat-v5-migration.js +1 -1
- package/dist/ai-react.d.ts +14 -9
- package/dist/ai-react.js +179 -29
- package/dist/ai-react.js.map +1 -1
- package/dist/{ai-types-D5YoPrBZ.d.ts → ai-types-81H_-Uxh.d.ts} +15 -7
- package/dist/{ai-types-B3aQaFv3.js → ai-types-CrMqkwc_.js} +5 -1
- package/dist/ai-types-CrMqkwc_.js.map +1 -0
- package/dist/ai-types.d.ts +1 -1
- package/dist/ai-types.js +1 -1
- package/dist/{client-BfiZ3HQd.js → client-B3SR12TQ.js} +2 -2
- package/dist/{client-BfiZ3HQd.js.map → client-B3SR12TQ.js.map} +1 -1
- package/dist/{client-CbWe9FBd.d.ts → client-BAQA84dr.d.ts} +2 -2
- package/dist/client-BZq9qau2.js +1093 -0
- package/dist/client-BZq9qau2.js.map +1 -0
- package/dist/client-CsaP9Irq.d.ts +1528 -0
- package/dist/client.d.ts +2 -2
- package/dist/client.js +2 -2
- package/dist/codemode/ai.js +5 -5
- package/dist/{do-oauth-client-provider-DGc5pP0l.d.ts → do-oauth-client-provider-C2CHH5x-.d.ts} +1 -1
- package/dist/{do-oauth-client-provider-CswoD5Lu.js → do-oauth-client-provider-CwqK5SXm.js} +2 -1
- package/dist/do-oauth-client-provider-CwqK5SXm.js.map +1 -0
- package/dist/{index-DhJCaDWd.d.ts → index-BUle9RiP.d.ts} +2 -2
- package/dist/{index-DEZFuVoU.d.ts → index-Bx5KK3VJ.d.ts} +44 -23
- package/dist/index.d.ts +6 -6
- package/dist/index.js +5 -5
- package/dist/mcp/client.d.ts +4 -4
- package/dist/mcp/client.js +2 -1
- package/dist/mcp/do-oauth-client-provider.d.ts +1 -1
- package/dist/mcp/do-oauth-client-provider.js +1 -1
- package/dist/mcp/index.d.ts +11 -11
- package/dist/mcp/index.js +21 -14
- package/dist/mcp/index.js.map +1 -1
- package/dist/mcp/x402.js +10 -6
- package/dist/mcp/x402.js.map +1 -1
- package/dist/{mcp-Dw5vDrY8.d.ts → mcp-BwPscEiF.d.ts} +1 -1
- package/dist/observability/index.d.ts +2 -2
- package/dist/observability/index.js +5 -5
- package/dist/{react-CrV9uP64.d.ts → react-CbwD4fBf.d.ts} +4 -4
- package/dist/react.d.ts +9 -9
- package/dist/react.js +2 -2
- package/dist/react.js.map +1 -1
- package/dist/{serializable-CymX8ovI.d.ts → serializable-faDkMCai.d.ts} +1 -1
- package/dist/serializable.d.ts +1 -1
- package/dist/{src-C8K3lu37.js → src-D_KKH_4c.js} +99 -146
- package/dist/src-D_KKH_4c.js.map +1 -0
- package/package.json +7 -7
- package/dist/ai-types-B3aQaFv3.js.map +0 -1
- package/dist/client-DZhjV_XA.js +0 -719
- package/dist/client-DZhjV_XA.js.map +0 -1
- package/dist/client-GfgZTqrS.d.ts +0 -5283
- package/dist/do-oauth-client-provider-CswoD5Lu.js.map +0 -1
- 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.
|
|
552
|
+
this.server.registerTool(
|
|
553
553
|
"add",
|
|
554
|
-
|
|
555
|
-
|
|
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
|
|
package/dist/ai-chat-agent.d.ts
CHANGED
|
@@ -1,14 +1,9 @@
|
|
|
1
|
-
import "./client-
|
|
2
|
-
import "./mcp-
|
|
3
|
-
import "./do-oauth-client-provider-
|
|
4
|
-
import "./index-
|
|
5
|
-
import "./ai-types-
|
|
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
|
}
|
package/dist/ai-chat-agent.js
CHANGED
|
@@ -1,11 +1,22 @@
|
|
|
1
|
-
import { t as MessageType } from "./ai-types-
|
|
2
|
-
import "./client-
|
|
3
|
-
import "./client-
|
|
4
|
-
import "./do-oauth-client-provider-
|
|
5
|
-
import { t as Agent } from "./src-
|
|
6
|
-
import { r as autoTransformMessages } from "./ai-chat-v5-migration-
|
|
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:
|
|
703
|
+
body: chunkBody,
|
|
472
704
|
done: false,
|
|
473
705
|
id,
|
|
474
706
|
type: MessageType.CF_AGENT_USE_CHAT_RESPONSE
|
|
475
707
|
});
|
|
476
|
-
} catch (
|
|
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:
|
|
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
|
};
|