History

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
  1. apps/cloud/src/mcp/session-durable-object.ts +4 −0
  2. apps/host-cloudflare/src/mcp/session-durable-object.ts +4 −0
  3. packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts +90 −0
  4. packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts +39 −0
  5. packages/hosts/mcp/src/artifacts-tools.test.ts +142 −1
  6. packages/hosts/mcp/src/tool-server.ts +46 −3
apps/cloud/src/mcp/session-durable-object.ts+4−0
@@ -280,6 +280,10 @@ export class McpSessionDOSqlite extends McpAgentSessionDOBase<Env, CloudSessionD · 279 unmodified lines
280280 // flag existed carry no value; absent now means off, same as a fresh
281281 // connection that did not ask for `?artifacts=true`.
282282 artifactsEnabled: sessionMeta.artifactsEnabled ?? false,
283+ // Cold restores rebuild this server with no `initialize` to replay, so
284+ // the negotiated apps support comes back from storage instead.
285+ restoredAppsEnabled: sessionMeta.appsEnabled ?? false,
286+ onAppsEnabledChange: (appsEnabled) => self.persistAppsEnabled(appsEnabled),
283287 loadAppShellHtml,
284288 smokeRenderArtifact,
285289 artifactUrl: artifactUrlFor(
apps/host-cloudflare/src/mcp/session-durable-object.ts+4−0
@@ -149,6 +149,10 @@ export class McpSessionDO extends McpAgentSessionDOBase<CloudflareEnv, CfSession · 148 unmodified lines
149149 // flag existed carry no value; absent now means off, same as a fresh
150150 // connection that did not ask for `?artifacts=true`.
151151 artifactsEnabled: sessionMeta.artifactsEnabled ?? false,
152+ // Cold restores rebuild this server with no `initialize` to replay, so
153+ // the negotiated apps support comes back from storage instead.
154+ restoredAppsEnabled: sessionMeta.appsEnabled ?? false,
155+ onAppsEnabledChange: (appsEnabled) => self.persistAppsEnabled(appsEnabled),
152156 loadAppShellHtml: self.loadAppShellHtml,
153157 smokeRenderArtifact,
154158 ...(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
216216 return session;
217217 };
218218
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+
219309 describe("McpAgentSessionDOBase transport restore", () => {
220310 it("restores a same-session request after idle disposal leaves a stale server transport", async () => {
221311 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
109109 * `buildMcpServer` scopes the tool catalog to it. */
110110 readonly resource: McpResource;
111111 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;
112123 }
113124
114125 export interface BuiltMcpServer {
@@ -355,6 +366,28 @@ export abstract class McpAgentSessionDOBase< · 240 unmodified lines
355366 await this.ctx.storage.put(SESSION_META_KEY, sessionMeta);
356367 }
357368
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+
358391 private async markActivity(now = Date.now()): Promise<void> {
359392 this.lastActivityMs = now;
360393 await Promise.all([
@@ -466,9 +499,15 @@ export abstract class McpAgentSessionDOBase< · 105 unmodified lines
466499 const self = this;
467500 return Effect.gen(function* () {
468501 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();
469507 const sessionMeta: SessionMeta = {
470508 ...resolved,
471509 ...(token.webOrigin ? { webOrigin: token.webOrigin } : {}),
510+ ...(stored?.appsEnabled === undefined ? {} : { appsEnabled: stored.appsEnabled }),
472511 };
473512 yield* Effect.promise(() => self.saveSessionMeta(sessionMeta)).pipe(
474513 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";
22 import { Data, Effect } from "effect";
33 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
44 import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
@@ -236,6 +236,147 @@ describe("MCP host — artifact tool visibility", () => { · 231 unmodified lines
236236 );
237237 });
238238
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+
239380 it("registers no ui tools at all when no shell loader is configured", async () => {
240381 const store = makeArtifactStore();
241382 const mcpServer = await Effect.runPromise(
packages/hosts/mcp/src/tool-server.ts+46−3
@@ -211,6 +211,25 @@ type SharedMcpServerConfig = { · 210 unmodified lines
211211 * result. Hosts recording product analytics supply it; core stays agnostic.
212212 */
213213 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>;
214233 };
215234
216235 /**
@@ -1505,8 +1524,11 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 1288 unmodified lines
15051524 const artifacts = config.artifacts;
15061525
15071526 // Set from the client's advertised capabilities at `initialize`. Read by
1508− // the render handlers to choose inline widget vs. deep link.
1509− let appsEnabled = false;
1527+ // the render handlers to choose inline widget vs. deep link. Seeded from
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;
15101532 let executeActionTool: { enable: () => void; disable: () => void } | undefined;
15111533 let executeActionResumeTool: { enable: () => void; disable: () => void } | undefined;
15121534
@@ -1921,7 +1943,28 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 408 unmodified lines
19211943 | (ClientCapabilities & { extensions?: Record<string, unknown> })
19221944 | null,
19231945 );
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+ }
19251968
19261969 // `execute-action` is only callable from inside a rendered app, so a
19271970 // client that can't render one should never see it. `create-artifact`,