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
8 changes: 7 additions & 1 deletion src/mcp/dispatch-span-registry.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,13 @@
// means every tool call runs unwrapped, at zero cost. That asymmetry is deliberate and is why this
// is a registry rather than a direct import: Workers has no OTel collector to export to, and pulling
// the tracer into that bundle would cost real bytes for a capability it cannot use.
export type McpDispatchSpanRunner = <T>(name: string, attributes: Record<string, unknown>, fn: () => Promise<T>) => Promise<T>;
import type { SetMcpSpanOutcomeAttributes } from "./dispatch-telemetry";

export type McpDispatchSpanRunner = <T>(
name: string,
attributes: Record<string, unknown>,
fn: (setOutcomeAttributes: SetMcpSpanOutcomeAttributes) => Promise<T>,
) => Promise<T>;

let runner: McpDispatchSpanRunner | null = null;

Expand Down
11 changes: 8 additions & 3 deletions src/mcp/dispatch-telemetry-sink.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ import {
type McpAnalyticsContext,
type McpInitializeTelemetry,
} from "@loopover/contract";
import type { DispatchTelemetrySink } from "./dispatch-telemetry";
import type { DispatchTelemetrySink, SetMcpSpanOutcomeAttributes } from "./dispatch-telemetry";
import { getMcpDispatchSpanRunner } from "./dispatch-span-registry";

const DEFAULT_POSTHOG_HOST = "https://us.i.posthog.com";
Expand Down Expand Up @@ -84,7 +84,11 @@ async function captureEvents(
export function createDispatchTelemetrySink(
env: DispatchTelemetryEnv,
defer: DeferWork,
withSpan?: <T>(name: string, attributes: Record<string, unknown>, fn: () => Promise<T>) => Promise<T>,
withSpan?: <T>(
name: string,
attributes: Record<string, unknown>,
fn: (setOutcomeAttributes: SetMcpSpanOutcomeAttributes) => Promise<T>,
) => Promise<T>,
context: McpAnalyticsContext = {},
): DispatchTelemetrySink {
return {
Expand All @@ -102,7 +106,8 @@ export function createDispatchTelemetrySink(
},
// The registry is consulted per call rather than captured at construction so a self-host boot
// that fills the slot after the first request still traces.
withSpan: (name, attributes, fn) => (withSpan ?? getMcpDispatchSpanRunner() ?? ((_n, _a, run) => run()))(name, attributes, fn),
withSpan: (name, attributes, fn) =>
(withSpan ?? getMcpDispatchSpanRunner() ?? ((_n, _a, run) => run(() => {})))(name, attributes, fn),
};
}

Expand Down
62 changes: 39 additions & 23 deletions src/mcp/dispatch-telemetry.ts
Original file line number Diff line number Diff line change
Expand Up @@ -43,13 +43,20 @@ function log(level: "warn" | "error", event: string, fields: Record<string, unkn
/** The shape a tool handler returns. `isError` distinguishes "the tool answered no" from a throw. */
type ToolResultLike = { isError?: boolean; structuredContent?: unknown } | undefined;

/** Outcome attributes discovered after the handler runs. Never throws. */
export type SetMcpSpanOutcomeAttributes = (attributes: Record<string, unknown>) => void;

export type DispatchTelemetrySink = {
/** Both usage events. Never throws. */
recordToolCall: (call: McpToolCallTelemetry, properties: { usage: Record<string, unknown>; mcpToolCall: Record<string, unknown> }) => void;
/** A genuine throw. Never throws. */
captureException: (error: unknown, call: McpToolCallTelemetry) => void;
/** Wrap the call in a span when tracing is on; a no-op passthrough when it is not. */
withSpan: <T>(name: string, attributes: Record<string, unknown>, fn: () => Promise<T>) => Promise<T>;
withSpan: <T>(
name: string,
attributes: Record<string, unknown>,
fn: (setOutcomeAttributes: SetMcpSpanOutcomeAttributes) => Promise<T>,
) => Promise<T>;
/**
* Session/server/client identity for the canonical `$mcp_*` events (#10175).
*
Expand All @@ -65,7 +72,7 @@ export type DispatchTelemetrySink = {
export const NOOP_DISPATCH_SINK: DispatchTelemetrySink = {
recordToolCall: () => undefined,
captureException: () => undefined,
withSpan: async (_name, _attributes, fn) => fn(),
withSpan: async (_name, _attributes, fn) => fn(() => {}),
};

function describe(toolName: string): { category: string; excluded: boolean } {
Expand All @@ -91,7 +98,6 @@ export function instrumentToolDispatch<TArgs extends unknown[], TResult extends
return async (...args: TArgs): Promise<TResult> => {
const { category, excluded } = describe(toolName);
const startedAt = Date.now();
const attributes = { tool: toolName, category, surface: "remote" as const };

const emit = (call: McpToolCallTelemetry, payloads: { arguments?: unknown; result?: unknown }): void => {
try {
Expand All @@ -104,8 +110,16 @@ export function instrumentToolDispatch<TArgs extends unknown[], TResult extends
}
};

try {
return await sink.withSpan(mcpToolSpanName(toolName), attributes, async () => {
const publishSpanOutcome = (setOutcomeAttributes: SetMcpSpanOutcomeAttributes, call: McpToolCallTelemetry): void => {
try {
setOutcomeAttributes(buildMcpToolSpanAttributes(call));
} catch {
// Telemetry must never surface into the tool caller.
}
};

return await sink.withSpan(mcpToolSpanName(toolName), {}, async (setOutcomeAttributes) => {
try {
const result = await handler(...args);
const ok = result?.isError !== true;
const call: McpToolCallTelemetry = {
Expand All @@ -121,29 +135,31 @@ export function instrumentToolDispatch<TArgs extends unknown[], TResult extends
...(ok ? {} : { errorCode: resolveErrorCode(toolErrorEnvelope(result?.structuredContent)) }),
};
emit(call, { arguments: args[0], result: result?.structuredContent });
publishSpanOutcome(setOutcomeAttributes, call);
if (!ok) {
// One line per failed call, with the same closed property set and no payload content.
log("warn", "mcp_tool_call_failed", buildMcpToolSpanAttributes(call));
}
return result;
});
} catch (error) {
const call: McpToolCallTelemetry = {
tool: toolName,
category,
surface: "remote",
ok: false,
durationMs: Date.now() - startedAt,
errorCode: resolveErrorCode(error),
};
emit(call, { arguments: args[0] });
try {
sink.captureException(error, call);
} catch {
// Same guarantee on the crash path.
} catch (error) {
const call: McpToolCallTelemetry = {
tool: toolName,
category,
surface: "remote",
ok: false,
durationMs: Date.now() - startedAt,
errorCode: resolveErrorCode(error),
};
emit(call, { arguments: args[0] });
publishSpanOutcome(setOutcomeAttributes, call);
try {
sink.captureException(error, call);
} catch {
// Same guarantee on the crash path.
}
log("error", "mcp_tool_call_threw", buildMcpToolSpanAttributes(call));
throw error;
}
log("error", "mcp_tool_call_threw", buildMcpToolSpanAttributes(call));
throw error;
}
});
};
}
4 changes: 3 additions & 1 deletion src/server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -375,7 +375,9 @@ async function main(): Promise<void> {
// #9525: hand the MCP dispatch chokepoint a real span runner. Only this entry does -- the cloud
// Worker has no collector to export to, so its slot stays null and every tool call runs
// unwrapped. Registry rather than a direct import so ./selfhost/otel never enters that bundle.
setMcpDispatchSpanRunner((name, attributes, fn) => withOtelSpan(name, attributes, fn));
setMcpDispatchSpanRunner((name, attributes, fn) =>
withOtelSpan(name, attributes, () => fn((outcome) => setCurrentOtelSpanAttributes(outcome))),
);
}
/* v8 ignore stop */
const startedAt = Date.now();
Expand Down
32 changes: 27 additions & 5 deletions test/unit/mcp-dispatch-telemetry-sink.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,11 @@ afterEach(() => {
describe("MCP dispatch span registry (#9525)", () => {
it("is empty until a self-host boot fills it, and clears again", () => {
expect(getMcpDispatchSpanRunner()).toBeUndefined();
const runner = async <T>(_name: string, _attributes: Record<string, unknown>, fn: () => Promise<T>): Promise<T> => fn();
const runner = async <T>(
_name: string,
_attributes: Record<string, unknown>,
fn: (setOutcomeAttributes: (attributes: Record<string, unknown>) => void) => Promise<T>,
): Promise<T> => fn(() => {});
setMcpDispatchSpanRunner(runner);
expect(getMcpDispatchSpanRunner()).toBe(runner);
setMcpDispatchSpanRunner(null);
Expand Down Expand Up @@ -98,21 +102,35 @@ describe("MCP dispatch telemetry sink (#9525)", () => {
const seen: Array<{ name: string; attributes: Record<string, unknown> }> = [];
setMcpDispatchSpanRunner(async (name, attributes, fn) => {
seen.push({ name, attributes });
return fn();
return fn(() => {});
});
const sink = createDispatchTelemetrySink(env(), () => undefined);
await expect(sink.withSpan("mcp.tool/x", { tool: "x" }, async () => "wrapped")).resolves.toBe("wrapped");
expect(seen).toEqual([{ name: "mcp.tool/x", attributes: { tool: "x" } }]);
});

it("forwards setOutcomeAttributes through the registry runner (#10042)", async () => {
const outcomes: Record<string, unknown>[] = [];
setMcpDispatchSpanRunner(async (_name, _attributes, fn) => fn((attrs) => outcomes.push(attrs)));
const sink = createDispatchTelemetrySink(env(), () => undefined);
await sink.withSpan("mcp.tool/x", {}, async (setOutcome) => {
setOutcome({ ok: false, error_code: "timeout" });
});
expect(outcomes).toEqual([{ ok: false, error_code: "timeout" }]);
});

it("prefers an explicitly injected runner over the registry", async () => {
setMcpDispatchSpanRunner(async () => {
throw new Error("registry runner should not have been used");
});
let injectedCalls = 0;
const injected = async <T>(_name: string, _attributes: Record<string, unknown>, fn: () => Promise<T>): Promise<T> => {
const injected = async <T>(
_name: string,
_attributes: Record<string, unknown>,
fn: (setOutcomeAttributes: (attributes: Record<string, unknown>) => void) => Promise<T>,
): Promise<T> => {
injectedCalls += 1;
return fn();
return fn(() => {});
};
const sink = createDispatchTelemetrySink(env(), () => undefined, injected);
await expect(sink.withSpan("mcp.tool/x", {}, async () => "injected")).resolves.toBe("injected");
Expand All @@ -131,7 +149,11 @@ describe("LoopoverMcp telemetry-sink injection (#9525)", () => {
const sink = {
recordToolCall: (entry: McpToolCallTelemetry) => recorded.push(entry),
captureException: () => undefined,
withSpan: async <T>(_name: string, _attributes: Record<string, unknown>, fn: () => Promise<T>) => fn(),
withSpan: async <T>(
_name: string,
_attributes: Record<string, unknown>,
fn: (setOutcomeAttributes: (attributes: Record<string, unknown>) => void) => Promise<T>,
) => fn(() => {}),
};

const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
Expand Down
82 changes: 79 additions & 3 deletions test/unit/mcp-dispatch-telemetry.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -240,7 +240,7 @@ describe("MCP dispatch chokepoint (#9525)", () => {
sink: {
recordToolCall: (recorded) => calls.push(recorded),
captureException: (error) => exceptions.push(error),
withSpan: async (_name, _attributes, fn) => fn(),
withSpan: async (_name, _attributes, fn) => fn(() => {}),
},
};
};
Expand Down Expand Up @@ -316,7 +316,7 @@ describe("MCP dispatch chokepoint (#9525)", () => {
captureException: () => {
throw new Error("sink down");
},
withSpan: async (_name, _attributes, fn) => fn(),
withSpan: async (_name, _attributes, fn) => fn(() => {}),
};
const ok = instrumentToolDispatch("loopover_get_repo_context", hostile, async (_args: unknown) => ({ structuredContent: { fine: true } }));
await expect(ok({})).resolves.toMatchObject({ structuredContent: { fine: true } });
Expand All @@ -338,7 +338,10 @@ describe("MCP dispatch chokepoint (#9525)", () => {
// every single call, so "does nothing and returns nothing" is worth asserting outright.
expect(NOOP_DISPATCH_SINK.recordToolCall(call, { usage: {}, mcpToolCall: {} })).toBeUndefined();
expect(NOOP_DISPATCH_SINK.captureException(new Error("x"), call)).toBeUndefined();
await expect(NOOP_DISPATCH_SINK.withSpan("n", {}, async () => "through")).resolves.toBe("through");
await expect(NOOP_DISPATCH_SINK.withSpan("n", {}, async (setOutcome) => {
setOutcome({ ok: false });
return "through";
})).resolves.toBe("through");
});

it("falls back to the unknown category for a tool with no contract entry", async () => {
Expand All @@ -349,6 +352,79 @@ describe("MCP dispatch chokepoint (#9525)", () => {
});
});

describe("MCP dispatch span outcome attributes (#10042)", () => {
it("publishes buildMcpToolSpanAttributes onto the span on success and on throw", async () => {
const error = vi.spyOn(console, "error").mockImplementation(() => undefined);
const spans: Array<{ name: string; open: Record<string, unknown>; outcome?: Record<string, unknown> }> = [];
const recordingSink: DispatchTelemetrySink = {
recordToolCall: () => undefined,
captureException: () => undefined,
withSpan: async (name, attributes, fn) => {
let outcome: Record<string, unknown> | undefined;
try {
const result = await fn((attrs) => {
outcome = attrs;
});
spans.push({ name, open: attributes, ...(outcome !== undefined ? { outcome } : {}) });
return result;
} catch (error) {
spans.push({ name, open: attributes, ...(outcome !== undefined ? { outcome } : {}) });
throw error;
}
},
};

const okWrapped = instrumentToolDispatch("loopover_get_repo_context", recordingSink, async (_args: unknown) => ({
structuredContent: { ok: 1 },
}));
await okWrapped({});
expect(spans[0]).toMatchObject({ name: "mcp.tool/loopover_get_repo_context", open: {} });
expect(spans[0]!.outcome).toMatchObject({
tool: "loopover_get_repo_context",
category: "maintainer",
surface: "remote",
transport: "local",
ok: true,
});
expect("error_code" in spans[0]!.outcome!).toBe(false);

const throwWrapped = instrumentToolDispatch("loopover_get_repo_context", recordingSink, async (_args: unknown) => {
throw new Error("request timed out");
});
await expect(throwWrapped({})).rejects.toThrow("request timed out");
expect(spans[1]!.outcome).toMatchObject({ ok: false, error_code: "timeout" });
expect(MCP_TELEMETRY_ERROR_CODES).toContain(spans[1]!.outcome!.error_code);

error.mockRestore();
});

it("keeps NOOP_DISPATCH_SINK.withSpan a pure passthrough that records nothing", async () => {
let setterCalled = false;
await expect(
NOOP_DISPATCH_SINK.withSpan("mcp.tool/x", { tool: "x" }, async (setOutcome) => {
setOutcome({ ok: false, error_code: "timeout" });
setterCalled = true;
return "through";
}),
).resolves.toBe("through");
expect(setterCalled).toBe(true);
});

it("never lets a failing setOutcomeAttributes reach the caller", async () => {
const recordingSink: DispatchTelemetrySink = {
recordToolCall: () => undefined,
captureException: () => undefined,
withSpan: async (_name, _attributes, fn) => fn(() => {
throw new Error("span attrs down");
}),
};
const wrapped = instrumentToolDispatch("loopover_get_repo_context", recordingSink, async (_args: unknown) => ({
structuredContent: { ok: 1 },
}));
await expect(wrapped({})).resolves.toMatchObject({ structuredContent: { ok: 1 } });
});
});

describe("PostHog canonical MCP analytics contract (#10175)", () => {
const ctx = { sessionId: "ses_abc123", serverName: "loopover", serverVersion: "3.18.4", clientName: "claude-code", clientVersion: "1.2.3" };

Expand Down
30 changes: 30 additions & 0 deletions test/unit/selfhost-otel.test.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { beforeEach, describe, expect, it, vi } from "vitest";
import { buildMcpToolSpanAttributes } from "@loopover/contract";

const otelMocks = vi.hoisted(() => {
const exportedSpans: any[] = [];
Expand Down Expand Up @@ -484,6 +485,35 @@ describe("self-host OpenTelemetry", () => {
await expect(reviewTraceAttributes({})).resolves.toEqual({});
});

it("applies MCP dispatch span outcome attributes through otelSafeAttributes (#10042)", async () => {
await initOpenTelemetry(env({
OTEL_TRACES_EXPORTER: "otlp",
OTEL_EXPORTER_OTLP_TRACES_ENDPOINT: "http://collector/v1/traces",
}));
const call = {
tool: "loopover_get_repo_context",
category: "maintainer",
surface: "remote" as const,
ok: false,
durationMs: 7,
errorCode: "timeout" as const,
};
await withOtelSpan("mcp.tool/loopover_get_repo_context", {}, () =>
Promise.resolve().then(() => setCurrentOtelSpanAttributes(buildMcpToolSpanAttributes(call))),
);
await flushOpenTelemetry();
const span = otelMocks.exportedSpans.find((entry) => entry.name === "mcp.tool/loopover_get_repo_context");
expect(span.attributes).toMatchObject({
tool: "loopover_get_repo_context",
category: "maintainer",
surface: "remote",
transport: "local",
ok: false,
duration_ms: 7,
error_code: "timeout",
});
});

it("swallows exporter flush and shutdown failures", async () => {
await initOpenTelemetry(env({
OTEL_TRACES_EXPORTER: "otlp",
Expand Down