Keep MCP-Apps rendering across a cold session restore

The negotiated MCP-Apps capability lived only in the tool server instance, set from the client at initialize. A session whose Durable Object was evicted (deploy, idle) and cold-restored on the same mcp-session-id was rebuilt with no initialize to replay, so apps support read back as false and create-artifact/show-artifact answered with the deep-link fallback instead of the inline widget, mid-conversation. Persist the negotiated bit in SessionMeta and seed the rebuilt server from it. A real initialize still overrules the restored value in both directions, and init carries the stored bit through its re-resolve so restoring a session no longer erases it.

Rhys Sullivan committed 3 months ago

Commit ed7bf081 parent a7c46896 files · +325 −4Browse files

Unified · Split

6 files

apps/cloud/src/mcp/session-durable-object.ts +4 −0
apps/host-cloudflare/src/mcp/session-durable-object.ts +4 −0
packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts +90 −0
packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts +39 −0
packages/hosts/mcp/src/artifacts-tools.test.ts +142 −1
packages/hosts/mcp/src/tool-server.ts +46 −3

apps/cloud/src/mcp/session-durable-object.ts+4−0

         // flag existed carry no value; absent now means off, same as a fresh
         // connection that did not ask for `?artifacts=true`.
         artifactsEnabled: sessionMeta.artifactsEnabled ?? false,
+        // Cold restores rebuild this server with no `initialize` to replay, so
+        // the negotiated apps support comes back from storage instead.
+        restoredAppsEnabled: sessionMeta.appsEnabled ?? false,
+        onAppsEnabledChange: (appsEnabled) => self.persistAppsEnabled(appsEnabled),
         loadAppShellHtml,
         smokeRenderArtifact,
         artifactUrl: artifactUrlFor(

apps/host-cloudflare/src/mcp/session-durable-object.ts+4−0

         // flag existed carry no value; absent now means off, same as a fresh
         // connection that did not ask for `?artifacts=true`.
         artifactsEnabled: sessionMeta.artifactsEnabled ?? false,
+        // Cold restores rebuild this server with no `initialize` to replay, so
+        // the negotiated apps support comes back from storage instead.
+        restoredAppsEnabled: sessionMeta.appsEnabled ?? false,
+        onAppsEnabledChange: (appsEnabled) => self.persistAppsEnabled(appsEnabled),
         loadAppShellHtml: self.loadAppShellHtml,
         smokeRenderArtifact,
         ...(artifactOrigin

packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts+90−0

@@ -216,6 +216,96 @@ const makeHarnessSession = async (): Promise<HarnessSession> => { · 215 unmodified lines
216	return session;	216	return session;
217	};	217	};
218		218
		219	// The negotiated MCP-Apps capability arrives once, at `initialize`, and lives
		220	// in the rebuilt server's memory. These pin the storage round-trip that lets a
		221	// cold-restored session rebuild with it instead of silently downgrading every
		222	// artifact to a deep link.
		223	describe("McpAgentSessionDOBase apps capability persistence", () => {
		224	type CapabilitySession = HarnessSession & {
		225	persistAppsEnabled: (appsEnabled: boolean) => Effect.Effect<void>;
		226	loadSessionMeta: () => Effect.Effect<SessionMeta \| null>;
		227	resolveSessionMeta: (token: unknown) => Effect.Effect<SessionMeta>;
		228	resolveAndStoreSessionMeta: (token: unknown) => Effect.Effect<SessionMeta>;
		229	};
		230
		231	const baseMeta: SessionMeta = {
		232	organizationId: "org-1",
		233	organizationName: "Org 1",
		234	userId: "user-1",
		235	resource: defaultMcpResource,
		236	};
		237
		238	const makeCapabilitySession = async (
		239	stored: SessionMeta = baseMeta,
		240	): Promise<{ session: CapabilitySession; storage: MemoryStorage }> => {
		241	const storage = new MemoryStorage();
		242	await storage.put("session-meta", stored);
		243	const session = Object.create(McpAgentSessionDOBase.prototype) as CapabilitySession;
		244	session.ctx = storage;
		245	session.getSessionId = () => "session-caps";
		246	return { session, storage };
		247	};
		248
		249	it("persists the negotiated capability so a later restore can read it back", async () => {
		250	const { session, storage } = await makeCapabilitySession();
		251
		252	await Effect.runPromise(session.persistAppsEnabled(true));
		253
		254	expect(await storage.get<SessionMeta>("session-meta")).toMatchObject({
		255	organizationId: "org-1",
		256	appsEnabled: true,
		257	});
		258	});
		259
		260	it("records a client that loses apps support just as durably", async () => {
		261	const { session, storage } = await makeCapabilitySession({ ...baseMeta, appsEnabled: true });
		262
		263	await Effect.runPromise(session.persistAppsEnabled(false));
		264
		265	expect(await storage.get<SessionMeta>("session-meta")).toMatchObject({ appsEnabled: false });
		266	});
		267
		268	// `init` runs again on every cold restore and rebuilds meta from the bearer
		269	// token, which carries no capabilities. If that overwrite won, restoring the
		270	// session would erase the very bit meant to survive it.
		271	it("carries the stored capability through the re-resolve on cold restore", async () => {
		272	const { session, storage } = await makeCapabilitySession({ ...baseMeta, appsEnabled: true });
		273	// What the token resolves to: no `appsEnabled` anywhere in sight.
		274	session.resolveSessionMeta = () => Effect.succeed(baseMeta);
		275
		276	const resolved = await Effect.runPromise(
		277	session.resolveAndStoreSessionMeta({ organizationId: "org-1", userId: "user-1" }),
		278	);
		279
		280	expect(resolved.appsEnabled).toBe(true);
		281	expect(await storage.get<SessionMeta>("session-meta")).toMatchObject({ appsEnabled: true });
		282	});
		283
		284	it("leaves a session with no negotiated capability untouched", async () => {
		285	const { session, storage } = await makeCapabilitySession();
		286	session.resolveSessionMeta = () => Effect.succeed(baseMeta);
		287
		288	const resolved = await Effect.runPromise(
		289	session.resolveAndStoreSessionMeta({ organizationId: "org-1", userId: "user-1" }),
		290	);
		291
		292	expect(resolved.appsEnabled).toBeUndefined();
		293	expect(await storage.get<SessionMeta>("session-meta")).not.toHaveProperty("appsEnabled");
		294	});
		295
		296	// Persistence is best-effort observation of a capability, never a reason to
		297	// fail the session that was merely trying to render something.
		298	it("stays silent when there is no stored meta to merge into", async () => {
		299	const storage = new MemoryStorage();
		300	const session = Object.create(McpAgentSessionDOBase.prototype) as CapabilitySession;
		301	session.ctx = storage;
		302	session.getSessionId = () => "session-caps";
		303
		304	await expect(Effect.runPromise(session.persistAppsEnabled(true))).resolves.toBeUndefined();
		305	expect(await storage.get<SessionMeta>("session-meta")).toBeUndefined();
		306	});
		307	});
		308
219	describe("McpAgentSessionDOBase transport restore", () => {	309	describe("McpAgentSessionDOBase transport restore", () => {
220	it("restores a same-session request after idle disposal leaves a stale server transport", async () => {	310	it("restores a same-session request after idle disposal leaves a stale server transport", async () => {
221	const session = await makeHarnessSession();	311	const session = await makeHarnessSession();

packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts+39−0

@@ -109,6 +109,17 @@ export interface SessionMeta { · 108 unmodified lines
109	* `buildMcpServer` scopes the tool catalog to it. */	109	* `buildMcpServer` scopes the tool catalog to it. */
110	readonly resource: McpResource;	110	readonly resource: McpResource;
111	readonly webOrigin?: string;	111	readonly webOrigin?: string;
		112	/**
		113	* Whether this session's client advertised MCP-Apps support at `initialize`.
		114	*
		115	* Capabilities are negotiated once, into the server instance's memory. When
		116	* the DO is evicted (deploy, idle) and a later request cold-restores it, the
		117	* rebuilt server never sees an `initialize` — so without persisting this, an
		118	* apps-capable client silently drops to artifact deep links mid-conversation.
		119	* Absent — including for sessions persisted before this field existed — means
		120	* unknown, which behaves as disabled until the next `initialize`.
		121	*/
		122	readonly appsEnabled?: boolean;
112	}	123	}
113		124
114	export interface BuiltMcpServer {	125	export interface BuiltMcpServer {
@@ -355,6 +366,28 @@ export abstract class McpAgentSessionDOBase< · 240 unmodified lines
355	await this.ctx.storage.put(SESSION_META_KEY, sessionMeta);	366	await this.ctx.storage.put(SESSION_META_KEY, sessionMeta);
356	}	367	}
357		368
		369	/**
		370	* Persist the MCP-Apps support negotiated at `initialize`, so a later cold
		371	* restore can rebuild the server with it. Subclasses hand this to
		372	* `createExecutorMcpServer` as `onAppsEnabledChange`.
		373	*
		374	* A no-op before meta exists: `initialize` always follows `init`, so there is
		375	* nothing to merge into and nothing worth failing the session over.
		376	*/
		377	protected persistAppsEnabled(appsEnabled: boolean): Effect.Effect<void> {
		378	const self = this;
		379	return Effect.gen(function* () {
		380	const stored = yield* self.loadSessionMeta();
		381	if (!stored \|\| stored.appsEnabled === appsEnabled) return;
		382	yield* Effect.promise(() => self.saveSessionMeta({ ...stored, appsEnabled }));
		383	}).pipe(
		384	Effect.withSpan("mcp.session.persist_apps_enabled", {
		385	attributes: { "mcp.artifact.apps_enabled": appsEnabled },
		386	}),
		387	Effect.ignoreCause({ log: false }),
		388	);
		389	}
		390
358	private async markActivity(now = Date.now()): Promise<void> {	391	private async markActivity(now = Date.now()): Promise<void> {
359	this.lastActivityMs = now;	392	this.lastActivityMs = now;
360	await Promise.all([	393	await Promise.all([
@@ -466,9 +499,15 @@ export abstract class McpAgentSessionDOBase< · 105 unmodified lines
466	const self = this;	499	const self = this;
467	return Effect.gen(function* () {	500	return Effect.gen(function* () {
468	const resolved = yield* self.resolveSessionMeta(token);	501	const resolved = yield* self.resolveSessionMeta(token);
		502	// `init` runs again on every cold restore, and `resolveSessionMeta`
		503	// rebuilds meta from the bearer token — which carries no negotiated
		504	// capabilities. Carry the stored value forward, or restoring the session
		505	// would erase the very bit that survives the restore.
		506	const stored = yield* self.loadSessionMeta();
469	const sessionMeta: SessionMeta = {	507	const sessionMeta: SessionMeta = {
470	...resolved,	508	...resolved,
471	...(token.webOrigin ? { webOrigin: token.webOrigin } : {}),	509	...(token.webOrigin ? { webOrigin: token.webOrigin } : {}),
		510	...(stored?.appsEnabled === undefined ? {} : { appsEnabled: stored.appsEnabled }),
472	};	511	};
473	yield* Effect.promise(() => self.saveSessionMeta(sessionMeta)).pipe(	512	yield* Effect.promise(() => self.saveSessionMeta(sessionMeta)).pipe(
474	Effect.withSpan("mcp.session.save_meta"),	513	Effect.withSpan("mcp.session.save_meta"),

packages/hosts/mcp/src/artifacts-tools.test.ts+142−1

@@ -1,4 +1,4 @@
1	import { describe, expect, it } from "@effect/vitest";	1	import { describe, expect, it, vi } from "@effect/vitest";
2	import { Data, Effect } from "effect";	2	import { Data, Effect } from "effect";
3	import { Client } from "@modelcontextprotocol/sdk/client/index.js";	3	import { Client } from "@modelcontextprotocol/sdk/client/index.js";
4	import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";	4	import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
@@ -236,6 +236,147 @@ describe("MCP host — artifact tool visibility", () => { · 231 unmodified lines
236	);	236	);
237	});	237	});
238		238
		239	// A cold restore (deploy, DO eviction) rebuilds the server mid-conversation
		240	// on the same session id, and the client — already past `initialize` — never
		241	// sends another one. Everything the rebuilt server knows about the client's
		242	// capabilities has to come from what the host persisted.
		243	it("keeps the app-only tools visible on a cold restore that replays no initialize", async () => {
		244	const store = makeArtifactStore();
		245	const mcpServer = await Effect.runPromise(
		246	createExecutorMcpServer({
		247	engine: makeStubEngine({}),
		248	artifacts: store.port,
		249	artifactsEnabled: true,
		250	loadAppShellHtml: () => Promise.resolve(SHELL_HTML),
		251	// What the host read back out of DO storage for this session.
		252	restoredAppsEnabled: true,
		253	}),
		254	);
		255	const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
		256	await mcpServer.connect(serverTransport);
		257	// Deliberately NOT an SDK `Client`: `client.connect()` always sends
		258	// `initialize`, which is the one thing a cold restore never replays. The
		259	// client is mid-conversation and just keeps issuing calls against the same
		260	// mcp-session-id, so this drives raw JSON-RPC over the transport.
		261	const responses = new Map<number, Record<string, unknown>>();
		262	clientTransport.onmessage = (message) => {
		263	const envelope = message as { id?: number; result?: Record<string, unknown> };
		264	if (typeof envelope.id === "number" && envelope.result) {
		265	responses.set(envelope.id, envelope.result);
		266	}
		267	};
		268	await clientTransport.start();
		269	const call = async (id: number, method: string, params: Record<string, unknown>) => {
		270	await clientTransport.send({ jsonrpc: "2.0", id, method, params });
		271	await vi.waitFor(() => expect(responses.has(id)).toBe(true));
		272	return responses.get(id) ?? {};
		273	};
		274
		275	const listed = (await call(1, "tools/list", {})) as {
		276	tools: Array<{ name: string }>;
		277	};
		278	const names = listed.tools.map((tool) => tool.name);
		279	expect(names).toContain("execute-action");
		280	expect(names).toContain("execute-action-resume");
		281
		282	// The render path is the symptom that took this to production: the same
		283	// session rendered inline before the restore and a deep link after.
		284	const created = (await call(2, "tools/call", {
		285	name: "create-artifact",
		286	arguments: { code: COUNTER_CODE, title: "Active users dashboard" },
		287	})) as { structuredContent?: Record<string, unknown> };
		288	expect(created.structuredContent).toEqual({ code: COUNTER_CODE, artifactId: "art_1" });
		289	expect(created.structuredContent).not.toHaveProperty("status", "fallback_url");
		290
		291	await clientTransport.close();
		292	await serverTransport.close();
		293	});
		294
		295	it("still falls back to a deep link when the restored session had no apps support", async () => {
		296	const store = makeArtifactStore();
		297	await withClient(
		298	makeStubEngine({}),
		299	NO_APPS_CAPS,
		300	async (client) => {
		301	const result = await client.callTool({
		302	name: "create-artifact",
		303	arguments: { code: COUNTER_CODE, title: "Active users dashboard" },
		304	});
		305	expect(structuredOf(result)).toEqual({
		306	status: "fallback_url",
		307	url: "https://executor.test/artifacts/art_1",
		308	artifactId: "art_1",
		309	});
		310	},
		311	{
		312	artifacts: store.port,
		313	restoredAppsEnabled: false,
		314	artifactUrl: artifactUrlFor("https://executor.test"),
		315	},
		316	);
		317	});
		318
		319	// The restored value seeds the session; it must never outrank the client in
		320	// front of you. A client that reconnects without apps support gets the
		321	// fallback, restored bit or not.
		322	it("lets a real initialize overrule the restored value in both directions", async () => {
		323	const store = makeArtifactStore();
		324	await withClient(
		325	makeStubEngine({}),
		326	NO_APPS_CAPS,
		327	async (client) => {
		328	const names = await toolNames(client);
		329	expect(names).not.toContain("execute-action");
		330	const result = await client.callTool({
		331	name: "create-artifact",
		332	arguments: { code: COUNTER_CODE, title: "Active users dashboard" },
		333	});
		334	expect(structuredOf(result)).toMatchObject({ status: "fallback_url" });
		335	},
		336	{
		337	artifacts: store.port,
		338	restoredAppsEnabled: true,
		339	artifactUrl: artifactUrlFor("https://executor.test"),
		340	},
		341	);
		342	});
		343
		344	it("reports the negotiated value to the host so the next restore can seed itself", async () => {
		345	const store = makeArtifactStore();
		346	const seen: boolean[] = [];
		347	await withClient(
		348	makeStubEngine({}),
		349	APPS_CAPS,
		350	async (client) => {
		351	await toolNames(client);
		352	},
		353	{
		354	artifacts: store.port,
		355	restoredAppsEnabled: false,
		356	onAppsEnabledChange: (appsEnabled) => Effect.sync(() => void seen.push(appsEnabled)),
		357	},
		358	);
		359	expect(seen).toEqual([true]);
		360	});
		361
		362	it("does not re-persist when initialize confirms the restored value", async () => {
		363	const store = makeArtifactStore();
		364	const seen: boolean[] = [];
		365	await withClient(
		366	makeStubEngine({}),
		367	APPS_CAPS,
		368	async (client) => {
		369	await toolNames(client);
		370	},
		371	{
		372	artifacts: store.port,
		373	restoredAppsEnabled: true,
		374	onAppsEnabledChange: (appsEnabled) => Effect.sync(() => void seen.push(appsEnabled)),
		375	},
		376	);
		377	expect(seen).toEqual([]);
		378	});
		379
239	it("registers no ui tools at all when no shell loader is configured", async () => {	380	it("registers no ui tools at all when no shell loader is configured", async () => {
240	const store = makeArtifactStore();	381	const store = makeArtifactStore();
241	const mcpServer = await Effect.runPromise(	382	const mcpServer = await Effect.runPromise(

packages/hosts/mcp/src/tool-server.ts+46−3

@@ -211,6 +211,25 @@ type SharedMcpServerConfig = { · 210 unmodified lines
211	* result. Hosts recording product analytics supply it; core stays agnostic.	211	* result. Hosts recording product analytics supply it; core stays agnostic.
212	*/	212	*/
213	readonly onArtifactUsage?: (action: "created" \| "viewed" \| "updated") => Effect.Effect<void>;	213	readonly onArtifactUsage?: (action: "created" \| "viewed" \| "updated") => Effect.Effect<void>;
		214	/**
		215	* Whether the client this session belongs to can render MCP Apps, as
		216	* negotiated at a previous `initialize`.
		217	*
		218	* Capabilities normally arrive from the client at `initialize` and live only
		219	* in the server instance. A session whose host evicted and cold-restored it
		220	* (deploy, idle) is rebuilt mid-conversation with no `initialize` to replay,
		221	* so without this the rebuilt server assumes no apps support and silently
		222	* downgrades every artifact to a deep link. Hosts that persist the
		223	* negotiated value pass it back here; the next `initialize`, if one comes,
		224	* overwrites it.
		225	*/
		226	readonly restoredAppsEnabled?: boolean;
		227	/**
		228	* Called when `initialize` negotiates the client's MCP-Apps support, so the
		229	* host can persist it for {@link restoredAppsEnabled} on a later cold
		230	* restore. Best-effort: failures are swallowed and never affect the session.
		231	*/
		232	readonly onAppsEnabledChange?: (appsEnabled: boolean) => Effect.Effect<void>;
214	};	233	};
215		234
216	/**	235	/**
@@ -1505,8 +1524,11 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 1288 unmodified lines
1505	const artifacts = config.artifacts;	1524	const artifacts = config.artifacts;
1506		1525
1507	// Set from the client's advertised capabilities at `initialize`. Read by	1526	// Set from the client's advertised capabilities at `initialize`. Read by
1508	// the render handlers to choose inline widget vs. deep link.	1527	// the render handlers to choose inline widget vs. deep link. Seeded from
1509	let appsEnabled = false;	1528	// the host's persisted value so a cold-restored session — rebuilt with no
		1529	// `initialize` to replay — keeps rendering inline for a client that had
		1530	// already negotiated apps support.
		1531	let appsEnabled = config.restoredAppsEnabled ?? false;
1510	let executeActionTool: { enable: () => void; disable: () => void } \| undefined;	1532	let executeActionTool: { enable: () => void; disable: () => void } \| undefined;
1511	let executeActionResumeTool: { enable: () => void; disable: () => void } \| undefined;	1533	let executeActionResumeTool: { enable: () => void; disable: () => void } \| undefined;
1512		1534
@@ -1921,7 +1943,28 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 408 unmodified lines
1921	\| (ClientCapabilities & { extensions?: Record<string, unknown> })	1943	\| (ClientCapabilities & { extensions?: Record<string, unknown> })
1922	\| null,	1944	\| null,
1923	);	1945	);
1924	appsEnabled = Boolean(uiCapability?.mimeTypes?.includes(RESOURCE_MIME_TYPE));	1946	// Absent capabilities (the SDK returns `undefined`) mean `initialize`
		1947	// hasn't happened on THIS server instance — the construction-time call
		1948	// below, or a cold restore that resumed mid-conversation. Neither is
		1949	// evidence the client lost apps support, so the restored value stands
		1950	// until a real `initialize` replaces it. Reading `false` off an absent
		1951	// value here is exactly what made a cold-restored session fall back to
		1952	// deep links.
		1953	const negotiated = clientCapabilities
		1954	? Boolean(uiCapability?.mimeTypes?.includes(RESOURCE_MIME_TYPE))
		1955	: appsEnabled;
		1956	const changed = negotiated !== appsEnabled;
		1957	appsEnabled = negotiated;
		1958
		1959	// Persist only a real negotiation that moved the value, so the next cold
		1960	// restore seeds itself. Best-effort: the session must not fail on it.
		1961	const onAppsEnabledChange = config.onAppsEnabledChange;
		1962	if (clientCapabilities && changed && onAppsEnabledChange) {
		1963	// oxlint-disable-next-line executor/no-effect-escape-hatch -- boundary: `oninitialized` is a sync SDK hook; persistence is fire-and-forget and its failure must not fail the session
		1964	void Effect.runPromiseWith(context)(
		1965	onAppsEnabledChange(negotiated).pipe(Effect.ignoreCause({ log: false })),
		1966	);
		1967	}
1925		1968
1926	// `execute-action` is only callable from inside a rendered app, so a	1969	// `execute-action` is only callable from inside a rendered app, so a
1927	// client that can't render one should never see it. `create-artifact`,	1970	// client that can't render one should never see it. `create-artifact`,