machine-bridge-mcp 3.0.0-beta.20 → 3.0.0-beta.21

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 (44) hide show
  1. package/CHANGELOG.md +16 -0
  2. package/README.md +2 -0
  3. package/SECURITY.md +4 -4
  4. package/browser-extension/manifest.json +1 -1
  5. package/docs/ARCHITECTURE.md +11 -9
  6. package/docs/AUDIT.md +16 -0
  7. package/docs/ENGINEERING.md +1 -1
  8. package/docs/LOCAL_AUTHORIZATION.md +1 -1
  9. package/docs/LOGGING.md +3 -1
  10. package/docs/MULTI_ACCOUNT.md +1 -1
  11. package/docs/OPERATIONS.md +5 -3
  12. package/docs/OVERVIEW.md +1 -1
  13. package/docs/PRIVACY.md +1 -1
  14. package/docs/TESTING.md +9 -7
  15. package/docs/THREAT_MODEL.md +3 -2
  16. package/docs/UPGRADING.md +7 -1
  17. package/package.json +1 -1
  18. package/src/local/process-tree-ownership.mjs +17 -5
  19. package/src/local/process-tree.mjs +6 -2
  20. package/src/local/relay-connection.mjs +24 -11
  21. package/src/local/runtime-activation.mjs +16 -0
  22. package/src/local/runtime.mjs +3 -1
  23. package/src/local/worker-deployment.mjs +16 -6
  24. package/src/worker/daemon-socket-attachment.ts +52 -0
  25. package/src/worker/daemon-sockets.ts +17 -48
  26. package/src/worker/durable-stream-calls.ts +129 -0
  27. package/src/worker/durable-stream-result.ts +22 -0
  28. package/src/worker/index.ts +120 -88
  29. package/src/worker/mcp-pending-call-expiry.ts +20 -0
  30. package/src/worker/mcp-pending-call-inspection.ts +33 -0
  31. package/src/worker/mcp-pending-call-records.ts +58 -0
  32. package/src/worker/mcp-pending-call-storage.ts +18 -0
  33. package/src/worker/mcp-pending-call-store.ts +285 -0
  34. package/src/worker/mcp-resumption-index.ts +41 -0
  35. package/src/worker/mcp-resumption-records.ts +8 -2
  36. package/src/worker/mcp-resumption.ts +80 -87
  37. package/src/worker/mcp-stream-dispatch.ts +42 -54
  38. package/src/worker/pending-admission.ts +15 -0
  39. package/src/worker/pending-call-contract.ts +4 -7
  40. package/src/worker/pending-calls.ts +4 -10
  41. package/src/worker/runtime-alarm-storage.ts +35 -0
  42. package/src/worker/runtime-alarm.ts +29 -39
  43. package/src/worker/websocket-protocol.ts +4 -0
  44. package/src/worker/worker-entry.ts +50 -0
@@ -0,0 +1,285 @@
1
+ import {
2
+ STREAM_INDEX_KEY,
3
+ indexEntry,
4
+ readIndex,
5
+ streamKey,
6
+ validRecord,
7
+ type JsonRpcMessage,
8
+ type StreamIndex,
9
+ type StreamIndexEntry,
10
+ type StreamRecord,
11
+ } from "./mcp-resumption-records.ts";
12
+ import type { PendingStreamCall, PendingStreamTransform } from "./mcp-pending-call-records.ts";
13
+ import {
14
+ extendAttachedCallExpiry,
15
+ extendDetachedCallExpiry,
16
+ positiveDelay,
17
+ } from "./mcp-pending-call-expiry.ts";
18
+ import { requiredPendingCallRecord } from "./mcp-pending-call-storage.ts";
19
+ import {
20
+ activePendingCallCount,
21
+ pendingCallSnapshot,
22
+ type PendingStreamCallSnapshot,
23
+ } from "./mcp-pending-call-inspection.ts";
24
+
25
+ type PendingCallStorage = Pick<DurableObjectStorage, "get" | "put" | "transaction">;
26
+ type CompleteStream = (
27
+ streamId: string,
28
+ message: JsonRpcMessage,
29
+ expected: { callId: string; connectionId?: string },
30
+ ) => Promise<boolean>;
31
+ type StorageRowsWrittenListener = (rows: number) => void;
32
+
33
+ export type PendingStreamCallView = PendingStreamCall & {
34
+ streamId: string;
35
+ requestId: string | number;
36
+ };
37
+
38
+ export class McpPendingCallLimitError extends Error {
39
+ constructor() {
40
+ super("too many concurrent daemon tool calls");
41
+ this.name = "McpPendingCallLimitError";
42
+ }
43
+ }
44
+
45
+ export class McpPendingCallConflictError extends Error {
46
+ constructor(message: string) {
47
+ super(message);
48
+ this.name = "McpPendingCallConflictError";
49
+ }
50
+ }
51
+
52
+ const DEFAULT_MAXIMUM_PENDING_CALLS = 32;
53
+
54
+ export class McpPendingCallStore {
55
+ private readonly storage: PendingCallStorage;
56
+ private readonly now: () => number;
57
+ private readonly activateStream: (streamId: string) => void;
58
+ private readonly completeStream: CompleteStream;
59
+ private readonly onRowsWritten: StorageRowsWrittenListener;
60
+ private readonly terminalRetentionMs: number;
61
+
62
+ constructor(
63
+ storage: PendingCallStorage,
64
+ now: () => number,
65
+ activateStream: (streamId: string) => void,
66
+ completeStream: CompleteStream,
67
+ onRowsWritten: StorageRowsWrittenListener,
68
+ terminalRetentionMs: number,
69
+ ) {
70
+ this.storage = storage;
71
+ this.now = now;
72
+ this.activateStream = activateStream;
73
+ this.completeStream = completeStream;
74
+ this.onRowsWritten = onRowsWritten;
75
+ this.terminalRetentionMs = positiveDelay(terminalRetentionMs);
76
+ }
77
+
78
+ async activate(input: {
79
+ streamId: string;
80
+ callId: string;
81
+ daemonInstanceId: string;
82
+ connectionId: string;
83
+ clientRequestKey?: string;
84
+ tool: string;
85
+ timeoutMs: number;
86
+ transform?: PendingStreamTransform;
87
+ transientActiveCount?: number;
88
+ maximumPendingCalls?: number;
89
+ }): Promise<void> {
90
+ const now = this.now();
91
+ const timeoutMs = positiveDelay(input.timeoutMs);
92
+ const maximum = Math.max(1, Math.floor(input.maximumPendingCalls ?? DEFAULT_MAXIMUM_PENDING_CALLS));
93
+ const call: PendingStreamCall = {
94
+ call_id: input.callId,
95
+ daemon_instance_id: input.daemonInstanceId,
96
+ connection_id: input.connectionId,
97
+ ...(input.clientRequestKey ? { client_request_key: input.clientRequestKey } : {}),
98
+ tool: input.tool,
99
+ state: "attached",
100
+ started_at: now,
101
+ operation_deadline_at: now + timeoutMs,
102
+ remaining_timeout_ms: timeoutMs,
103
+ ...(input.transform ? { transform: structuredClone(input.transform) } : {}),
104
+ };
105
+ await this.storage.transaction(async (transaction) => {
106
+ const index = readIndex(await transaction.get<unknown>(STREAM_INDEX_KEY));
107
+ const activeCalls = index.entries.filter((entry) => entry.call).length;
108
+ if (activeCalls + Math.max(0, Math.floor(input.transientActiveCount ?? 0)) >= maximum) {
109
+ throw new McpPendingCallLimitError();
110
+ }
111
+ if (index.entries.some((entry) => entry.call?.call_id === input.callId)) {
112
+ throw new McpPendingCallConflictError("duplicate internal daemon call id");
113
+ }
114
+ if (input.clientRequestKey && index.entries.some((entry) => entry.call?.client_request_key === input.clientRequestKey)) {
115
+ throw new McpPendingCallConflictError("duplicate in-flight JSON-RPC request id within this MCP session");
116
+ }
117
+ const record = await transaction.get<unknown>(streamKey(input.streamId));
118
+ if (!validRecord(record) || record.status !== "pending") {
119
+ throw new Error("resumable MCP stream record is unavailable for activation");
120
+ }
121
+ if (record.call) throw new McpPendingCallConflictError("resumable MCP stream already has an active daemon call");
122
+ const updated: StreamRecord = {
123
+ ...record,
124
+ expires_at: extendAttachedCallExpiry(record.expires_at, call.operation_deadline_at, this.terminalRetentionMs),
125
+ call,
126
+ };
127
+ const entries = index.entries.map((entry) => entry.stream_id === input.streamId ? indexEntry(updated) : entry);
128
+ if (!entries.some((entry) => entry.stream_id === input.streamId)) {
129
+ throw new Error("resumable MCP stream index lost its record");
130
+ }
131
+ await transaction.put(streamKey(input.streamId), updated);
132
+ await transaction.put(STREAM_INDEX_KEY, { schema_version: 1, entries } satisfies StreamIndex);
133
+ });
134
+ this.onRowsWritten(2);
135
+ this.activateStream(input.streamId);
136
+ }
137
+
138
+ async activeCount(): Promise<number> {
139
+ return activePendingCallCount(readIndex(await this.storage.get<unknown>(STREAM_INDEX_KEY)).entries);
140
+ }
141
+
142
+ async get(callId: string): Promise<PendingStreamCallView | undefined> {
143
+ const index = readIndex(await this.storage.get<unknown>(STREAM_INDEX_KEY));
144
+ const entry = index.entries.find((candidate) => candidate.call?.call_id === callId);
145
+ return entry ? this.callView(entry) : undefined;
146
+ }
147
+
148
+ async getByRequestKey(requestKey: string): Promise<PendingStreamCallView | undefined> {
149
+ const index = readIndex(await this.storage.get<unknown>(STREAM_INDEX_KEY));
150
+ const entry = index.entries.find((candidate) => candidate.call?.client_request_key === requestKey);
151
+ return entry ? this.callView(entry) : undefined;
152
+ }
153
+
154
+ async snapshot(maximum = DEFAULT_MAXIMUM_PENDING_CALLS): Promise<PendingStreamCallSnapshot> {
155
+ const index = readIndex(await this.storage.get<unknown>(STREAM_INDEX_KEY));
156
+ return pendingCallSnapshot(index.entries, this.now(), maximum);
157
+ }
158
+
159
+ async nextDeadlineDelayMs(): Promise<number> {
160
+ const index = readIndex(await this.storage.get<unknown>(STREAM_INDEX_KEY));
161
+ const now = this.now();
162
+ return Math.min(Number.POSITIVE_INFINITY, ...index.entries.flatMap((entry) => {
163
+ if (!entry.call) return [];
164
+ const deadline = entry.call.state === "attached"
165
+ ? entry.call.operation_deadline_at
166
+ : entry.call.reconnect_deadline_at;
167
+ return Number.isSafeInteger(deadline) ? [Math.max(0, Number(deadline) - now)] : [];
168
+ }));
169
+ }
170
+
171
+ async due(now = this.now()): Promise<PendingStreamCallView[]> {
172
+ const index = readIndex(await this.storage.get<unknown>(STREAM_INDEX_KEY));
173
+ const due: PendingStreamCallView[] = [];
174
+ for (const entry of index.entries) {
175
+ const call = entry.call;
176
+ if (!call) continue;
177
+ const deadline = call.state === "attached" ? call.operation_deadline_at : call.reconnect_deadline_at;
178
+ if (Number.isSafeInteger(deadline) && Number(deadline) <= now) due.push(await this.callView(entry));
179
+ }
180
+ return due;
181
+ }
182
+
183
+ async detach(connectionId: string, graceMs: number): Promise<number> {
184
+ const now = this.now();
185
+ const reconnectDeadline = now + positiveDelay(graceMs);
186
+ let rows = 0;
187
+ const count = await this.storage.transaction(async (transaction) => {
188
+ const index = readIndex(await transaction.get<unknown>(STREAM_INDEX_KEY));
189
+ const entries = [...index.entries];
190
+ let changed = 0;
191
+ for (let position = 0; position < entries.length; position += 1) {
192
+ const entry = entries[position];
193
+ if (entry.call?.state !== "attached" || entry.call.connection_id !== connectionId) continue;
194
+ const record = await requiredPendingCallRecord(transaction, entry as StreamIndexEntry & { call: PendingStreamCall });
195
+ const remaining = Math.max(1, record.call.operation_deadline_at - now);
196
+ const updated: StreamRecord = {
197
+ ...record,
198
+ expires_at: extendDetachedCallExpiry(
199
+ record.expires_at,
200
+ reconnectDeadline,
201
+ remaining,
202
+ this.terminalRetentionMs,
203
+ ),
204
+ call: {
205
+ ...record.call,
206
+ state: "detached",
207
+ remaining_timeout_ms: remaining,
208
+ reconnect_deadline_at: reconnectDeadline,
209
+ },
210
+ };
211
+ await transaction.put(streamKey(entry.stream_id), updated);
212
+ entries[position] = indexEntry(updated);
213
+ changed += 1;
214
+ }
215
+ if (changed > 0) {
216
+ await transaction.put(STREAM_INDEX_KEY, { schema_version: 1, entries } satisfies StreamIndex);
217
+ rows = changed + 1;
218
+ }
219
+ return changed;
220
+ });
221
+ this.onRowsWritten(rows);
222
+ return count;
223
+ }
224
+
225
+ async rebind(daemonInstanceId: string, connectionId: string): Promise<string[]> {
226
+ if (!daemonInstanceId || !connectionId) return [];
227
+ const now = this.now();
228
+ let rows = 0;
229
+ const rebound = await this.storage.transaction(async (transaction) => {
230
+ const index = readIndex(await transaction.get<unknown>(STREAM_INDEX_KEY));
231
+ const entries = [...index.entries];
232
+ const ids: string[] = [];
233
+ for (let position = 0; position < entries.length; position += 1) {
234
+ const entry = entries[position];
235
+ if (!entry.call || entry.call.daemon_instance_id !== daemonInstanceId || entry.call.connection_id === connectionId) continue;
236
+ const record = await requiredPendingCallRecord(transaction, entry as StreamIndexEntry & { call: PendingStreamCall });
237
+ const remaining = record.call.state === "detached"
238
+ ? record.call.remaining_timeout_ms
239
+ : Math.max(1, record.call.operation_deadline_at - now);
240
+ const updatedCall: PendingStreamCall = {
241
+ ...record.call,
242
+ connection_id: connectionId,
243
+ state: "attached",
244
+ operation_deadline_at: now + remaining,
245
+ remaining_timeout_ms: remaining,
246
+ };
247
+ delete updatedCall.reconnect_deadline_at;
248
+ const updated: StreamRecord = {
249
+ ...record,
250
+ expires_at: extendAttachedCallExpiry(
251
+ record.expires_at,
252
+ updatedCall.operation_deadline_at,
253
+ this.terminalRetentionMs,
254
+ ),
255
+ call: updatedCall,
256
+ };
257
+ await transaction.put(streamKey(entry.stream_id), updated);
258
+ entries[position] = indexEntry(updated);
259
+ ids.push(updatedCall.call_id);
260
+ }
261
+ if (ids.length > 0) {
262
+ await transaction.put(STREAM_INDEX_KEY, { schema_version: 1, entries } satisfies StreamIndex);
263
+ rows = ids.length + 1;
264
+ }
265
+ return ids;
266
+ });
267
+ this.onRowsWritten(rows);
268
+ return rebound;
269
+ }
270
+
271
+ async complete(callId: string, connectionId: string | undefined, message: JsonRpcMessage): Promise<boolean> {
272
+ const view = await this.get(callId);
273
+ if (!view || (connectionId && view.connection_id !== connectionId)) return false;
274
+ return this.completeStream(view.streamId, message, { callId, connectionId });
275
+ }
276
+
277
+ private async callView(entry: StreamIndexEntry): Promise<PendingStreamCallView> {
278
+ if (!entry.call) throw new Error("stream index entry has no pending call");
279
+ const record = await this.storage.get<unknown>(streamKey(entry.stream_id));
280
+ if (!validRecord(record) || !record.call || record.call.call_id !== entry.call.call_id) {
281
+ throw new Error("persisted pending call record is corrupt");
282
+ }
283
+ return { ...structuredClone(record.call), streamId: record.stream_id, requestId: record.request_id };
284
+ }
285
+ }
@@ -0,0 +1,41 @@
1
+ import {
2
+ streamKey,
3
+ type StreamIndexEntry,
4
+ } from "./mcp-resumption-records.ts";
5
+
6
+ type TransactionStorage = Pick<DurableObjectTransaction, "delete">;
7
+
8
+ export async function pruneExpiredStreams(
9
+ storage: TransactionStorage,
10
+ entries: StreamIndexEntry[],
11
+ now: number,
12
+ ): Promise<{ entries: StreamIndexEntry[]; removedStreamIds: string[] }> {
13
+ const retained: StreamIndexEntry[] = [];
14
+ const removedStreamIds: string[] = [];
15
+ for (const entry of entries) {
16
+ if (entry.expires_at <= now) {
17
+ await storage.delete(streamKey(entry.stream_id));
18
+ removedStreamIds.push(entry.stream_id);
19
+ } else retained.push(entry);
20
+ }
21
+ return { entries: retained, removedStreamIds };
22
+ }
23
+
24
+ export async function freeCompletedStreamSlots(
25
+ storage: TransactionStorage,
26
+ entries: StreamIndexEntry[],
27
+ maximumRetained: number,
28
+ ): Promise<{ available: boolean; entries: StreamIndexEntry[]; removedStreamIds: string[] }> {
29
+ const retained = [...entries];
30
+ const removedStreamIds: string[] = [];
31
+ const completed = retained
32
+ .filter((entry) => entry.status === "ready")
33
+ .sort((left, right) => left.created_at - right.created_at);
34
+ for (const entry of completed) {
35
+ if (retained.length <= maximumRetained) break;
36
+ await storage.delete(streamKey(entry.stream_id));
37
+ retained.splice(retained.findIndex((candidate) => candidate.stream_id === entry.stream_id), 1);
38
+ removedStreamIds.push(entry.stream_id);
39
+ }
40
+ return { available: retained.length <= maximumRetained, entries: retained, removedStreamIds };
41
+ }
@@ -1,4 +1,5 @@
1
1
  import relayContract from "../shared/relay-contract.json" with { type: "json" };
2
+ import { validPendingStreamCall, type PendingStreamCall } from "./mcp-pending-call-records.ts";
2
3
  export type JsonRpcMessage = Record<string, unknown> | null;
3
4
  export type StreamStatus = "pending" | "ready";
4
5
  export type StreamIndexEntry = {
@@ -6,6 +7,7 @@ export type StreamIndexEntry = {
6
7
  status: StreamStatus;
7
8
  created_at: number;
8
9
  expires_at: number;
10
+ call?: PendingStreamCall;
9
11
  };
10
12
  export type StreamIndex = { schema_version: 1; entries: StreamIndexEntry[] };
11
13
  export type StreamRecord = StreamIndexEntry & {
@@ -55,6 +57,7 @@ export function validRecord(value: unknown): value is StreamRecord {
55
57
  return false;
56
58
  }
57
59
  if (record.status === "pending") return record.message_json === undefined && record.message_sha256 === undefined;
60
+ if (record.call !== undefined) return false;
58
61
  return typeof record.message_json === "string"
59
62
  && new TextEncoder().encode(record.message_json).byteLength <= DEFAULT_MAXIMUM_MESSAGE_BYTES
60
63
  && typeof record.message_sha256 === "string"
@@ -67,6 +70,7 @@ export function indexEntry(record: StreamRecord): StreamIndexEntry {
67
70
  status: record.status,
68
71
  created_at: record.created_at,
69
72
  expires_at: record.expires_at,
73
+ ...(record.call ? { call: structuredClone(record.call) } : {}),
70
74
  };
71
75
  }
72
76
 
@@ -90,8 +94,9 @@ export function readyRecord(
90
94
  messageSha256: string,
91
95
  expiresAt: number,
92
96
  ): StreamRecord {
97
+ const { call: _call, ...base } = record;
93
98
  return {
94
- ...record,
99
+ ...base,
95
100
  status: "ready",
96
101
  expires_at: expiresAt,
97
102
  message_json: messageJson,
@@ -150,7 +155,8 @@ function validIndexEntry(value: unknown): value is StreamIndexEntry {
150
155
  return typeof entry.stream_id === "string" && isStreamId(entry.stream_id)
151
156
  && (entry.status === "pending" || entry.status === "ready")
152
157
  && Number.isSafeInteger(entry.created_at) && Number.isSafeInteger(entry.expires_at)
153
- && entry.expires_at! > entry.created_at!;
158
+ && entry.expires_at! > entry.created_at!
159
+ && (entry.call === undefined || (entry.status === "pending" && validPendingStreamCall(entry.call)));
154
160
  }
155
161
 
156
162
  function storedIntegrityMessage(requestId: string | number): Record<string, unknown> {
@@ -13,28 +13,32 @@ import {
13
13
  workerRestartMessage,
14
14
  type JsonRpcMessage,
15
15
  type StreamIndex,
16
- type StreamIndexEntry,
17
16
  type StreamRecord,
18
17
  } from "./mcp-resumption-records.ts";
18
+ import { McpPendingCallStore } from "./mcp-pending-call-store.ts";
19
+ import { freeCompletedStreamSlots, pruneExpiredStreams } from "./mcp-resumption-index.ts";
19
20
  import { resumptionLimits, type McpResumptionOptions } from "./mcp-resumption-config.ts";
20
21
  export type { JsonRpcMessage } from "./mcp-resumption-records.ts";
22
+
21
23
  type ResumptionStorage = Pick<DurableObjectStorage, "get" | "put" | "delete" | "transaction">;
22
- type TransactionStorage = Pick<DurableObjectTransaction, "get" | "put" | "delete">;
23
24
  type StreamReadyListener = (streamId: string, message: JsonRpcMessage) => void;
24
25
  type StorageRowsWrittenListener = (rows: number) => void;
25
26
  export type StreamResumeResult =
26
27
  | { kind: "invalid" | "not_found" | "expired" }
27
28
  | { kind: "complete"; streamId: string }
28
29
  | { kind: "message"; streamId: string };
30
+
29
31
  export class McpStreamLimitError extends Error {
30
32
  constructor() {
31
33
  super("too many resumable MCP streams are active");
32
34
  this.name = "McpStreamLimitError";
33
35
  }
34
36
  }
37
+
35
38
  const EVENT_ID_PATTERN = /^(stream_[A-Za-z0-9_-]{43}):([01])$/;
39
+
36
40
  export class McpResumptionStore {
37
- private readonly storage: ResumptionStorage;
41
+ readonly calls: McpPendingCallStore;
38
42
  private readonly active = new Set<string>();
39
43
  private readonly transientReady = new Map<string, JsonRpcMessage>();
40
44
  private readonly now: () => number;
@@ -42,20 +46,35 @@ export class McpResumptionStore {
42
46
  private readonly pendingRetentionMs: number;
43
47
  private readonly maximumStreams: number;
44
48
  private readonly maximumMessageBytes: number;
49
+ private readonly storage: ResumptionStorage;
45
50
  private readonly onReady: StreamReadyListener;
46
51
  private readonly onRowsWritten: StorageRowsWrittenListener;
47
- constructor(storage: ResumptionStorage, options: McpResumptionOptions = {},
48
- onReady: StreamReadyListener = () => {}, onRowsWritten: StorageRowsWrittenListener = () => {}) {
52
+
53
+ constructor(
54
+ storage: ResumptionStorage,
55
+ options: McpResumptionOptions = {},
56
+ onReady: StreamReadyListener = () => {},
57
+ onRowsWritten: StorageRowsWrittenListener = () => {},
58
+ ) {
49
59
  this.storage = storage;
50
- this.now = options.now ?? Date.now;
51
60
  this.onReady = onReady;
52
61
  this.onRowsWritten = onRowsWritten;
62
+ this.now = options.now ?? Date.now;
53
63
  const limits = resumptionLimits(options);
54
64
  this.retentionMs = limits.retentionMs;
55
65
  this.pendingRetentionMs = limits.pendingRetentionMs;
56
66
  this.maximumStreams = limits.maximumStreams;
57
67
  this.maximumMessageBytes = limits.maximumMessageBytes;
68
+ this.calls = new McpPendingCallStore(
69
+ storage,
70
+ this.now,
71
+ (streamId) => this.activate(streamId),
72
+ (streamId, message, expected) => this.completeInternal(streamId, message, expected),
73
+ onRowsWritten,
74
+ this.retentionMs,
75
+ );
58
76
  }
77
+
59
78
  async begin(input: {
60
79
  streamId: string;
61
80
  tokenKey: string;
@@ -76,8 +95,8 @@ export class McpResumptionStore {
76
95
  };
77
96
  const removedStreamIds = await this.storage.transaction(async (transaction) => {
78
97
  const index = readIndex(await transaction.get<unknown>(STREAM_INDEX_KEY));
79
- const pruned = await pruneExpired(transaction, index.entries, now);
80
- const freed = await freeCompletedSlots(transaction, pruned.entries, this.maximumStreams - 1);
98
+ const pruned = await pruneExpiredStreams(transaction, index.entries, now);
99
+ const freed = await freeCompletedStreamSlots(transaction, pruned.entries, this.maximumStreams - 1);
81
100
  if (!freed.available) throw new McpStreamLimitError();
82
101
  const entries = [...freed.entries, indexEntry(record)];
83
102
  await transaction.put(streamKey(input.streamId), record);
@@ -87,6 +106,7 @@ export class McpResumptionStore {
87
106
  this.onRowsWritten(2 + removedStreamIds.length);
88
107
  for (const removedStreamId of removedStreamIds) this.clearMemory(removedStreamId);
89
108
  }
109
+
90
110
  activate(streamId: string): void {
91
111
  if (!isStreamId(streamId)) throw new Error("invalid MCP stream id");
92
112
  this.transientReady.delete(streamId);
@@ -94,43 +114,12 @@ export class McpResumptionStore {
94
114
  }
95
115
 
96
116
  async complete(streamId: string, message: JsonRpcMessage): Promise<void> {
97
- try {
98
- const initial = await this.storage.get<unknown>(streamKey(streamId));
99
- if (initial === undefined) throw new Error("resumable MCP stream record disappeared before completion");
100
- if (!validRecord(initial)) throw new Error("resumable MCP stream record is corrupt");
101
- const messageJson = resumableMessageJson(message, initial.request_id, this.maximumMessageBytes);
102
- const digest = await messageSha256(messageJson);
103
- const expiresAt = this.now() + this.retentionMs;
104
- const rowsWritten = await this.storage.transaction(async (transaction) => {
105
- const record = await transaction.get<unknown>(streamKey(streamId));
106
- if (record === undefined) return -1;
107
- if (!validRecord(record)) throw new Error("resumable MCP stream record is corrupt");
108
- if (record.status === "ready") return 0;
109
- const index = readIndex(await transaction.get<unknown>(STREAM_INDEX_KEY));
110
- if (!index.entries.some((candidate) => candidate.stream_id === streamId)) {
111
- throw new Error("resumable MCP stream index lost its record");
112
- }
113
- const ready = readyRecord(record, messageJson, digest, expiresAt);
114
- const entries = index.entries.map((candidate) => candidate.stream_id === streamId
115
- ? indexEntry(ready)
116
- : candidate);
117
- await transaction.put(streamKey(streamId), ready);
118
- await transaction.put(STREAM_INDEX_KEY, { schema_version: 1, entries } satisfies StreamIndex);
119
- return 2;
120
- });
121
- if (rowsWritten < 0) throw new Error("resumable MCP stream record disappeared during completion");
122
- this.onRowsWritten(rowsWritten);
123
- this.transientReady.delete(streamId);
124
- } catch (error) {
125
- this.transientReady.set(streamId, message);
126
- throw error;
127
- } finally {
128
- this.active.delete(streamId);
129
- try { this.onReady(streamId, message); } catch { /* Persistent state remains authoritative. */ }
130
- }
117
+ await this.completeInternal(streamId, message);
131
118
  }
132
119
 
133
- async pollMessage(streamId: string): Promise<{ kind: "pending" } | { kind: "not_found" } | { kind: "message"; message: JsonRpcMessage }> {
120
+ async pollMessage(streamId: string): Promise<
121
+ { kind: "pending" } | { kind: "not_found" } | { kind: "message"; message: JsonRpcMessage }
122
+ > {
134
123
  if (!isStreamId(streamId)) return { kind: "not_found" };
135
124
  const transient = this.transientReady.get(streamId);
136
125
  if (transient) return { kind: "message", message: transient };
@@ -142,18 +131,13 @@ export class McpResumptionStore {
142
131
  return { kind: "not_found" };
143
132
  }
144
133
  if (record.status === "ready") return { kind: "message", message: await storedMessage(record) };
145
- if (this.active.has(streamId)) return { kind: "pending" };
146
-
134
+ if (record.call || this.active.has(streamId)) return { kind: "pending" };
147
135
  const unavailable = workerRestartMessage(record.request_id);
148
136
  await this.complete(streamId, unavailable);
149
137
  return { kind: "message", message: unavailable };
150
138
  }
151
139
 
152
- async resume(input: {
153
- lastEventId: string;
154
- tokenKey: string;
155
- sessionId: string;
156
- }): Promise<StreamResumeResult> {
140
+ async resume(input: { lastEventId: string; tokenKey: string; sessionId: string }): Promise<StreamResumeResult> {
157
141
  const event = parseStreamEventId(input.lastEventId);
158
142
  if (!event) return { kind: "invalid" };
159
143
  const record = await this.storage.get<unknown>(streamKey(event.streamId));
@@ -165,15 +149,59 @@ export class McpResumptionStore {
165
149
  return { kind: "expired" };
166
150
  }
167
151
  if (event.sequence >= 1) return { kind: "complete", streamId: event.streamId };
168
- if (record.status === "ready" || this.transientReady.has(event.streamId) || this.active.has(event.streamId)) {
152
+ if (record.status === "ready" || this.transientReady.has(event.streamId) || record.call || this.active.has(event.streamId)) {
169
153
  return { kind: "message", streamId: event.streamId };
170
154
  }
171
-
172
155
  const unavailable = workerRestartMessage(record.request_id);
173
156
  await this.complete(event.streamId, unavailable);
174
157
  return { kind: "message", streamId: event.streamId };
175
158
  }
176
159
 
160
+ private async completeInternal(
161
+ streamId: string,
162
+ message: JsonRpcMessage,
163
+ expected?: { callId: string; connectionId?: string },
164
+ ): Promise<boolean> {
165
+ try {
166
+ const initial = await this.storage.get<unknown>(streamKey(streamId));
167
+ if (initial === undefined) throw new Error("resumable MCP stream record disappeared before completion");
168
+ if (!validRecord(initial)) throw new Error("resumable MCP stream record is corrupt");
169
+ const messageJson = resumableMessageJson(message, initial.request_id, this.maximumMessageBytes);
170
+ const digest = await messageSha256(messageJson);
171
+ const expiresAt = this.now() + this.retentionMs;
172
+ const rowsWritten = await this.storage.transaction(async (transaction) => {
173
+ const record = await transaction.get<unknown>(streamKey(streamId));
174
+ if (record === undefined) return -1;
175
+ if (!validRecord(record)) throw new Error("resumable MCP stream record is corrupt");
176
+ if (record.status === "ready") return 0;
177
+ if (expected && (
178
+ record.call?.call_id !== expected.callId
179
+ || (expected.connectionId && record.call.connection_id !== expected.connectionId)
180
+ )) return -2;
181
+ const index = readIndex(await transaction.get<unknown>(STREAM_INDEX_KEY));
182
+ if (!index.entries.some((candidate) => candidate.stream_id === streamId)) {
183
+ throw new Error("resumable MCP stream index lost its record");
184
+ }
185
+ const ready = readyRecord(record, messageJson, digest, expiresAt);
186
+ const entries = index.entries.map((candidate) => candidate.stream_id === streamId ? indexEntry(ready) : candidate);
187
+ await transaction.put(streamKey(streamId), ready);
188
+ await transaction.put(STREAM_INDEX_KEY, { schema_version: 1, entries } satisfies StreamIndex);
189
+ return 2;
190
+ });
191
+ if (rowsWritten < -1) return false;
192
+ if (rowsWritten < 0) throw new Error("resumable MCP stream record disappeared during completion");
193
+ this.onRowsWritten(rowsWritten);
194
+ if (rowsWritten === 0) return false;
195
+ this.clearMemory(streamId);
196
+ try { this.onReady(streamId, message); } catch { /* Persistent state remains authoritative. */ }
197
+ return true;
198
+ } catch (error) {
199
+ this.transientReady.set(streamId, message);
200
+ try { this.onReady(streamId, message); } catch { /* In-memory terminal state remains available. */ }
201
+ throw error;
202
+ }
203
+ }
204
+
177
205
  private async remove(streamId: string): Promise<void> {
178
206
  await this.storage.transaction(async (transaction) => {
179
207
  const index = readIndex(await transaction.get<unknown>(STREAM_INDEX_KEY));
@@ -202,38 +230,3 @@ export function parseStreamEventId(value: string): { streamId: string; sequence:
202
230
  if (!match) return null;
203
231
  return { streamId: match[1], sequence: Number(match[2]) as 0 | 1 };
204
232
  }
205
-
206
- async function pruneExpired(
207
- storage: TransactionStorage,
208
- entries: StreamIndexEntry[],
209
- now: number,
210
- ): Promise<{ entries: StreamIndexEntry[]; removedStreamIds: string[] }> {
211
- const retained: StreamIndexEntry[] = [];
212
- const removedStreamIds: string[] = [];
213
- for (const entry of entries) {
214
- if (entry.expires_at <= now) {
215
- await storage.delete(streamKey(entry.stream_id));
216
- removedStreamIds.push(entry.stream_id);
217
- } else retained.push(entry);
218
- }
219
- return { entries: retained, removedStreamIds };
220
- }
221
-
222
- async function freeCompletedSlots(
223
- storage: TransactionStorage,
224
- entries: StreamIndexEntry[],
225
- maximumRetained: number,
226
- ): Promise<{ available: boolean; entries: StreamIndexEntry[]; removedStreamIds: string[] }> {
227
- const retained = [...entries];
228
- const removedStreamIds: string[] = [];
229
- const completed = retained
230
- .filter((entry) => entry.status === "ready")
231
- .sort((left, right) => left.created_at - right.created_at);
232
- for (const entry of completed) {
233
- if (retained.length <= maximumRetained) break;
234
- await storage.delete(streamKey(entry.stream_id));
235
- retained.splice(retained.findIndex((candidate) => candidate.stream_id === entry.stream_id), 1);
236
- removedStreamIds.push(entry.stream_id);
237
- }
238
- return { available: retained.length <= maximumRetained, entries: retained, removedStreamIds };
239
- }