fix(lineage): correlate by the raw wire arguments and keep undigested windows in contention (review)

Author: ScriptedAlchemy

.changeset/424-cursor-mcp-correlation-arguments.md

@@ -3,4 +3,4 @@
 "agent-bundle": patch
 ---
 
-Resolve `request.lineage` for concurrent Cursor MCP calls by their arguments. Cursor's `tools/call` `_meta` names no conversation, so a generated MCP server correlated a call only through the open `MCP:<tool>` pre-tool hook and reported `id-not-resolvable` whenever several conversations had the same tool open. The pre-tool hook's `tool_input` is the call's arguments verbatim, so the lineage registry now records their digest on each open window (`inputDigest`) and the generated server passes the call's arguments to `resolveToolCall`; a concurrent call with different arguments resolves (`resolution: inferred`, provenance `derived`), identical arguments still refuse, and a single open conversation is unaffected. (#483)
+Resolve `request.lineage` for concurrent Cursor MCP calls by their arguments. Cursor's `tools/call` `_meta` names no conversation, so a generated MCP server correlated a call only through the open `MCP:<tool>` pre-tool hook and reported `id-not-resolvable` whenever several conversations had the same tool open. The pre-tool hook's `tool_input` is the call's arguments verbatim, so the lineage registry now records their digest on each open window (`inputDigest`) and the generated server passes the call's raw wire arguments (captured before schema parsing, so input defaults never make two calls look alike) to `resolveToolCall`; a concurrent call with different arguments resolves (`resolution: inferred`, provenance `derived`), identical arguments still refuse, a window recorded without a digest stays in contention, and a single open conversation is unaffected. (#483)

packages/agent-bundle/src/mcp-server-runtime.ts

@@ -16,7 +16,7 @@
  */
 import { Worker } from 'node:worker_threads';
 
-import { McpServer, ProtocolError, ProtocolErrorCode } from '@modelcontextprotocol/server';
+import { McpServer, ProtocolError, ProtocolErrorCode, isJSONRPCRequest, type Transport } from '@modelcontextprotocol/server';
 import {
   AgentRuntimeError,
   agent,
@@ -84,6 +84,8 @@ export interface GeneratedRouteRequestContext {
   readonly mcpReq: {
     /** Request `_meta`: the progress token plus host-specific correlation keys (`claudecode/toolUseId`, `x-codex-turn-metadata`). */
     readonly _meta?: { readonly progressToken?: McpProgressToken } & Readonly<Record<string, unknown>>;
+    /** The JSON-RPC request id, the key the raw `tools/call` arguments were captured under. */
+    readonly id?: number | string;
     readonly notify?: (notification: {
       readonly method: 'notifications/progress';
       readonly params: McpProgressNotificationParams;
@@ -109,25 +111,79 @@ interface GeneratedRouteIdentity {
   readonly workspace: Observed<AgentWorkspaceIdentity>;
 }
 
+/** A captured `tools/call` `params.arguments` value; `value` is `undefined` when the call carried none. */
+export interface RawToolArguments {
+  readonly value: unknown;
+}
+
+/** Raw `tools/call` arguments by request, consumed once by the tool callback that serves the request. */
+export interface RawToolArgumentsCapture {
+  take(requestId: number | string | undefined): RawToolArguments | undefined;
+}
+
+const requestKey = (requestId: number | string): string => `${typeof requestId}:${String(requestId)}`;
+/** Calls that never reach a registered tool (unknown tool, rejected params) are forgotten past this many. */
+const RAW_ARGUMENTS_RETENTION = 1024;
+
+/**
+ * Captures every `tools/call`'s arguments off the wire, before the SDK parses
+ * them against the route's input schema. Lineage correlates a Cursor call with
+ * its pre-tool hook by those raw arguments (the hook's `tool_input`); schema
+ * defaults would make two different calls look alike, so the parsed input is
+ * never what is compared. The capture wraps the transport's `onmessage` the
+ * moment the server connects and keeps one entry per request id until the
+ * tool callback takes it.
+ */
+const captureRawToolArguments = (server: McpServer): RawToolArgumentsCapture => {
+  const captured = new Map<string, RawToolArguments>();
+  const connect = server.connect.bind(server);
+  server.connect = async (transport: Transport): Promise<void> => {
+    await connect(transport);
+    const inner = transport.onmessage;
+    transport.onmessage = (message, extra) => {
+      if (isJSONRPCRequest(message) && message.method === 'tools/call') {
+        const params = message.params;
+        const value = params !== undefined && typeof params === 'object' && params !== null && !Array.isArray(params)
+          ? (params as { readonly arguments?: unknown }).arguments
+          : undefined;
+        captured.set(requestKey(message.id), Object.freeze({ value }));
+        if (captured.size > RAW_ARGUMENTS_RETENTION) captured.delete(captured.keys().next().value!);
+      }
+      inner?.(message, extra);
+    };
+  };
+  return Object.freeze({
+    take(requestId: number | string | undefined): RawToolArguments | undefined {
+      if (requestId === undefined) return undefined;
+      const key = requestKey(requestId);
+      const raw = captured.get(key);
+      captured.delete(key);
+      return raw;
+    },
+  });
+};
+
 /**
  * Lineage for one MCP tool call: Codex names it in `_meta`, Claude names the
  * pre-tool hook's `tool_use_id` in `_meta`, Cursor names nothing — so the
  * registry falls back to the open `MCP:<tool>` pre-tool hook, told apart from
- * a concurrent call in another conversation by the arguments the hook
- * recorded. Without a registry (a project with no event routes, or the
- * in-memory proof level) the axis is honestly absent.
+ * a concurrent call in another conversation by the raw arguments the hook
+ * recorded. A call whose raw arguments were not captured (a transport the
+ * server did not connect itself) is correlated by tool name alone, never by
+ * its schema-parsed input. Without a registry (a project with no event
+ * routes, or the in-memory proof level) the axis is honestly absent.
  */
 const toolCallLineage = async (
   registry: AgentLineageRegistry | undefined,
   context: GeneratedRouteRequestContext,
   toolName: string,
-  input: unknown,
+  rawArguments: RawToolArguments | undefined,
   clientName: string | undefined,
   fallbackHost: LineageHost | undefined,
 ): Promise<Observed<AgentLineage>> => {
   if (registry === undefined) return unavailable<AgentLineage>('not-provided');
   return registry.resolveToolCall({
-    arguments: input,
+    ...(rawArguments === undefined ? {} : { arguments: rawArguments.value }),
     host: lineageHostFromClient(clientName) ?? fallbackHost,
     meta: context.mcpReq._meta,
     toolName,
@@ -291,6 +347,8 @@ export interface RegisterGeneratedRoutesOptions {
   readonly lineage?: AgentLineageRegistry;
   /** The artifact's host, used when the negotiated client name maps to none. */
   readonly lineageHost?: LineageHost;
+  /** Raw `tools/call` arguments captured off the wire, for lineage correlation. */
+  readonly rawArguments?: RawToolArgumentsCapture;
 }
 
 /** Registers the compiled MCP routes on a server, keyed by route kind. */
@@ -311,13 +369,14 @@ export const registerGeneratedRoutes = (
           ...(outputSchema === undefined ? {} : { outputSchema }),
         } as never, (async (input: unknown, context: GeneratedRouteRequestContext) => settled(async () => {
           const clientName = server.server.getClientVersion()?.name;
+          const rawArguments = options.rawArguments?.take(context.mcpReq.id);
           const rendered = await renderGeneratedRoute(
             dispatcher,
             artifactEpoch,
             route,
             input,
             context,
-            { clientName, lineage: await toolCallLineage(options.lineage, context, route.name, input, clientName, options.lineageHost) },
+            { clientName, lineage: await toolCallLineage(options.lineage, context, route.name, rawArguments, clientName, options.lineageHost) },
           );
           return attachMcpStructuredContent(rendered.toolResult, rendered.result);
         }, options.afterRender)) as never);
@@ -840,7 +899,7 @@ export const createGeneratedRouteMcpServer = async (
     : await startEventRuntime(options.events, dispatcher, options.host, afterRender, options.lineage);
   registerGeneratedRoutes(server, options.routes, dispatcher, options.artifactEpoch, {
     ...(afterRender === undefined ? {} : { afterRender }),
-    ...(options.lineage === undefined ? {} : { lineage: options.lineage }),
+    ...(options.lineage === undefined ? {} : { lineage: options.lineage, rawArguments: captureRawToolArguments(server) }),
     ...(options.events === undefined || lineageHostFor(options.events.target) === undefined
       ? {}
       : { lineageHost: lineageHostFor(options.events.target) }),

packages/agent-bundle/tests/mcp-server-runtime.test.ts

@@ -1,3 +1,5 @@
+import { unavailable } from '@agent-bundle/runtime';
+import type { AgentLineageRegistry, LineageToolCallQuery } from '@agent-bundle/runtime/lineage';
 import { Client, InMemoryTransport } from '@modelcontextprotocol/client';
 import { describe, expect, it } from '@rstest/core';
 import { z } from 'zod';
@@ -75,6 +77,59 @@ const stubs = (options: {
   return { host, notices, order };
 };
 
+describe('generated server lineage correlation', () => {
+  it('hands the registry the raw tools/call arguments, not the schema-parsed input with defaults applied', async () => {
+    // Cursor's hook records the arguments as sent (`tool_input`); a schema default
+    // would make `{}` and `{ label: 'probe' }` parse alike and misattribute the
+    // omitted-argument call, so the capture must read the wire, not the callback input.
+    const queries: LineageToolCallQuery[] = [];
+    const lineage: AgentLineageRegistry = {
+      observe: async () => unavailable('id-not-resolvable'),
+      resolveToolCall: async (query) => {
+        queries.push(query);
+        return unavailable('id-not-resolvable');
+      },
+      snapshot: () => ({ nodes: {}, openCalls: [], pendingChildren: [], pendingSpawns: [], seenStarts: [] }),
+    };
+    const { host } = stubs();
+    const server = await createGeneratedRouteMcpServer({
+      artifactEpoch: 'epoch',
+      host,
+      lineage,
+      plugin: { name: 'raw-arguments', version: '0.0.0' },
+      routes: {
+        'mcp/raw/tools/probe': {
+          config: {},
+          id: 'mcp/raw/tools/probe',
+          kind: 'tool',
+          module: {
+            default: () => undefined,
+            inputSchema: z.object({ label: z.string().default('probe') }).strict(),
+            resultSchema: z.object({ ok: z.boolean() }).strict(),
+          },
+          name: 'probe',
+        },
+      },
+    });
+    const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+    const client = new Client({ name: 'cursor-vscode', version: '1.0.0' });
+    await Promise.all([server.connect(serverTransport), client.connect(clientTransport)]);
+    try {
+      await client.callTool({ arguments: {}, name: 'probe' }, { signal: AbortSignal.timeout(5_000) });
+      await client.callTool({ arguments: { label: 'probe' }, name: 'probe' }, { signal: AbortSignal.timeout(5_000) });
+      await client.callTool({ arguments: { label: 'other' }, name: 'probe' }, { signal: AbortSignal.timeout(5_000) });
+      expect(queries.map((query) => [Object.hasOwn(query, 'arguments'), query.arguments, query.host, query.toolName])).toEqual([
+        [true, {}, 'cursor', 'probe'],
+        [true, { label: 'probe' }, 'cursor', 'probe'],
+        [true, { label: 'other' }, 'cursor', 'probe'],
+      ]);
+    } finally {
+      await client.close();
+      await server.close();
+    }
+  });
+});
+
 describe('generated server render completion', () => {
   it('answers a completed render while the inbox observation is still pending on another connection', async () => {
     // The signaller renews a hold for as long as a notification write takes,

packages/rsc-runtime/src/lineage/registry.ts

@@ -39,10 +39,13 @@ export interface LineageObservation {
 
 export interface LineageToolCallQuery {
   /**
-   * The `tools/call` arguments. Without a conversation id in `_meta` (Cursor
-   * sends only `progressToken`, #424) they are the one payload fact shared
-   * with the pre-tool hook's `tool_input`, and narrow the open windows when
-   * several conversations have the same tool open at once.
+   * The raw `tools/call` arguments, exactly as the client sent them. Without a
+   * conversation id in `_meta` (Cursor sends only `progressToken`, #424) they
+   * are the one payload fact shared with the pre-tool hook's `tool_input`, and
+   * narrow the open windows when several conversations have the same tool open
+   * at once. Omit the property (rather than passing `undefined`) when the raw
+   * arguments are unknown: an absent call argument is `undefined` and digests
+   * like `{}`, an absent property disables the narrowing.
    */
   readonly arguments?: unknown;
   readonly host: LineageHost | undefined;
@@ -562,7 +565,7 @@ export const createAgentLineageRegistry = (
         }
       }
       const { host, meta, toolName } = query;
-      const argumentsDigest = inputDigest(query.arguments);
+      const argumentsDigest = Object.hasOwn(query, 'arguments') ? inputDigest(query.arguments) : undefined;
       if (host === undefined) return unavailable('id-not-resolvable');
       if (host === 'codex') {
         const turn = meta?.['x-codex-turn-metadata'];
@@ -611,15 +614,18 @@ export const createAgentLineageRegistry = (
         // `mcp__<server>__<tool>` on Codex, `mcp__plugin_<p>_<s>__<tool>` on
         // Claude. Several from one conversation share a lineage; several from
         // different conversations are told apart only by the arguments the
-        // hook recorded — identical arguments stay ambiguous, never guessed.
+        // hook recorded — identical arguments stay ambiguous, never guessed,
+        // and a window opened before digests were recorded (or from a hook
+        // whose `tool_input` was not an object) could own any call, so it is
+        // never excluded.
         const matches = state.openCalls.filter((candidate) =>
           candidate.toolName === `MCP:${toolName}`
           || candidate.toolName.endsWith(`__${toolName}`)
           || candidate.toolName === toolName);
         const conversations = (candidates: readonly OpenToolCall[]): number => new Set(candidates.map((candidate) => candidate.conversation)).size;
         let narrowed = matches;
         if (conversations(matches) > 1 && argumentsDigest !== undefined) {
-          narrowed = matches.filter((candidate) => candidate.inputDigest === argumentsDigest);
+          narrowed = matches.filter((candidate) => candidate.inputDigest === undefined || candidate.inputDigest === argumentsDigest);
         }
         if (conversations(narrowed) > 1) return unavailable('id-not-resolvable');
         call = narrowed[narrowed.length - 1];

packages/rsc-runtime/tests/lineage-registry.test.ts

@@ -625,6 +625,26 @@ describe('lineage registry ambiguity refusals (review round 4)', () => {
     expect(await registry.resolveToolCall({ arguments: { limit: 10 }, host: 'cursor', toolName: 'dump' })).toMatchObject({ value: { conversation: 'root-a' } });
     expect(registry.snapshot().openCalls.find((call) => call.toolCallId === 'pa')?.inputDigest).toBeDefined();
   });
+
+  it('keeps a window without a recorded digest in contention, so an upgraded journal never attributes by elimination', async () => {
+    // A durable registry upgraded mid-session still holds pre-upgrade windows with no
+    // inputDigest (the field is optional for v1 journals). Such a window could own any
+    // call for its tool, so it is never filtered out: with a digested competitor in
+    // another conversation the call stays id-not-resolvable.
+    const registry = createAgentLineageRegistry();
+    const observe = (event: string, key: string, native: Record<string, unknown>) =>
+      registry.observe({ event, host: 'cursor', idempotencyKey: key, native });
+    await observe('prompt/submit', 'a', { conversation_id: 'root-a', hook_event_name: 'beforeSubmitPrompt' });
+    await observe('prompt/submit', 'b', { conversation_id: 'root-b', hook_event_name: 'beforeSubmitPrompt' });
+    // A pre-upgrade window: the hook carried no object tool_input, so no digest was recorded.
+    await observe('tool/before', 'legacy', { conversation_id: 'root-a', hook_event_name: 'preToolUse', tool_input: 'opaque', tool_name: 'MCP:probe', tool_use_id: 'legacy' });
+    expect(registry.snapshot().openCalls.find((call) => call.toolCallId === 'legacy')?.inputDigest).toBeUndefined();
+    await observe('tool/before', 'pb', { conversation_id: 'root-b', hook_event_name: 'preToolUse', tool_input: { note: 'nested' }, tool_name: 'MCP:probe', tool_use_id: 'pb' });
+    expect(await registry.resolveToolCall({ arguments: { note: 'nested' }, host: 'cursor', toolName: 'probe' })).toEqual(unavailable('id-not-resolvable'));
+    // Once the digested competitor closes, the legacy window resolves alone, whatever the arguments.
+    await observe('tool/after', 'pb-close', { conversation_id: 'root-b', hook_event_name: 'postToolUse', tool_input: { note: 'nested' }, tool_name: 'MCP:probe', tool_output: '{}', tool_use_id: 'pb' });
+    expect(await registry.resolveToolCall({ arguments: { note: 'nested' }, host: 'cursor', toolName: 'probe' })).toMatchObject({ value: { conversation: 'root-a' } });
+  });
 });
 
 describe('lineage registry retirement and cohorts (review round 5)', () => {