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 @@
1import { describe, expect, it } from "@effect/vitest";
2import { Effect } from "effect";
3import type * as Tracer from "effect/Tracer";
4
5import { 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.
9const 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
47const 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
54const annotate = (request: Request) =>
55 annotateMcpRequest(request, { token: null, parseBody: true });
56
57describe("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
105const NamedParams = Schema.Struct({ name: Schema.optional(Schema.String) });105const NamedParams = Schema.Struct({ name: Schema.optional(Schema.String) });
106const UriParams = Schema.Struct({ uri: Schema.optional(Schema.String) });106const UriParams = Schema.Struct({ uri: Schema.optional(Schema.String) });
107 107
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.
113const CancelledParams = Schema.Struct({
114 requestId: Schema.optional(Schema.Union([Schema.String, Schema.Number])),
115 reason: Schema.optional(Schema.String),
116});
117
108const decodeJsonRpcEnvelopeString = Schema.decodeUnknownOption(118const decodeJsonRpcEnvelopeString = Schema.decodeUnknownOption(
109 Schema.fromJsonString(JsonRpcEnvelope),119 Schema.fromJsonString(JsonRpcEnvelope),
110);120);
111const decodeInitializeParams = Schema.decodeUnknownOption(InitializeParams);121const decodeInitializeParams = Schema.decodeUnknownOption(InitializeParams);
112const decodeNamedParams = Schema.decodeUnknownOption(NamedParams);122const decodeNamedParams = Schema.decodeUnknownOption(NamedParams);
113const decodeUriParams = Schema.decodeUnknownOption(UriParams);123const decodeUriParams = Schema.decodeUnknownOption(UriParams);
124const decodeCancelledParams = Schema.decodeUnknownOption(CancelledParams);
114const decodeElicitationReplyResult = Schema.decodeUnknownOption(ElicitationReplyResult);125const decodeElicitationReplyResult = Schema.decodeUnknownOption(ElicitationReplyResult);
115 126
116const readJsonRpcEnvelope = (request: Request): Effect.Effect<Option.Option<JsonRpcEnvelope>> =>127const readJsonRpcEnvelope = (request: Request): Effect.Effect<Option.Option<JsonRpcEnvelope>> =>
@@ -158,6 +169,15 @@ const methodAttrs = (envelope: JsonRpcEnvelope): Record<string, unknown> => { · 41 unmodified lines
158 onSome: ({ name }) => (name ? { "mcp.prompt.name": name } : {}),169 onSome: ({ name }) => (name ? { "mcp.prompt.name": name } : {}),
159 }),170 }),
160 ),171 ),
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 ),
161 Match.option,181 Match.option,
162 Option.getOrElse(() => ({}) as Record<string, unknown>),182 Option.getOrElse(() => ({}) as Record<string, unknown>),
163 );183 );
packages/hosts/mcp/src/tool-server.test.ts+113−0
@@ -1,5 +1,6 @@
1import { describe, expect, it } from "@effect/vitest";1import { describe, expect, it } from "@effect/vitest";
2import { Data, Deferred, Effect } from "effect";2import { Data, Deferred, Effect } from "effect";
3import type * as Tracer from "effect/Tracer";
3import { Client } from "@modelcontextprotocol/sdk/client/index.js";4import { Client } from "@modelcontextprotocol/sdk/client/index.js";
4import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";5import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
5import { ElicitRequestSchema } from "@modelcontextprotocol/sdk/types.js";6import { ElicitRequestSchema } from "@modelcontextprotocol/sdk/types.js";
@@ -88,6 +89,72 @@ const withNativeClient = async <E extends Cause.YieldableError>( · 82 unmodified lines
88 fn: (client: Client) => Promise<void>,89 fn: (client: Client) => Promise<void>,
89) => withClient(engine, capabilities, fn, { elicitationMode: { mode: "native" } });90) => withClient(engine, capabilities, fn, { elicitationMode: { mode: "native" } });
90 91
92type 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. */
99const 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. */
137const 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
91const ELICITATION_CAPS: ClientCapabilities = {158const ELICITATION_CAPS: ClientCapabilities = {
92 elicitation: { form: {}, url: {} },159 elicitation: { form: {}, url: {} },
93};160};
@@ -1758,3 +1825,49 @@ describe("MCP host server — skills tool", () => { · 1664 unmodified lines
1758 });1825 });
1759 });1826 });
1760});1827});
1828
1829describe("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
647 return index === -1 ? "" : description.slice(index).trimEnd();647 return index === -1 ? "" : description.slice(index).trimEnd();
648};648};
649 649
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
665type McpRequestJoinKeys = {
666 readonly requestId: string | number;
667 readonly sessionId?: string | undefined;
668};
669
670const joinKeyAttributes = (extra: McpRequestJoinKeys): Record<string, unknown> => ({
671 "mcp.rpc.id": String(extra.requestId),
672 ...(extra.sessionId ? { "mcp.request.session_id": extra.sessionId } : {}),
673});
674
675const withJoinKeys = <A, EffE>(
676 effect: Effect.Effect<A, EffE>,
677 extra: McpRequestJoinKeys,
678): Effect.Effect<A, EffE> => Effect.annotateSpans(effect, joinKeyAttributes(extra));
679
680const startMarker = (name: string, attributes: Record<string, unknown>): Effect.Effect<void> =>
681 Effect.void.pipe(Effect.withSpan(name, { attributes }));
682
650const JsonObjectFromString = Schema.fromJsonString(Schema.Record(Schema.String, Schema.Unknown));683const JsonObjectFromString = Schema.fromJsonString(Schema.Record(Schema.String, Schema.Unknown));
651const decodeJsonObjectString = Schema.decodeUnknownOption(JsonObjectFromString);684const decodeJsonObjectString = Schema.decodeUnknownOption(JsonObjectFromString);
652 685
@@ -764,8 +797,15 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 111 unmodified lines
764 ),797 ),
765 ).pipe(Effect.withSpan("mcp.host.create_server"));798 ).pipe(Effect.withSpan("mcp.host.create_server"));
766 799
767 const executeCode = (code: string): Effect.Effect<McpToolResult, E> =>800 const executeCode = (
801 code: string,
802 extra: McpRequestJoinKeys,
803 ): Effect.Effect<McpToolResult, E> =>
768 Effect.gen(function* () {804 Effect.gen(function* () {
805 yield* startMarker("mcp.host.tool.execute.start", {
806 "mcp.tool.name": "execute",
807 "mcp.execute.code_length": code.length,
808 });
769 debugLog("execute.call", {809 debugLog("execute.call", {
770 elicitationMode: elicitationMode.mode,810 elicitationMode: elicitationMode.mode,
771 elicitationSupport: getElicitationSupport(server),811 elicitationSupport: getElicitationSupport(server),
@@ -807,14 +847,20 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 35 unmodified lines
807 "mcp.execute.code_length": code.length,847 "mcp.execute.code_length": code.length,
808 },848 },
809 }),849 }),
850 (effect) => withJoinKeys(effect, extra),
810 );851 );
811 852
812 const resumeExecution = (853 const resumeExecution = (
813 executionId: string,854 executionId: string,
814 action: "accept" | "decline" | "cancel",855 action: "accept" | "decline" | "cancel",
815 content: Record<string, unknown> | undefined,856 content: Record<string, unknown> | undefined,
857 extra: McpRequestJoinKeys,
816 ): Effect.Effect<McpToolResult, E> =>858 ): Effect.Effect<McpToolResult, E> =>
817 Effect.gen(function* () {859 Effect.gen(function* () {
860 yield* startMarker("mcp.host.tool.resume.start", {
861 "mcp.tool.name": "resume",
862 "mcp.execute.execution_id": executionId,
863 });
818 debugLog("resume.call", {864 debugLog("resume.call", {
819 executionId,865 executionId,
820 action,866 action,
@@ -855,6 +901,7 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 34 unmodified lines
855 "mcp.execute.execution_id": executionId,901 "mcp.execute.execution_id": executionId,
856 },902 },
857 }),903 }),
904 (effect) => withJoinKeys(effect, extra),
858 );905 );
859 906
860 const requireUserResumeApproval = (executionId: string): Effect.Effect<McpToolResult> =>907 const requireUserResumeApproval = (executionId: string): Effect.Effect<McpToolResult> =>
@@ -898,8 +945,15 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 37 unmodified lines
898 );945 );
899 };946 };
900 947
901 const resumeAfterBrowserApproval = (executionId: string): Effect.Effect<McpToolResult, E> =>948 const resumeAfterBrowserApproval = (
949 executionId: string,
950 extra: McpRequestJoinKeys,
951 ): Effect.Effect<McpToolResult, E> =>
902 Effect.gen(function* () {952 Effect.gen(function* () {
953 yield* startMarker("mcp.host.tool.resume.start", {
954 "mcp.tool.name": "resume",
955 "mcp.execute.execution_id": executionId,
956 });
903 const response = yield* waitForBrowserApprovalResponse(executionId);957 const response = yield* waitForBrowserApprovalResponse(executionId);
904 if (!response) return yield* requireUserResumeApproval(executionId);958 if (!response) return yield* requireUserResumeApproval(executionId);
905 959
@@ -926,6 +980,7 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 20 unmodified lines
926 "mcp.execute.execution_id": executionId,980 "mcp.execute.execution_id": executionId,
927 },981 },
928 }),982 }),
983 (effect) => withJoinKeys(effect, extra),
929 );984 );
930 985
931 // --- tools ---986 // --- tools ---
@@ -937,7 +992,7 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 5 unmodified lines
937 description,992 description,
938 inputSchema: { code: z.string().trim().min(1) },993 inputSchema: { code: z.string().trim().min(1) },
939 },994 },
940 ({ code }) => runToolEffect(executeCode(code)),995 ({ code }, extra) => runToolEffect(executeCode(code, extra)),
941 ),996 ),
942 ).pipe(997 ).pipe(
943 Effect.withSpan("mcp.host.register_tool", {998 Effect.withSpan("mcp.host.register_tool", {
@@ -993,8 +1048,10 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 49 unmodified lines
993 .default("{}"),1048 .default("{}"),
994 },1049 },
995 },1050 },
996 ({ executionId, action, content: rawContent }) =>1051 ({ executionId, action, content: rawContent }, extra) =>
997 runToolEffect(resumeExecution(executionId, action, parseJsonContent(rawContent))),1052 runToolEffect(
1053 resumeExecution(executionId, action, parseJsonContent(rawContent), extra),
1054 ),
998 );1055 );
999 }1056 }
1000 1057
@@ -1010,7 +1067,7 @@ export const createExecutorMcpServer = <E extends Cause.YieldableError>( · 9 unmodified lines
1010 executionId: z.string().describe("The execution ID from the paused result"),1067 executionId: z.string().describe("The execution ID from the paused result"),
1011 },1068 },
1012 },1069 },
1013 ({ executionId }) => runToolEffect(resumeAfterBrowserApproval(executionId)),1070 ({ executionId }, extra) => runToolEffect(resumeAfterBrowserApproval(executionId, extra)),
1014 );1071 );
1015 }).pipe(1072 }).pipe(
1016 Effect.withSpan("mcp.host.register_tool", {1073 Effect.withSpan("mcp.host.register_tool", {