Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions .changeset/mcp-progress-projector.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
---
"@agent-bundle/runtime": minor
"agent-bundle": patch
---

Add the standards-compatible MCP progress/final projector and warm-runtime
fail-closed host. Generated tool calls emit `notifications/progress` only when
the caller supplied a token, return one `CallToolResult`, and refuse silent
rich-content drops, epoch mismatch, and a missing or restarted runtime.
7 changes: 6 additions & 1 deletion packages/agent-bundle/src/build/entries.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import { emitPlanEntries, resolveArtifactDestination } from './emit.ts';
import { scanEntryExports } from './entry-exports.ts';
import {
generatedExecutableEntrySource,
generatedRouteArtifactEpoch,
generatedRouteFlightWorkerSource,
generatedRouteMcpEntrySource,
generatedStdioMcpEntrySource,
Expand Down Expand Up @@ -215,7 +216,11 @@ export const compileMcpEntries = async (
const server = servers.find((candidate) => candidate.id === entry.id);
return server?.generatedRoutes === undefined
? undefined
: generatedRouteFlightWorkerSource({ routes: server.generatedRoutes, serverName: server.name });
: generatedRouteFlightWorkerSource({
artifactEpoch: generatedRouteArtifactEpoch(options.plugin),
routes: server.generatedRoutes,
serverName: server.name,
});
});
// Factory-exporting entries (default export) are wrapped in the framework
// stdio lifecycle shell; self-connecting entries keep today's behavior byte
Expand Down
104 changes: 67 additions & 37 deletions packages/agent-bundle/src/build/entry-shell.ts
Original file line number Diff line number Diff line change
Expand Up @@ -81,10 +81,16 @@ export interface GeneratedRouteMcpEntryOptions {
}

export interface GeneratedRouteFlightWorkerOptions {
readonly artifactEpoch: string;
readonly routes: readonly CompiledAgentRoute[];
readonly serverName: string;
}

export const generatedRouteArtifactEpoch = (plugin: {
readonly name: string;
readonly version: string;
}): string => `${plugin.name}@${plugin.version}`;

const routeProtocolName = (route: CompiledAgentRoute): string =>
route.id.slice(route.id.lastIndexOf('/') + 1);

Expand Down Expand Up @@ -118,20 +124,30 @@ export const generatedRouteFlightWorkerSource = (options: GeneratedRouteFlightWo
'// Generated routes contain only intrinsic Agent protocol elements, so no client references exist.',
'globalThis.__rspack_rsc_manifest__ ??= Object.freeze({ clientManifest: Object.freeze({}) });',
"if (parentPort === null) throw new Error('Generated Flight worker requires a parent port.');",
'process.stdout.write = process.stderr.write.bind(process.stderr);',
`const ARTIFACT_EPOCH = ${JSON.stringify(options.artifactEpoch)};`,
'const processLifetime = { hits: 0, instanceId: crypto.randomUUID(), pid: process.pid };',
'const routes = Object.freeze({',
...routeRecords(routes),
'});',
'const requests = new Map();',
'',
'const render = async (message) => {',
' if (message.artifactEpoch !== undefined && message.artifactEpoch !== ARTIFACT_EPOCH) {',
" parentPort.postMessage({ code: 'artifact-epoch-mismatch', id: message.id, message: `Runtime artifact epoch ${JSON.stringify(ARTIFACT_EPOCH)} does not match request epoch ${JSON.stringify(message.artifactEpoch)}`, receivedEpoch: message.artifactEpoch, type: 'error' });",
' return;',
' }',
' const route = routes[message.invocation.props.operationId];',
" if (route === undefined || typeof route.module.default !== 'function') throw new TypeError('Generated MCP route must default-export an async Server Component.');",
' const controller = new AbortController();',
' requests.set(message.id, controller);',
' processLifetime.hits += 1;',
' try {',
' const bytes = await runAgentRequest({',
' ...(message.actor === undefined ? {} : { actor: message.actor }),',
' invocation: { kind: \'tool\', operationId: route.id, surface: route.name },',
' invocation: { artifactEpoch: ARTIFACT_EPOCH, kind: \'tool\', operationId: route.id, surface: route.name },',
' progress: { report: async (update) => { parentPort.postMessage({ id: message.id, type: \'progress\', update }); } },',
' providers: { processLifetime: { hits: processLifetime.hits, instanceId: processLifetime.instanceId, pid: processLifetime.pid } },',
' ...(message.session === undefined ? {} : { session: message.session }),',
' signal: controller.signal,',
' }, async () => {',
Expand Down Expand Up @@ -166,7 +182,10 @@ const routeRegistrations = (routes: readonly CompiledAgentRoute[]): readonly str
` ...${stableJson(config)},`,
` inputSchema: ${access}.module.inputSchema,`,
` outputSchema: ${access}.module.resultSchema,`,
` }, async (input, context) => projectToolResult(await renderRoute(dispatcher, ${access}, input, context)));`,
` }, async (input, context) => {`,
` const rendered = await renderRoute(dispatcher, ${access}, input, context);`,
' return attachMcpStructuredContent(rendered.toolResult, rendered.result);',
' });',
].join('\n'));
break;
}
Expand Down Expand Up @@ -210,90 +229,101 @@ const routeRegistrations = (routes: readonly CompiledAgentRoute[]): readonly str
};

/**
* The generated MCP entry owns the final-only dispatcher and one warm Flight
* The generated MCP entry owns the stream projector and one warm Flight
* worker. The worker is split only to satisfy React's react-server condition;
* it is reused for every request until the MCP server closes.
*/
export const generatedRouteMcpEntrySource = (options: GeneratedRouteMcpEntryOptions): string => {
const routes = executableMcpRoutes(options.routes);
const artifactEpoch = generatedRouteArtifactEpoch(options.plugin);
return [
"import { Worker } from 'node:worker_threads';",
"import { McpServer } from '@modelcontextprotocol/server';",
"import { agent, available, createAgentRenderDispatcher, runAgentRequest } from '@agent-bundle/runtime';",
"import { AgentRuntimeError, agent, attachMcpStructuredContent, available, createAgentRenderDispatcher, createWarmFlightHost, projectMcpRenderStream, runAgentRequest } from '@agent-bundle/runtime';",
"import mcpApps from 'agent-bundle/mcp-apps';",
...routeImports(routes),
'',
`const ARTIFACT_EPOCH = ${JSON.stringify(artifactEpoch)};`,
'const routes = Object.freeze({',
...routeRecords(routes),
'});',
'',
'const workerError = (message) => {',
" if (message.code === 'artifact-epoch-mismatch') return new AgentRuntimeError('artifact-epoch-mismatch', message.message, { expectedEpoch: ARTIFACT_EPOCH, receivedEpoch: message.receivedEpoch });",
" if (message.code === 'runtime-unavailable') return new AgentRuntimeError('runtime-unavailable', message.message);",
" if (message.code === 'runtime-restarted') return new AgentRuntimeError('runtime-restarted', message.message);",
' return new Error(message.message);',
'};',
'',
'const createWorkerHost = () => {',
` const worker = new Worker(new URL(${JSON.stringify(`./${options.workerFile}`)}, import.meta.url));`,
` const worker = new Worker(new URL(${JSON.stringify(`./${options.workerFile}`)}, import.meta.url), { stderr: true, stdout: true });`,
' worker.stdout?.on(\'data\', (chunk) => process.stderr.write(chunk));',
' worker.stderr?.on(\'data\', (chunk) => process.stderr.write(chunk));',
' const pending = new Map();',
' let sequence = 0;',
' let exited = false;',
' const failPending = (error) => { for (const request of pending.values()) request.reject(error); pending.clear(); };',
' worker.on(\'error\', failPending);',
' worker.on(\'exit\', (code) => { if (code !== 0) failPending(new Error(`Generated Flight worker exited with code ${String(code)}.`)); });',
' worker.on(\'error\', (error) => { exited = true; failPending(error); });',
' worker.on(\'exit\', (code) => {',
' exited = true;',
" failPending(new AgentRuntimeError(code === 0 ? 'runtime-unavailable' : 'runtime-restarted', code === 0 ? 'The MCP render runtime is unavailable' : `The MCP render runtime restarted; worker exited with code ${String(code)}.`));",
' });',
' worker.on(\'message\', (message) => {',
' const request = pending.get(message.id);',
' if (request === undefined) return;',
" if (message.type === 'progress') { void request.progress?.report(message.update); return; }",

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Await worker progress forwarding

When a generated tool reports more than the session's buffered progress limit (10,000 by default) before returning, request.progress.report() rejects with event-count-exceeded, but this discarded promise becomes an unhandled rejection in the MCP server process. Under the supported Node versions, that can terminate the entire stdio server instead of failing only the offending tool call; forwarding needs to be awaited/serialized (with worker acknowledgement for backpressure) or its rejection must be routed into the pending request.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in #210 (merged as 1c36813): the generated shell no longer voids progress forwarding — rejections (and sync throws) from entry.progress.report route into the request failure path via entry.fail, so a progress-cap rejection fails the request instead of becoming an unhandled worker rejection. Regression test in entry-shell.test.ts runs the generated factory with a rejecting reporter and asserts the request fails with no unhandled rejection.

' pending.delete(message.id);',
' request.signal.removeEventListener(\'abort\', request.abort);',
" if (message.type === 'error') { request.reject(new Error(message.message)); return; }",
" if (message.type === 'error') { request.reject(workerError(message)); return; }",
' request.resolve(new ReadableStream({ start(controller) { controller.enqueue(message.bytes); controller.close(); } }));',
' });',
' return Object.freeze({',
' const host = Object.freeze({',
' close: async () => { await worker.terminate(); },',
' execute: async ({ invocation, signal }) => {',
' execute: async ({ artifactEpoch, invocation, progress, signal }) => {',
" if (exited) throw new AgentRuntimeError('runtime-unavailable', 'The MCP render runtime is unavailable');",
' const context = await agent();',
' const id = ++sequence;',
' return new Promise((resolve, reject) => {',
" const abort = () => { worker.postMessage({ id, type: 'cancel' }); pending.delete(id); reject(new DOMException('Agent render was aborted', 'AbortError')); };",
' pending.set(id, { abort, reject, resolve, signal });',
' pending.set(id, { abort, progress, reject, resolve, signal });',
" signal.addEventListener('abort', abort, { once: true });",
' if (signal.aborted) { abort(); return; }',
" worker.postMessage({ actor: context.actor, id, invocation, session: context.session, type: 'render' });",
' worker.postMessage({ actor: context.actor, artifactEpoch: artifactEpoch ?? ARTIFACT_EPOCH, id, invocation, session: context.session, type: \'render\' });',
' });',
' },',
' });',
' return createWarmFlightHost({ artifactEpoch: ARTIFACT_EPOCH, close: host.close, host });',
'};',
'',
'const requestIdentity = (context) => ({',
' ...(context.http?.authInfo?.clientId === undefined ? {} : { actor: available({ id: context.http.authInfo.clientId }, \'native\') }),',
" ...(typeof context.sessionId === 'string' && context.sessionId.trim() !== '' ? { session: available({ sessionId: context.sessionId }, 'native') } : {}),",
'});',
'',
'const mcpProjectorOptions = (context) => {',
' const progressToken = context.mcpReq._meta?.progressToken;',
' return {',
' signal: context.mcpReq.signal,',
' ...(progressToken === undefined ? {} : {',
' progressToken,',
" sendProgress: (params) => context.mcpReq.notify({ method: 'notifications/progress', params }),",
' }),',
' };',
'};',
'',
'const renderRoute = async (dispatcher, route, input, context) => runAgentRequest({',
' ...requestIdentity(context),',
" invocation: { kind: 'tool', operationId: route.id, surface: route.name },",
" invocation: { artifactEpoch: ARTIFACT_EPOCH, kind: 'tool', operationId: route.id, surface: route.name },",
' signal: context.mcpReq.signal,',
'}, async () => {',
' const document = await dispatcher.dispatch({ invocation: { kind: \'tool\', props: { input, operationId: route.id } }, signal: context.mcpReq.signal });',
' return { document, result: route.module.resultSchema.parse(document.value) };',
' const projected = await projectMcpRenderStream(dispatcher.stream({',
' artifactEpoch: ARTIFACT_EPOCH,',
" invocation: { kind: 'tool', props: { input, operationId: route.id } },",
' signal: context.mcpReq.signal,',
' }), mcpProjectorOptions(context));',
' return { document: projected.document, result: route.module.resultSchema.parse(projected.document.value), toolResult: projected.result };',
'});',
'',
'const appendNode = (node, content) => {',
' switch (node.kind) {',
" case 'result': for (const child of node.children) appendNode(child, content); break;",
" case 'context':",
" case 'markdown':",
" case 'text': content.push({ text: node.text, type: 'text' }); break;",
" case 'json': content.push({ text: JSON.stringify(node.value), type: 'text' }); break;",
" case 'progress': content.push({ text: node.message ?? `Progress: ${node.completed}${node.total === undefined ? '' : `/${node.total}`}`, type: 'text' }); break;",
" case 'image': content.push({ data: node.data, mimeType: node.mimeType, type: 'image' }); break;",
" case 'audio': content.push({ data: node.data, mimeType: node.mimeType, type: 'audio' }); break;",
" case 'resource': content.push({ ...(node.mimeType === undefined ? {} : { mimeType: node.mimeType }), name: node.name, type: 'resource_link', uri: node.uri }); break;",
" case 'error': content.push({ text: `[${node.code}] ${node.message}`, type: 'text' }); break;",
" default: throw new TypeError(`Unsupported Agent Document node: ${String(node.kind)}`);",
' }',
'};',
'',
'const projectToolResult = ({ document, result }) => {',
' const content = [];',
' appendNode(document.root, content);',
' return { content, ...(document.status === \'represented-error\' ? { isError: true } : {}), ...(result !== null && typeof result === \'object\' && !Array.isArray(result) ? { structuredContent: result } : {}) };',
'};',
'',
'const createGeneratedRouteServer = () => {',
` const server = new McpServer(${stableJson(options.plugin)});`,
' const workerHost = createWorkerHost();',
Expand Down
7 changes: 6 additions & 1 deletion packages/agent-bundle/src/build/inspect-bundler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ import type { AgentBundleToolsConfig, NormalizedPlugin } from '../core/types.ts'
import { scanEntryExports } from './entry-exports.ts';
import {
generatedExecutableEntrySource,
generatedRouteArtifactEpoch,
generatedRouteFlightWorkerSource,
generatedRouteMcpEntrySource,
generatedStdioMcpEntrySource,
Expand Down Expand Up @@ -211,7 +212,11 @@ const mcpEntryEntries = async (
rscManifest: true,
source: entry.source,
sourceInputs: [],
virtualSource: generatedRouteFlightWorkerSource({ routes: generatedRoutes, serverName }),
virtualSource: generatedRouteFlightWorkerSource({
artifactEpoch: generatedRouteArtifactEpoch({ name: model.metadata.name, version: model.metadata.version }),
routes: generatedRoutes,
serverName,
}),
},
kind: 'mcp-entry',
name: `${serverName}:flight`,
Expand Down
11 changes: 11 additions & 0 deletions packages/agent-bundle/tests/entry-shell.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -127,12 +127,20 @@ it('generates one final-only Flight MCP factory from filesystem routes', () => {
expect(source).toContain("from '@modelcontextprotocol/server'");
expect(source).toContain("from 'agent-bundle/mcp-apps'");
expect(source).toContain('createAgentRenderDispatcher');
expect(source).toContain('createWarmFlightHost');
expect(source).toContain('{ stderr: true, stdout: true }');
expect(source).toContain('projectMcpRenderStream');
expect(source).toContain('attachMcpStructuredContent');
expect(source).toContain('runAgentRequest');
expect(source).toContain('notifications/progress');
expect(source).toContain('ARTIFACT_EPOCH');
expect(source).toContain('route-fixture@1.2.3');
expect(source).toContain('server.registerTool("inspect"');
expect(source).toContain('server.registerResource("catalog", "catalog://books"');
expect(source).toContain('server.registerPrompt("curate"');
expect(source).toContain('export default createGeneratedRouteServer');
expect(source).not.toContain('lowerMcpResult');
expect(source).not.toContain('projectToolResult');
});


Expand All @@ -143,6 +151,7 @@ it('generates the warm react-server Flight worker separately from the MCP dispat
expect(typeof generate).toBe('function');
if (generate === undefined) return;
const source = generate({
artifactEpoch: 'route-fixture@1.2.3',
routes: [{
config: {},
id: 'tool:curator/inspect',
Expand All @@ -154,5 +163,7 @@ it('generates the warm react-server Flight worker separately from the MCP dispat
expect(source).toContain("from '@agent-bundle/runtime/flight/server'");
expect(source).toContain("from 'node:worker_threads'");
expect(source).toContain('runAgentRequest');
expect(source).toContain('processLifetime');
expect(source).toContain('route-fixture@1.2.3');
expect(source).toContain('/project/src/mcp/curator/tools/inspect.tsx');
});
Loading
Loading