History

mcp: make hung executions detectable in tracing

Rhys Sullivan committed 3 months ago

Commit 774ce451 parent fc1e5894 files · +289 −6Browse files

Unified · Split

4 files
  1. apps/cloud/src/mcp/telemetry.test.ts +93 −0
  2. apps/cloud/src/mcp/telemetry.ts +20 −0
  3. packages/hosts/mcp/src/tool-server.test.ts +113 −0
  4. packages/hosts/mcp/src/tool-server.ts +63 −6
apps/cloud/src/mcp/telemetry.test.ts+93−0
@@ -0,0 +1,93 @@
1+import { describe, expect, it } from "@effect/vitest";
2+import { Effect } from "effect";
3+import type * as Tracer from "effect/Tracer";
4+
5+import { annotateMcpRequest } from "./telemetry";
6+
7+// A tracer that records attributes per span so the test can assert what
8+// `annotateMcpRequest` stamps on the active request span.
9+const makeRecordingTracer = (): {
10+ tracer: Tracer.Tracer;
11+ attributesOf: (name: string) => ReadonlyMap<string, unknown> | undefined;
12+} => {
13+ const spans = new Map<string, Map<string, unknown>>();
14+ const tracer: Tracer.Tracer = {
15+ span: (options) => {
16+ const attributes = spans.get(options.name) ?? new Map<string, unknown>();
17+ spans.set(options.name, attributes);
18+ let status: Tracer.SpanStatus = { _tag: "Started", startTime: options.startTime };
19+ return {
20+ _tag: "Span",
21+ name: options.name,
22+ spanId: `span-${spans.size}`,
23+ traceId: "trace-1",
24+ parent: options.parent,
25+ annotations: options.annotations,
26+ get status() {
27+ return status;
28+ },
29+ attributes,
30+ links: options.links,
31+ sampled: options.sampled,
32+ kind: options.kind,
33+ end: (endTime, exit) => {
34+ status = { _tag: "Ended", startTime: options.startTime, endTime, exit };
35+ },
36+ attribute: (key, value) => {
37+ attributes.set(key, value);
38+ },
39+ event: () => undefined,
40+ addLinks: () => undefined,
41+ };
42+ },
43+ };
44+ return { tracer, attributesOf: (name) => spans.get(name) };
45+};
46+
47+const postRequest = (body: unknown): Request =>
48+ new Request("https://executor.sh/mcp", {
49+ method: "POST",
50+ headers: { "content-type": "application/json", "mcp-session-id": "sess-123" },
51+ body: JSON.stringify(body),
52+ });
53+
54+const annotate = (request: Request) =>
55+ annotateMcpRequest(request, { token: null, parseBody: true });
56+
57+describe("annotateMcpRequest — cancellation join keys", () => {
58+ it.effect("stamps the cancelled request id and reason from notifications/cancelled", () => {
59+ const { tracer, attributesOf } = makeRecordingTracer();
60+ return Effect.gen(function* () {
61+ yield* annotate(
62+ postRequest({
63+ jsonrpc: "2.0",
64+ method: "notifications/cancelled",
65+ params: { requestId: 42, reason: "client timeout" },
66+ }),
67+ );
68+ const attributes = attributesOf("mcp.request")!;
69+ expect(attributes.get("mcp.rpc.method")).toBe("notifications/cancelled");
70+ expect(attributes.get("mcp.rpc.cancelled_id")).toBe("42");
71+ expect(attributes.get("mcp.rpc.cancelled_reason")).toBe("client timeout");
72+ expect(attributes.get("mcp.request.session_id")).toBe("sess-123");
73+ }).pipe(Effect.withSpan("mcp.request"), Effect.withTracer(tracer));
74+ });
75+
76+ it.effect("still stamps mcp.rpc.id on ordinary requests", () => {
77+ const { tracer, attributesOf } = makeRecordingTracer();
78+ return Effect.gen(function* () {
79+ yield* annotate(
80+ postRequest({
81+ jsonrpc: "2.0",
82+ id: 7,
83+ method: "tools/call",
84+ params: { name: "execute" },
85+ }),
86+ );
87+ const attributes = attributesOf("mcp.request")!;
88+ expect(attributes.get("mcp.rpc.id")).toBe("7");
89+ expect(attributes.get("mcp.tool.name")).toBe("execute");
90+ expect(attributes.get("mcp.rpc.cancelled_id")).toBeUndefined();
91+ }).pipe(Effect.withSpan("mcp.request"), Effect.withTracer(tracer));
92+ });
93+});
apps/cloud/src/mcp/telemetry.ts+20−0
@@ -105,12 +105,23 @@ const InitializeParams = Schema.Struct({ · 104 unmodified lines
105105 const NamedParams = Schema.Struct({ name: Schema.optional(Schema.String) });
106106 const UriParams = Schema.Struct({ uri: Schema.optional(Schema.String) });
107107
108+// `notifications/cancelled` carries no JSON-RPC id of its own, but its params
109+// name the request the client gave up on. Surfacing that id as
110+// `mcp.rpc.cancelled_id` makes "a client gave up" joinable to the exact
111+// `tools/call` (and the execution spans, which carry `mcp.rpc.id`) that hung —
112+// the best available proxy for a killed execution.
113+const CancelledParams = Schema.Struct({
114+ requestId: Schema.optional(Schema.Union([Schema.String, Schema.Number])),
115+ reason: Schema.optional(Schema.String),
116+});
117+
108118 const decodeJsonRpcEnvelopeString = Schema.decodeUnknownOption(
109119 Schema.fromJsonString(JsonRpcEnvelope),
110120 );
111121 const decodeInitializeParams = Schema.decodeUnknownOption(InitializeParams);
112122 const decodeNamedParams = Schema.decodeUnknownOption(NamedParams);
113123 const decodeUriParams = Schema.decodeUnknownOption(UriParams);
124+const decodeCancelledParams = Schema.decodeUnknownOption(CancelledParams);
114125 const decodeElicitationReplyResult = Schema.decodeUnknownOption(ElicitationReplyResult);
115126
116127 const readJsonRpcEnvelope = (request: Request): Effect.Effect<Option.Option<JsonRpcEnvelope>> =>
@@ -158,6 +169,15 @@ const methodAttrs = (envelope: JsonRpcEnvelope): Record<string, unknown> => { · 41 unmodified lines
158169 onSome: ({ name }) => (name ? { "mcp.prompt.name": name } : {}),
159170 }),
160171 ),
172+ Match.when("notifications/cancelled", () =>
173+ Option.match(decodeCancelledParams(params), {
174+ onNone: () => ({}) as Record<string, unknown>,
175+ onSome: ({ requestId, reason }) => ({
176+ ...(requestId !== undefined && { "mcp.rpc.cancelled_id": String(requestId) }),
177+ ...(reason && { "mcp.rpc.cancelled_reason": reason }),
178+ }),
179+ }),
180+ ),
161181 Match.option,
162182 Option.getOrElse(() => ({}) as Record<string, unknown>),
163183 );
packages/hosts/mcp/src/tool-server.test.ts+113−0
@@ -1,5 +1,6 @@
11 import { describe, expect, it } from "@effect/vitest";
22 import { Data, Deferred, Effect } from "effect";
3+import type * as Tracer from "effect/Tracer";
34 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
45 import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
56 import { ElicitRequestSchema } from "@modelcontextprotocol/sdk/types.js";
@@ -88,6 +89,72 @@ const withNativeClient = async <E extends Cause.YieldableError>( · 82 unmodified lines
8889 fn: (client: Client) => Promise<void>,
8990 ) => withClient(engine, capabilities, fn, { elicitationMode: { mode: "native" } });
9091
92+type RecordedSpan = {
93+ readonly name: string;
94+ readonly attributes: ReadonlyMap<string, unknown>;
95+ ended: boolean;
96+};
97+
98+/** A tracer that records every span + its attributes and end state. */
99+const makeRecordingTracer = (): { tracer: Tracer.Tracer; spans: RecordedSpan[] } => {
100+ const spans: RecordedSpan[] = [];
101+ const tracer: Tracer.Tracer = {
102+ span: (options) => {
103+ const attributes = new Map<string, unknown>();
104+ const recorded: RecordedSpan = { name: options.name, attributes, ended: false };
105+ spans.push(recorded);
106+ let status: Tracer.SpanStatus = { _tag: "Started", startTime: options.startTime };
107+ return {
108+ _tag: "Span",
109+ name: options.name,
110+ spanId: `span-${spans.length}`,
111+ traceId: "trace-1",
112+ parent: options.parent,
113+ annotations: options.annotations,
114+ get status() {
115+ return status;
116+ },
117+ attributes,
118+ links: options.links,
119+ sampled: options.sampled,
120+ kind: options.kind,
121+ end: (endTime, exit) => {
122+ recorded.ended = true;
123+ status = { _tag: "Ended", startTime: options.startTime, endTime, exit };
124+ },
125+ attribute: (key, value) => {
126+ attributes.set(key, value);
127+ },
128+ event: () => undefined,
129+ addLinks: () => undefined,
130+ };
131+ },
132+ };
133+ return { tracer, spans };
134+};
135+
136+/** withClient, but with a recording tracer installed for the whole server. */
137+const withTracedClient = async <E extends Cause.YieldableError>(
138+ engine: ExecutionEngine<E>,
139+ fn: (client: Client, spans: RecordedSpan[]) => Promise<void>,
140+) => {
141+ const { tracer, spans } = makeRecordingTracer();
142+ const mcpServer = await Effect.runPromise(
143+ createExecutorMcpServer({ engine }).pipe(Effect.withTracer(tracer)),
144+ );
145+ const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
146+ const client = new Client({ name: "test-client", version: "1.0.0" }, { capabilities: NO_CAPS });
147+ await mcpServer.connect(serverTransport);
148+ await client.connect(clientTransport);
149+ // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: test helper must close MCP transports after async client assertions
150+ try {
151+ await fn(client, spans);
152+ } finally {
153+ await clientTransport.close();
154+ await serverTransport.close();
155+ }
156+};
157+
91158 const ELICITATION_CAPS: ClientCapabilities = {
92159 elicitation: { form: {}, url: {} },
93160 };
@@ -1758,3 +1825,49 @@ describe("MCP host server — skills tool", () => { · 1664 unmodified lines
17581825 });
17591826 });
17601827 });
1828+
1829+describe("MCP host server — hang-visibility tracing", () => {
1830+ it("execute emits a start marker and stamps the JSON-RPC id on execution spans", async () => {
1831+ const engine = makeStubEngine({});
1832+
1833+ await withTracedClient(engine, async (client, spans) => {
1834+ await client.callTool({ name: "execute", arguments: { code: "1+1" } });
1835+
1836+ // The start marker exports immediately: a start without a matching
1837+ // completed `mcp.host.tool.execute` is the killed-execution signal.
1838+ const start = spans.find((span) => span.name === "mcp.host.tool.execute.start");
1839+ expect(start).toBeDefined();
1840+ expect(start!.ended).toBe(true);
1841+ expect(start!.attributes.get("mcp.rpc.id")).toBeDefined();
1842+
1843+ // The completion span carries the same join key.
1844+ const execute = spans.find((span) => span.name === "mcp.host.tool.execute");
1845+ expect(execute).toBeDefined();
1846+ expect(execute!.attributes.get("mcp.rpc.id")).toBe(start!.attributes.get("mcp.rpc.id"));
1847+ expect(execute!.ended).toBe(true);
1848+ });
1849+ });
1850+
1851+ it("resume emits a start marker carrying the execution id and rpc id", async () => {
1852+ const engine = makeStubEngine({
1853+ resume: () => Effect.succeed({ status: "completed", result: { result: "resumed" } }),
1854+ });
1855+
1856+ await withTracedClient(engine, async (client, spans) => {
1857+ await client.callTool({
1858+ name: "resume",
1859+ arguments: { executionId: "exec-1", action: "accept" },
1860+ });
1861+
1862+ const start = spans.find((span) => span.name === "mcp.host.tool.resume.start");
1863+ expect(start).toBeDefined();
1864+ expect(start!.ended).toBe(true);
1865+ expect(start!.attributes.get("mcp.execute.execution_id")).toBe("exec-1");
1866+ expect(start!.attributes.get("mcp.rpc.id")).toBeDefined();
1867+
1868+ const resume = spans.find((span) => span.name === "mcp.host.tool.resume");
1869+ expect(resume).toBeDefined();
1870+ expect(resume!.attributes.get("mcp.rpc.id")).toBe(start!.attributes.get("mcp.rpc.id"));
1871+ });
1872+ });
1873+});
packages/hosts/mcp/src/tool-server.ts+63−6
@@ -647,6 +647,39 @@ const extractInventory = (description: string): string => { · 646 unmodified lines
647647 return index === -1 ? "" : description.slice(index).trimEnd();
648648 };
649649
650+// ---------------------------------------------------------------------------
651+// Hang-visibility join keys
652+// ---------------------------------------------------------------------------
653+// A killed execution exports nothing: OTEL only ships a span when it ends, and
654+// a Cloudflare deploy/eviction cancels the request without an error, so a hung
655+// `execute` is invisible in the trace store. Two mitigations live here:
656+// 1. Every execution-path span carries the JSON-RPC id + transport session id
657+// (`mcp.rpc.id`, `mcp.request.session_id`), so a client's
658+// `notifications/cancelled` — which names the cancelled request id — can
659+// be joined to the exact call it gave up on.
660+// 2. A zero-duration start marker span is emitted the moment execution
661+// begins. It ends (and therefore exports) immediately, so a start marker
662+// without a matching completion span is a true positive for an execution
663+// that died mid-flight.
664+
665+type McpRequestJoinKeys = {
666+ readonly requestId: string | number;
667+ readonly sessionId?: string | undefined;
668+};
669+
670+const joinKeyAttributes = (extra: McpRequestJoinKeys): Record<string, unknown> => ({
671+ "mcp.rpc.id": String(extra.requestId),
672+ ...(extra.sessionId ? { "mcp.request.session_id": extra.sessionId } : {}),
673+});
674+
675+const withJoinKeys = <A, EffE>(
676+ effect: Effect.Effect<A, EffE>,
677+ extra: McpRequestJoinKeys,
678+): Effect.Effect<A, EffE> => Effect.annotateSpans(effect, joinKeyAttributes(extra));
679+
680+const startMarker = (name: string, attributes: Record<string, unknown>): Effect.Effect<void> =>
681+ Effect.void.pipe(Effect.withSpan(name, { attributes }));
682+
650683 const JsonObjectFromString = Schema.fromJsonString(Schema.Record(Schema.String, Schema.Unknown));
651684 const decodeJsonObjectString = Schema.decodeUnknownOption(JsonObjectFromString);
652685
@@ -764,8 +797,15 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 111 unmodified lines
764797 ),
765798 ).pipe(Effect.withSpan("mcp.host.create_server"));
766799
767− const executeCode = (code: string): Effect.Effect<McpToolResult, E> =>
800+ const executeCode = (
801+ code: string,
802+ extra: McpRequestJoinKeys,
803+ ): Effect.Effect<McpToolResult, E> =>
768804 Effect.gen(function* () {
805+ yield* startMarker("mcp.host.tool.execute.start", {
806+ "mcp.tool.name": "execute",
807+ "mcp.execute.code_length": code.length,
808+ });
769809 debugLog("execute.call", {
770810 elicitationMode: elicitationMode.mode,
771811 elicitationSupport: getElicitationSupport(server),
@@ -807,14 +847,20 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 35 unmodified lines
807847 "mcp.execute.code_length": code.length,
808848 },
809849 }),
850+ (effect) => withJoinKeys(effect, extra),
810851 );
811852
812853 const resumeExecution = (
813854 executionId: string,
814855 action: "accept" | "decline" | "cancel",
815856 content: Record<string, unknown> | undefined,
857+ extra: McpRequestJoinKeys,
816858 ): Effect.Effect<McpToolResult, E> =>
817859 Effect.gen(function* () {
860+ yield* startMarker("mcp.host.tool.resume.start", {
861+ "mcp.tool.name": "resume",
862+ "mcp.execute.execution_id": executionId,
863+ });
818864 debugLog("resume.call", {
819865 executionId,
820866 action,
@@ -855,6 +901,7 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 34 unmodified lines
855901 "mcp.execute.execution_id": executionId,
856902 },
857903 }),
904+ (effect) => withJoinKeys(effect, extra),
858905 );
859906
860907 const requireUserResumeApproval = (executionId: string): Effect.Effect<McpToolResult> =>
@@ -898,8 +945,15 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 37 unmodified lines
898945 );
899946 };
900947
901− const resumeAfterBrowserApproval = (executionId: string): Effect.Effect<McpToolResult, E> =>
948+ const resumeAfterBrowserApproval = (
949+ executionId: string,
950+ extra: McpRequestJoinKeys,
951+ ): Effect.Effect<McpToolResult, E> =>
902952 Effect.gen(function* () {
953+ yield* startMarker("mcp.host.tool.resume.start", {
954+ "mcp.tool.name": "resume",
955+ "mcp.execute.execution_id": executionId,
956+ });
903957 const response = yield* waitForBrowserApprovalResponse(executionId);
904958 if (!response) return yield* requireUserResumeApproval(executionId);
905959
@@ -926,6 +980,7 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 20 unmodified lines
926980 "mcp.execute.execution_id": executionId,
927981 },
928982 }),
983+ (effect) => withJoinKeys(effect, extra),
929984 );
930985
931986 // --- tools ---
@@ -937,7 +992,7 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 5 unmodified lines
937992 description,
938993 inputSchema: { code: z.string().trim().min(1) },
939994 },
940− ({ code }) => runToolEffect(executeCode(code)),
995+ ({ code }, extra) => runToolEffect(executeCode(code, extra)),
941996 ),
942997 ).pipe(
943998 Effect.withSpan("mcp.host.register_tool", {
@@ -993,8 +1048,10 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 49 unmodified lines
9931048 .default("{}"),
9941049 },
9951050 },
996− ({ executionId, action, content: rawContent }) =>
997− runToolEffect(resumeExecution(executionId, action, parseJsonContent(rawContent))),
1051+ ({ executionId, action, content: rawContent }, extra) =>
1052+ runToolEffect(
1053+ resumeExecution(executionId, action, parseJsonContent(rawContent), extra),
1054+ ),
9981055 );
9991056 }
10001057
@@ -1010,7 +1067,7 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 9 unmodified lines
10101067 executionId: z.string().describe("The execution ID from the paused result"),
10111068 },
10121069 },
1013− ({ executionId }) => runToolEffect(resumeAfterBrowserApproval(executionId)),
1070+ ({ executionId }, extra) => runToolEffect(resumeAfterBrowserApproval(executionId, extra)),
10141071 );
10151072 }).pipe(
10161073 Effect.withSpan("mcp.host.register_tool", {