History

Fix cross-session MCP resume leases (#1314)

Rhys Sullivan committed 4 months ago

Commit 52f30f01 parent a50ec9d2 files · +210 −12Browse files

Unified · Split

2 files
  1. apps/host-cloudflare/src/worker.e2e.node.test.ts +92 −0
  2. packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts +118 −12
apps/host-cloudflare/src/worker.e2e.node.test.ts+92−0
@@ -325,4 +325,96 @@ describe("cloudflare host e2e (workerd/miniflare)", () => { · 324 unmodified lines
325325 }>(call);
326326 expect(result.result?.structuredContent?.result).toBe(42);
327327 }, 60_000);
328+
329+ it("resumes a model-paused execution from a fresh MCP session", async () => {
330+ const accept = "application/json, text/event-stream";
331+ const rpc = (sessionId: string | null, body: unknown) =>
332+ worker.fetch("/mcp?elicitation_mode=model", {
333+ method: "POST",
334+ headers: {
335+ "content-type": "application/json",
336+ accept,
337+ ...(sessionId ? { "mcp-session-id": sessionId } : {}),
338+ },
339+ body: JSON.stringify(body),
340+ });
341+ const initialize = async (id: number): Promise<string> => {
342+ const init = await rpc(null, {
343+ jsonrpc: "2.0",
344+ id,
345+ method: "initialize",
346+ params: {
347+ protocolVersion: "2025-03-26",
348+ capabilities: {},
349+ clientInfo: { name: "test", version: "1" },
350+ },
351+ });
352+ expect(init.status).toBe(200);
353+ const sessionId = init.headers.get("mcp-session-id");
354+ expect(sessionId).toBeTruthy();
355+ await rpc(sessionId, {
356+ jsonrpc: "2.0",
357+ method: "notifications/initialized",
358+ });
359+ return sessionId!;
360+ };
361+
362+ const sessionA = await initialize(1);
363+ const execute = await rpc(sessionA, {
364+ jsonrpc: "2.0",
365+ id: 2,
366+ method: "tools/call",
367+ params: {
368+ name: "execute",
369+ arguments: {
370+ code: [
371+ "return await tools.executor.coreTools.policies.create({",
372+ ' owner: "org",',
373+ ' pattern: "microsoft.*",',
374+ ' action: "require_approval"',
375+ "});",
376+ ].join("\n"),
377+ },
378+ },
379+ });
380+ expect(execute.status).toBe(200);
381+ const paused = await readMcpJson<{
382+ result?: {
383+ structuredContent?: {
384+ status?: string;
385+ executionId?: string;
386+ };
387+ };
388+ }>(execute);
389+ expect(paused.result?.structuredContent?.status).toBe("waiting_for_interaction");
390+ const executionId = paused.result?.structuredContent?.executionId;
391+ expect(executionId).toBeTruthy();
392+
393+ const sessionB = await initialize(3);
394+ expect(sessionB).not.toBe(sessionA);
395+ const resume = await rpc(sessionB, {
396+ jsonrpc: "2.0",
397+ id: 4,
398+ method: "tools/call",
399+ params: {
400+ name: "resume",
401+ arguments: { executionId, action: "accept", content: "{}" },
402+ },
403+ });
404+ expect(resume.status).toBe(200);
405+ const resumed = await readMcpJson<{
406+ result?: {
407+ isError?: boolean;
408+ structuredContent?: {
409+ status?: string;
410+ recovery?: string;
411+ result?: unknown;
412+ };
413+ };
414+ }>(resume);
415+ expect(resumed.result?.structuredContent?.status).not.toBe("execution_not_found");
416+ expect(resumed.result?.structuredContent?.recovery).not.toBe("re_execute");
417+ expect(resumed.result?.isError).toBeFalsy();
418+ expect(resumed.result?.structuredContent?.status).toBe("completed");
419+ }, 60_000);
328420 });
packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts+118−12
@@ -436,6 +436,82 @@ export abstract class McpAgentSessionDOBase< · 435 unmodified lines
436436 });
437437 }
438438
439+ private logExecutionOwnerDirectoryFailure(input: {
440+ readonly operation: "put" | "get" | "delete";
441+ readonly executionId: string;
442+ readonly cause: Cause.Cause<unknown>;
443+ }): Effect.Effect<void> {
444+ const self = this;
445+ return Effect.gen(function* () {
446+ const first = Cause.prettyErrors(input.cause)[0];
447+ console.error(
448+ JSON.stringify({
449+ event: "mcp_execution_owner_directory_error",
450+ operation: input.operation,
451+ executionId: input.executionId,
452+ sessionId: self.sessionId,
453+ exceptionType: first?.name ?? "Error",
454+ exceptionMessage: first?.message ?? "unknown",
455+ cause: Cause.pretty(input.cause),
456+ }),
457+ );
458+ yield* Effect.annotateCurrentSpan({
459+ "mcp.execution_owner.directory.operation": input.operation,
460+ });
461+ yield* self.recordCauseOnSpan(input.cause);
462+ });
463+ }
464+
465+ private logModelResumeForwardFailure(input: {
466+ readonly executionId: string;
467+ readonly owner: McpExecutionOwnerRoute;
468+ readonly cause: Cause.Cause<unknown>;
469+ }): Effect.Effect<void> {
470+ const self = this;
471+ return Effect.gen(function* () {
472+ const first = Cause.prettyErrors(input.cause)[0];
473+ console.error(
474+ JSON.stringify({
475+ event: "mcp_model_resume_forward_error",
476+ executionId: input.executionId,
477+ sessionId: self.sessionId,
478+ ownerSessionId: input.owner.sessionId,
479+ exceptionType: first?.name ?? "Error",
480+ exceptionMessage: first?.message ?? "unknown",
481+ cause: Cause.pretty(input.cause),
482+ }),
483+ );
484+ yield* Effect.annotateCurrentSpan({
485+ "mcp.execution_owner.forward.owner_session_id": input.owner.sessionId,
486+ });
487+ yield* self.recordCauseOnSpan(input.cause);
488+ });
489+ }
490+
491+ private logModelResumeForwardTimeout(input: {
492+ readonly executionId: string;
493+ readonly owner: McpExecutionOwnerRoute;
494+ readonly timeoutMs: number;
495+ }): Effect.Effect<void> {
496+ const self = this;
497+ return Effect.gen(function* () {
498+ console.error(
499+ JSON.stringify({
500+ event: "mcp_model_resume_forward_error",
501+ reason: "timeout",
502+ executionId: input.executionId,
503+ sessionId: self.sessionId,
504+ ownerSessionId: input.owner.sessionId,
505+ timeoutMs: input.timeoutMs,
506+ }),
507+ );
508+ yield* Effect.annotateCurrentSpan({
509+ "mcp.execution_owner.forward.owner_session_id": input.owner.sessionId,
510+ "mcp.execution_owner.forward.error": "timeout",
511+ });
512+ });
513+ }
514+
439515 private withSpanFlush<A, E>(effect: Effect.Effect<A, E>): Effect.Effect<A, E> {
440516 const self = this;
441517 return effect.pipe(Effect.ensuring(Effect.promise(() => self.flushTelemetry())));
@@ -813,6 +889,9 @@ export abstract class McpAgentSessionDOBase< · 371 unmodified lines
813889 Effect.map((record) =>
814890 record ? { expiresAt: record.expiresAt, ttlMs: record.ttlMs } : undefined,
815891 ),
892+ Effect.tapCause((cause) =>
893+ this.logExecutionOwnerDirectoryFailure({ operation: "get", executionId, cause }),
894+ ),
816895 Effect.catchCause(() => noDeadline),
817896 );
818897 }
@@ -835,13 +914,26 @@ export abstract class McpAgentSessionDOBase< · 16 unmodified lines
835914 expiresAt: deadline.expiresAt,
836915 ttlMs: deadline.ttlMs,
837916 };
838− yield* directory.put(record);
917+ yield* directory
918+ .put(record)
919+ .pipe(
920+ Effect.tapCause((cause) =>
921+ self.logExecutionOwnerDirectoryFailure({ operation: "put", executionId, cause }),
922+ ),
923+ );
839924 }).pipe(Effect.ignore);
840925 }
841926
842927 private deleteExecutionOwnerEntry(executionId: string): Effect.Effect<void> {
843928 const directory = this.executionOwnerDirectory();
844− return directory?.delete(executionId).pipe(Effect.ignore) ?? Effect.void;
929+ return (
930+ directory?.delete(executionId).pipe(
931+ Effect.tapCause((cause) =>
932+ this.logExecutionOwnerDirectoryFailure({ operation: "delete", executionId, cause }),
933+ ),
934+ Effect.ignore,
935+ ) ?? Effect.void
936+ );
845937 }
846938
847939 private resumeEngineWithLifecycle(
@@ -864,9 +956,12 @@ export abstract class McpAgentSessionDOBase< · 16 unmodified lines
864956 if (!directory) return Effect.succeed(null);
865957 const self = this;
866958 return Effect.gen(function* () {
867− const record = yield* directory
868− .get(executionId)
869− .pipe(Effect.catchCause(() => Effect.succeed(null)));
959+ const record = yield* directory.get(executionId).pipe(
960+ Effect.tapCause((cause) =>
961+ self.logExecutionOwnerDirectoryFailure({ operation: "get", executionId, cause }),
962+ ),
963+ Effect.catchCause(() => Effect.succeed(null)),
964+ );
870965 if (!record) return null;
871966
872967 const sessionMeta = yield* self.loadSessionMeta();
@@ -893,10 +988,24 @@ export abstract class McpAgentSessionDOBase< · 20 unmodified lines
893988 Effect.timeoutOrElse({
894989 duration: `${MODEL_RESUME_FORWARD_TIMEOUT_MS} millis`,
895990 orElse: () =>
896− Effect.succeed({ status: "execution_expired" as const, ttlMs: record.ttlMs }),
991+ Effect.gen(function* () {
992+ yield* self.logModelResumeForwardTimeout({
993+ executionId,
994+ owner: record.owner,
995+ timeoutMs: MODEL_RESUME_FORWARD_TIMEOUT_MS,
996+ });
997+ return { status: "execution_expired" as const, ttlMs: record.ttlMs };
998+ }),
897999 }),
898− Effect.catchCause(() =>
899− Effect.succeed({ status: "execution_expired" as const, ttlMs: record.ttlMs }),
1000+ Effect.catchCause((cause) =>
1001+ Effect.gen(function* () {
1002+ yield* self.logModelResumeForwardFailure({
1003+ executionId,
1004+ owner: record.owner,
1005+ cause,
1006+ });
1007+ return { status: "execution_expired" as const, ttlMs: record.ttlMs };
1008+ }),
9001009 ),
9011010 );
9021011 if (
@@ -994,10 +1103,7 @@ export abstract class McpAgentSessionDOBase< · 91 unmodified lines
9941103 const self = this;
9951104 return Effect.gen(function* () {
9961105 const lease = self.pendingApprovalLeases.get(executionId);
997− if (!lease) {
998− yield* self.deleteExecutionOwnerEntry(executionId);
999− return;
1000− }
1106+ if (!lease) return;
10011107 if (lease.timeout) clearTimeout(lease.timeout);
10021108 self.pendingApprovalLeases.delete(executionId);
10031109 lease.disposeKeepAlive();