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']); + }); +});