From c06c34f0a26513491516ea5e695c545f3ed897fb Mon Sep 17 00:00:00 2001 From: Michal Date: Tue, 25 Aug 2026 23:32:10 +0100 Subject: [PATCH] feat(mcplocal): a trace code that actually reaches the audit trail The plumbing for request tracing was all present and almost entirely inert. correlationId was generated at the endpoint and threaded into traffic events, but: - emitTrace() dropped it, so every durable tool_call_trace row had correlationId = null and could not be joined to anything; - ExecuteOptions.correlationId existed but its only production caller never passed it, making every stage_execution and pipeline_execution row unjoinable too -- dead code that looked like a feature; - plugin events (all three gate emits) had no access to it at all; - and the value itself was `:`, which nobody can retype from a screenshot. Also fixed: onToolCallBefore intercepts returned without emitTrace, so every tool call made while a session was gated produced no trace whatsoever. The trace code now IS the correlationId. Nothing parses the old format -- checked across mcpd's audit repository and route, mcplocal's router and traffic capture, and the CLI console -- sessionId is its own audit column and the JSON-RPC id is in the traffic body, so nothing is lost. AuditEvent.correlationId is already an indexed Postgres column, so lookup works with no migration. 8 chars of Crockford base32 without I/L/O/U: readable aloud, typeable from a screenshot, 40 bits (a collision shows two traces, it corrupts nothing). Getting it to the plugins and the executor needed AsyncLocalStorage, not a parameter. Both are constructed per SESSION, not per request -- processContent is a closure with no RouteContext -- so threading an argument would have touched PluginSessionContext, every plugin, ExecuteOptions and StageContext. ALS also stays correct with two requests in flight on one session, which a mutable field on the session-scoped context would not; there is a test for exactly that. PluginContextImpl.emitAuditEvent auto-fills it the way it already auto-fills sessionId, so this covers every gate event and every plugin written later without a per-callsite edit. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01GqMidYEGUJG5fxeoTELBu2 --- src/mcplocal/src/http/project-mcp-endpoint.ts | 22 +++++-- src/mcplocal/src/http/traffic.ts | 5 ++ src/mcplocal/src/proxymodel/plugin-context.ts | 7 ++- src/mcplocal/src/request-context.ts | 37 +++++++++++ src/mcplocal/src/router.ts | 16 ++++- src/mcplocal/src/util/trace-code.ts | 28 +++++++++ src/mcplocal/tests/request-context.test.ts | 62 +++++++++++++++++++ 7 files changed, 170 insertions(+), 7 deletions(-) create mode 100644 src/mcplocal/src/request-context.ts create mode 100644 src/mcplocal/src/util/trace-code.ts create mode 100644 src/mcplocal/tests/request-context.test.ts diff --git a/src/mcplocal/src/http/project-mcp-endpoint.ts b/src/mcplocal/src/http/project-mcp-endpoint.ts index 22d239a..7f64636 100644 --- a/src/mcplocal/src/http/project-mcp-endpoint.ts +++ b/src/mcplocal/src/http/project-mcp-endpoint.ts @@ -28,6 +28,8 @@ import { composePlugins } from '../proxymodel/plugins/compose.js'; import type { ProxyModelPlugin } from '../proxymodel/plugin.js'; import { AuditCollector } from '../audit/collector.js'; import { WireNameCodec, routeWithWireNames } from '../util/wire-names.js'; +import { newTraceCode } from '../util/trace-code.js'; +import { runInRequestScope } from '../request-context.js'; interface ProjectCacheEntry { router: McpRouter; @@ -325,7 +327,11 @@ export function registerProjectMcpEndpoint(app: FastifyInstance, mcpdClient: Mcp const requestId = message.id as string | number; const sid = transport.sessionId ?? 'unknown'; const method = (message as { method?: string }).method; - const correlationId = `${sid}:${requestId}`; + // Short, transcribable, and the audit correlationId itself. Nothing + // parses the old `:` form, sessionId is its own audit column + // and the JSON-RPC id is in the traffic body, so no migration is needed + // and `mcpctl trace ` works against the existing index. + const correlationId = newTraceCode(); requestCorrelations.set(requestId, correlationId); // Capture client request @@ -347,10 +353,16 @@ export function registerProjectMcpEndpoint(app: FastifyInstance, mcpdClient: Mcp codec = new WireNameCodec(); wireCodecs.set(projectName, codec); } - const response = await routeWithWireNames( - codec, - (req) => router.route(req, ctx), - message as unknown as JsonRpcRequest, + // Everything downstream — plugins, stages, the pipeline executor — is + // built per SESSION, not per request, so the trace code reaches them + // through the async scope rather than a parameter on every signature. + const response = await runInRequestScope( + { correlationId, sessionId: sid, projectName, method: method ?? 'unknown' }, + () => routeWithWireNames( + codec, + (req) => router.route(req, ctx), + message as unknown as JsonRpcRequest, + ), ); // Forward queued notifications BEFORE the response — the response send diff --git a/src/mcplocal/src/http/traffic.ts b/src/mcplocal/src/http/traffic.ts index f41e94f..745aba8 100644 --- a/src/mcplocal/src/http/traffic.ts +++ b/src/mcplocal/src/http/traffic.ts @@ -36,6 +36,8 @@ export interface ActiveSession { export interface TrafficFilter { project?: string | undefined; session?: string | undefined; + /** Trace code — narrows the buffer to one request's events. */ + correlationId?: string | undefined; } type Listener = (event: TrafficEvent) => void; @@ -97,6 +99,9 @@ export class TrafficCapture { if (filter?.session) { events = events.filter((e) => e.sessionId === filter.session); } + if (filter?.correlationId) { + events = events.filter((e) => e.correlationId === filter.correlationId); + } return events; } diff --git a/src/mcplocal/src/proxymodel/plugin-context.ts b/src/mcplocal/src/proxymodel/plugin-context.ts index 5575692..fbbc338 100644 --- a/src/mcplocal/src/proxymodel/plugin-context.ts +++ b/src/mcplocal/src/proxymodel/plugin-context.ts @@ -9,6 +9,7 @@ import type { LLMProvider, CacheProvider, StageLogger, Section, ToolDefinition, import type { PluginSessionContext, VirtualToolHandler, VirtualServer, PromptIndexEntry } from './plugin.js'; import type { AuditCollector } from '../audit/collector.js'; import type { AuditEvent } from '../audit/types.js'; +import { currentCorrelationId } from '../request-context.js'; /** Dependencies injected from the router into each context. */ export interface PluginContextDeps { @@ -119,11 +120,15 @@ export class PluginContextImpl implements PluginSessionContext { return this.deps.getFromMcpd(path); } - /** Emit an audit event, auto-filling sessionId and projectName. */ + /** Emit an audit event, auto-filling sessionId, projectName and correlationId. */ emitAuditEvent(event: Omit): void { + // Auto-filled here rather than at each call site: this covers every gate + // event and every plugin written later, with no per-callsite edit. + const correlationId = event.correlationId ?? currentCorrelationId(); this.deps.auditCollector?.emit({ ...event, sessionId: this.sessionId, + ...(correlationId !== undefined ? { correlationId } : {}), }); } } diff --git a/src/mcplocal/src/request-context.ts b/src/mcplocal/src/request-context.ts new file mode 100644 index 0000000..25f8318 --- /dev/null +++ b/src/mcplocal/src/request-context.ts @@ -0,0 +1,37 @@ +/** + * Per-request scope, carried through code that has no request parameter. + * + * The plugin context and the pipeline executor are both constructed **per + * session**, not per request: `getOrCreatePluginContext` builds one + * PluginContextDeps per session, and `processContent` is a closure with no + * RouteContext argument. Threading a correlationId parameter would touch + * PluginSessionContext, every plugin, ExecuteOptions and StageContext. + * + * AsyncLocalStorage gets it there with no signature churn, and — unlike a + * mutable field on the session-scoped context — stays correct when two + * requests are in flight on the same session. + */ +import { AsyncLocalStorage } from 'node:async_hooks'; + +export interface RequestScope { + /** The trace code; also the audit correlationId. */ + correlationId: string; + sessionId: string; + projectName: string; + method: string; +} + +const storage = new AsyncLocalStorage(); + +export function runInRequestScope(scope: RequestScope, fn: () => T): T { + return storage.run(scope, fn); +} + +export function currentScope(): RequestScope | undefined { + return storage.getStore(); +} + +/** Convenience: the current trace code, or undefined outside a request. */ +export function currentCorrelationId(): string | undefined { + return storage.getStore()?.correlationId; +} diff --git a/src/mcplocal/src/router.ts b/src/mcplocal/src/router.ts index 3715a01..06fbbf1 100644 --- a/src/mcplocal/src/router.ts +++ b/src/mcplocal/src/router.ts @@ -6,6 +6,7 @@ import type { PromptIndexEntry } from './gate/tag-matcher.js'; import { LinkResolver } from './services/link-resolver.js'; import type { LLMProvider, CacheProvider, Section } from './proxymodel/types.js'; import { executePipeline } from './proxymodel/executor.js'; +import { currentCorrelationId } from './request-context.js'; import { getProxyModel } from './proxymodel/loader.js'; import type { ProxyModelPlugin, PluginSessionContext } from './proxymodel/plugin.js'; import { PluginContextImpl, type PluginContextDeps } from './proxymodel/plugin-context.js'; @@ -175,6 +176,11 @@ export class McpRouter { getSystemPrompt: (name, fallback) => this.getSystemPrompt(name, fallback), ...(this.auditCollector ? { auditCollector: this.auditCollector } : {}), ...(serverName !== undefined ? { serverName } : {}), + // processContent is a per-session closure with no RouteContext, so + // the id comes from the request scope. Without this, + // ExecuteOptions.correlationId was dead code and every + // stage_execution / pipeline_execution row landed unjoinable. + ...(currentCorrelationId() !== undefined ? { correlationId: currentCorrelationId()! } : {}), }); // Pause queue: if paused, hold the result until released/edited/dropped @@ -868,6 +874,9 @@ export class McpRouter { eventKind: 'tool_call_trace' as const, source: 'mcplocal' as const, verified: true, + // Already in scope and previously dropped, which left every durable + // tool_call_trace row unjoinable to the rest of its request. + ...(context.correlationId !== undefined ? { correlationId: context.correlationId } : {}), payload: { toolName, argKeys: Object.keys(toolArgs).join(', '), @@ -903,7 +912,12 @@ export class McpRouter { // onToolCallBefore — can intercept and return a response directly if (this.plugin.onToolCallBefore) { const intercepted = await this.plugin.onToolCallBefore(toolName ?? '', toolArgs, request, ctx); - if (intercepted) return intercepted; + if (intercepted) { + // Intercepts were previously invisible in tool_call_trace, so every + // call made while the session was gated produced no trace at all. + emitTrace(intercepted, 'plugin'); + return intercepted; + } } // Route to upstream diff --git a/src/mcplocal/src/util/trace-code.ts b/src/mcplocal/src/util/trace-code.ts new file mode 100644 index 0000000..75d6742 --- /dev/null +++ b/src/mcplocal/src/util/trace-code.ts @@ -0,0 +1,28 @@ +/** + * Short, human-transcribable code identifying one MCP request end to end. + * + * This IS the correlationId — it replaces the previous `:` + * form. Nothing parses that format (checked across every consumer: mcpd's + * audit repository and route, mcplocal's router and traffic capture, and the + * CLI console), `sessionId` is its own column on every audit event, and the + * JSON-RPC id is in the traffic body — so nothing is lost and no migration is + * needed. `AuditEvent.correlationId` is already an indexed Postgres column, so + * lookup by code works from day one. + * + * Crockford base32 without I/L/O/U: unambiguous read aloud, typed from a + * screenshot, or pasted into a bug report. 40 bits of randomness — a collision + * shows two traces, it does not corrupt anything. + */ +import { randomBytes } from 'node:crypto'; + +const ALPHABET = '0123456789ABCDEFGHJKMNPQRSTVWXYZ'; + +export function newTraceCode(): string { + const buf = randomBytes(5); + const value = buf.readUIntBE(0, 5); + let out = ''; + for (let i = 0; i < 8; i++) { + out += ALPHABET[Math.floor(value / 32 ** (7 - i)) % 32]; + } + return out; +} diff --git a/src/mcplocal/tests/request-context.test.ts b/src/mcplocal/tests/request-context.test.ts new file mode 100644 index 0000000..aec215d --- /dev/null +++ b/src/mcplocal/tests/request-context.test.ts @@ -0,0 +1,62 @@ +import { describe, it, expect } from 'vitest'; +import { runInRequestScope, currentScope, currentCorrelationId } from '../src/request-context.js'; +import { newTraceCode } from '../src/util/trace-code.js'; + +describe('newTraceCode', () => { + it('is 8 unambiguous Crockford characters', () => { + for (let i = 0; i < 200; i++) { + expect(newTraceCode()).toMatch(/^[0-9ABCDEFGHJKMNPQRSTVWXYZ]{8}$/); + } + }); + + it('excludes I, L, O and U so a code can be read aloud or retyped', () => { + const codes = Array.from({ length: 500 }, () => newTraceCode()).join(''); + expect(codes).not.toMatch(/[ILOU]/); + }); + + it('does not collide over a realistic request volume', () => { + const seen = new Set(Array.from({ length: 20_000 }, () => newTraceCode())); + expect(seen.size).toBe(20_000); + }); +}); + +describe('request scope', () => { + it('is undefined outside a request', () => { + expect(currentScope()).toBeUndefined(); + expect(currentCorrelationId()).toBeUndefined(); + }); + + it('survives awaits', async () => { + await runInRequestScope( + { correlationId: 'ABC12345', sessionId: 's', projectName: 'p', method: 'tools/call' }, + async () => { + await new Promise((r) => setTimeout(r, 5)); + expect(currentCorrelationId()).toBe('ABC12345'); + await new Promise((r) => setTimeout(r, 5)); + expect(currentScope()?.method).toBe('tools/call'); + }, + ); + }); + + it('does not cross-contaminate concurrent requests on one session', async () => { + // The case a mutable field on the session-scoped plugin context would get + // wrong — two in-flight requests share a session but not a trace code. + const seen: string[] = []; + const one = runInRequestScope( + { correlationId: 'AAAAAAAA', sessionId: 'same', projectName: 'p', method: 'tools/call' }, + async () => { + await new Promise((r) => setTimeout(r, 20)); + seen.push(currentCorrelationId() ?? 'none'); + }, + ); + const two = runInRequestScope( + { correlationId: 'BBBBBBBB', sessionId: 'same', projectName: 'p', method: 'tools/call' }, + async () => { + await new Promise((r) => setTimeout(r, 5)); + seen.push(currentCorrelationId() ?? 'none'); + }, + ); + await Promise.all([one, two]); + expect(seen.sort()).toEqual(['AAAAAAAA', 'BBBBBBBB']); + }); +});