From d94eaa96af2e59382848d12153dd4d2c2cd11782 Mon Sep 17 00:00:00 2001 From: galuis116 Date: Sun, 19 Jul 2026 08:01:17 -0400 Subject: [PATCH] fix(mcp): defer tool-call telemetry via waitUntil instead of firing it forgotten recordMcpToolTelemetry was invoked synchronously and un-awaited right before handleMcpRequest returned its response. On Cloudflare Workers, any in-flight I/O not registered via waitUntil can be torn down once the request context finishes, silently dropping the PostHog capture before its network request lands. Both call sites now pass the telemetry promise to the execution context's waitUntil, matching the scheduleAfterResponse pattern already used in src/orb/webhook.ts. recordMcpToolCall now awaits client.flush() so waitUntil has real async work to hold the context open for. Closes #7233 --- src/mcp/server.ts | 15 ++++--- src/mcp/telemetry.ts | 16 +++++-- test/unit/mcp-server-telemetry.test.ts | 51 ++++++++++++++++++++++ test/unit/mcp-telemetry.test.ts | 58 +++++++++++++++++--------- 4 files changed, 111 insertions(+), 29 deletions(-) diff --git a/src/mcp/server.ts b/src/mcp/server.ts index c493366092..e5b1237821 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -1654,10 +1654,11 @@ export async function handleMcpRequest(c: AppContext): Promise { const usageMetadata = await describeMcpUsageRequest(c.req.raw, telemetry.metadata); const startedAt = Date.now(); const server = new LoopoverMcp(c.env, identity).createServer(); + const executionCtx = getExecutionContext(c); try { - const response = await createMcpHandler(server, { route: "/mcp", enableJsonResponse: true })(c.req.raw, c.env, getExecutionContext(c)); + const response = await createMcpHandler(server, { route: "/mcp", enableJsonResponse: true })(c.req.raw, c.env, executionCtx); if (typeof usageMetadata.toolName === "string") { - recordMcpToolTelemetry(c.env, usageMetadata.toolName, response.status < 400, Date.now() - startedAt); + executionCtx.waitUntil(recordMcpToolTelemetry(c.env, usageMetadata.toolName, response.status < 400, Date.now() - startedAt)); } await recordProductUsageEvent(c.env, { surface: "mcp", @@ -1675,7 +1676,7 @@ export async function handleMcpRequest(c: AppContext): Promise { return response; } catch (error) { if (typeof usageMetadata.toolName === "string") { - recordMcpToolTelemetry(c.env, usageMetadata.toolName, false, Date.now() - startedAt); + executionCtx.waitUntil(recordMcpToolTelemetry(c.env, usageMetadata.toolName, false, Date.now() - startedAt)); } await recordProductUsageEvent(c.env, { surface: "mcp", @@ -1697,10 +1698,12 @@ export async function handleMcpRequest(c: AppContext): Promise { // Single chokepoint for the #6228 PostHog tool-call telemetry (#6237): every `tools/call` request that // reaches handleMcpRequest routes through here exactly once, whether it succeeds or throws. Pure // observability -- never lets a telemetry failure reach the caller, matching recordMcpToolCall's own -// no-op guarantee (#6235) with a second, defensive layer at the actual call site. -function recordMcpToolTelemetry(env: Env, tool: string, ok: boolean, durationMs: number): void { +// no-op guarantee (#6235) with a second, defensive layer at the actual call site. Called sites pass the +// returned promise to `waitUntil` (#7233) rather than awaiting it inline, so a slow PostHog flush can't +// delay the MCP tool response. +async function recordMcpToolTelemetry(env: Env, tool: string, ok: boolean, durationMs: number): Promise { try { - recordMcpToolCall(env, { tool, callerType: "remote", ok, durationMs }); + await recordMcpToolCall(env, { tool, callerType: "remote", ok, durationMs }); } catch { // Telemetry must never affect the tool response (#6237). } diff --git a/src/mcp/telemetry.ts b/src/mcp/telemetry.ts index 33cf1eafd4..12dffa0ec8 100644 --- a/src/mcp/telemetry.ts +++ b/src/mcp/telemetry.ts @@ -49,8 +49,15 @@ export interface McpToolCallEvent { export type McpTelemetryEnv = Pick; /** Record a single MCP tool call to PostHog. Safe no-op when telemetry is unconfigured (no POSTHOG_API_KEY), - * and never throws — a PostHog init/capture failure degrades to recording nothing (#6235). */ -export function recordMcpToolCall(env: McpTelemetryEnv, event: McpToolCallEvent): void { + * and never throws — a PostHog init/capture failure degrades to recording nothing (#6235). + * + * Returns a promise that resolves once the event has actually been flushed to PostHog (#7233) — the + * `flushAt: 1, flushInterval: 0` config makes `capture()` queue the network request immediately, but + * `capture()` itself is fire-and-forget and returns before that request lands. Awaiting `client.flush()` + * gives the caller (src/mcp/server.ts's `recordMcpToolTelemetry`, deferred via Cloudflare's `waitUntil`) + * real async work to hold the execution context open for, so the event isn't silently dropped if the + * Workers runtime tears the request down before an un-awaited background fetch completes. */ +export async function recordMcpToolCall(env: McpTelemetryEnv, event: McpToolCallEvent): Promise { const apiKey = trimmedOrUndefined(env.POSTHOG_API_KEY); // Unconfigured ⇒ record nothing, byte-identical to before this module existed. if (!apiKey) return; @@ -71,9 +78,10 @@ export function recordMcpToolCall(env: McpTelemetryEnv, event: McpToolCallEvent) // No IP-based geo enrichment: the event is anonymous fleet telemetry, not a user location. disableGeoip: true, }); + await client.flush(); } catch { - // Telemetry is best-effort and MUST NOT throw into the MCP tool caller (#6235): a PostHog init/capture - // failure degrades to recording nothing, identical to the unconfigured path above. + // Telemetry is best-effort and MUST NOT throw into the MCP tool caller (#6235): a PostHog init/capture/ + // flush failure degrades to recording nothing, identical to the unconfigured path above. } } diff --git a/test/unit/mcp-server-telemetry.test.ts b/test/unit/mcp-server-telemetry.test.ts index db5f8b8b2b..449ad65cc0 100644 --- a/test/unit/mcp-server-telemetry.test.ts +++ b/test/unit/mcp-server-telemetry.test.ts @@ -187,6 +187,57 @@ describe("MCP server telemetry", () => { ); }); + it("defers tool-call telemetry via executionCtx.waitUntil instead of firing it synchronously and forgetting it (#7233)", async () => { + vi.resetModules(); + vi.doMock("agents/mcp", () => ({ + createMcpHandler: () => () => Response.json({ ok: true }), + })); + let flushRecordMcpToolCall: (() => void) | undefined; + const recordMcpToolCallStarted = new Promise((resolve) => { + flushRecordMcpToolCall = resolve; + }); + const recordMcpToolCall = vi.fn(async () => { + flushRecordMcpToolCall?.(); + }); + vi.doMock("../../src/mcp/telemetry", () => ({ recordMcpToolCall })); + const { handleMcpRequest } = await import("../../src/mcp/server"); + const env = createTestEnv({ PRODUCT_USAGE_HASH_SALT: "mcp-wait-until-telemetry-salt" }); + const request = new Request("https://api.test/mcp", { + method: "POST", + headers: { + authorization: `Bearer ${env.LOOPOVER_MCP_TOKEN}`, + "content-type": "application/json", + }, + body: JSON.stringify({ jsonrpc: "2.0", id: "tool-call", method: "tools/call", params: { name: "loopover_local_status" } }), + }); + + const waitUntilTasks: Promise[] = []; + await expect( + handleMcpRequest({ + env, + executionCtx: { + waitUntil(task: Promise) { + waitUntilTasks.push(task); + }, + passThroughOnException() {}, + }, + req: { + method: "POST", + raw: request, + header: (name: string) => request.headers.get(name) ?? undefined, + }, + json: (body: unknown, status?: number) => Response.json(body, status === undefined ? undefined : { status }), + } as never), + ).resolves.toMatchObject({ status: 200 }); + + // The telemetry promise was handed to waitUntil (not just fired-and-forgotten inline) before the response + // returned, and awaiting it resolves cleanly once the deferred recordMcpToolCall actually runs. + expect(waitUntilTasks).toHaveLength(1); + await recordMcpToolCallStarted; + await expect(waitUntilTasks[0]).resolves.toBeUndefined(); + expect(recordMcpToolCall).toHaveBeenCalledTimes(1); + }); + it("does not let a throwing recordMcpToolCall affect the tool response (#6237)", async () => { vi.resetModules(); vi.doMock("agents/mcp", () => ({ diff --git a/test/unit/mcp-telemetry.test.ts b/test/unit/mcp-telemetry.test.ts index 8a4446ebe1..793cc116f7 100644 --- a/test/unit/mcp-telemetry.test.ts +++ b/test/unit/mcp-telemetry.test.ts @@ -1,11 +1,13 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; -// Mock the PostHog Node SDK so nothing hits the network: the class records every constructor + capture call -// on hoisted spies, and per-test flags let us force an init/capture failure to exercise the never-throw path. +// Mock the PostHog Node SDK so nothing hits the network: the class records every constructor + capture + +// flush call on hoisted spies, and per-test flags let us force an init/capture/flush failure to exercise +// the never-throw path. const h = vi.hoisted(() => ({ constructSpy: vi.fn(), captureSpy: vi.fn(), - state: { throwOnConstruct: false, throwOnCapture: false }, + flushSpy: vi.fn(), + state: { throwOnConstruct: false, throwOnCapture: false, throwOnFlush: false }, })); vi.mock("posthog-node", () => ({ @@ -18,6 +20,10 @@ vi.mock("posthog-node", () => ({ h.captureSpy(message); if (h.state.throwOnCapture) throw new Error("posthog capture failed"); } + async flush(): Promise { + h.flushSpy(); + if (h.state.throwOnFlush) throw new Error("posthog flush failed"); + } }, })); @@ -29,24 +35,27 @@ describe("recordMcpToolCall", () => { beforeEach(() => { h.constructSpy.mockClear(); h.captureSpy.mockClear(); + h.flushSpy.mockClear(); h.state.throwOnConstruct = false; h.state.throwOnCapture = false; + h.state.throwOnFlush = false; }); - it("is a safe no-op when POSTHOG_API_KEY is unset (unconfigured deployment)", () => { - recordMcpToolCall({}, EVENT); + it("is a safe no-op when POSTHOG_API_KEY is unset (unconfigured deployment)", async () => { + await recordMcpToolCall({}, EVENT); expect(h.constructSpy).not.toHaveBeenCalled(); expect(h.captureSpy).not.toHaveBeenCalled(); + expect(h.flushSpy).not.toHaveBeenCalled(); }); - it("treats a blank/whitespace API key as unconfigured", () => { - recordMcpToolCall({ POSTHOG_API_KEY: " " }, EVENT); + it("treats a blank/whitespace API key as unconfigured", async () => { + await recordMcpToolCall({ POSTHOG_API_KEY: " " }, EVENT); expect(h.constructSpy).not.toHaveBeenCalled(); expect(h.captureSpy).not.toHaveBeenCalled(); }); - it("captures exactly the allowlisted fields against the US-cloud default host when configured", () => { - recordMcpToolCall({ POSTHOG_API_KEY: "phc_test" }, EVENT); + it("captures exactly the allowlisted fields against the US-cloud default host when configured", async () => { + await recordMcpToolCall({ POSTHOG_API_KEY: "phc_test" }, EVENT); expect(h.constructSpy).toHaveBeenCalledTimes(1); expect(h.constructSpy).toHaveBeenCalledWith("phc_test", { @@ -73,10 +82,12 @@ describe("recordMcpToolCall", () => { }); // The allowlist is the whole payload — no argument/source/wallet/hotkey/trust-score field can ride along. expect(Object.keys(message.properties).sort()).toEqual(["caller_type", "duration_ms", "ok", "tool"]); + // #7233: the event is actually flushed, not just queued, before recordMcpToolCall's promise resolves. + expect(h.flushSpy).toHaveBeenCalledTimes(1); }); - it("honors a POSTHOG_HOST override and carries a local caller / failed call verbatim", () => { - recordMcpToolCall( + it("honors a POSTHOG_HOST override and carries a local caller / failed call verbatim", async () => { + await recordMcpToolCall( { POSTHOG_API_KEY: "phc_test", POSTHOG_HOST: "https://eu.i.posthog.com" }, { tool: "check_slop_risk", callerType: "local", ok: false, durationMs: 0 }, ); @@ -95,8 +106,8 @@ describe("recordMcpToolCall", () => { }); }); - it("trims surrounding whitespace from the API key and host", () => { - recordMcpToolCall({ POSTHOG_API_KEY: " phc_test ", POSTHOG_HOST: " https://eu.i.posthog.com " }, EVENT); + it("trims surrounding whitespace from the API key and host", async () => { + await recordMcpToolCall({ POSTHOG_API_KEY: " phc_test ", POSTHOG_HOST: " https://eu.i.posthog.com " }, EVENT); expect(h.constructSpy).toHaveBeenCalledWith("phc_test", { host: "https://eu.i.posthog.com", flushAt: 1, @@ -104,8 +115,8 @@ describe("recordMcpToolCall", () => { }); }); - it("falls back to the default host when POSTHOG_HOST is blank", () => { - recordMcpToolCall({ POSTHOG_API_KEY: "phc_test", POSTHOG_HOST: " " }, EVENT); + it("falls back to the default host when POSTHOG_HOST is blank", async () => { + await recordMcpToolCall({ POSTHOG_API_KEY: "phc_test", POSTHOG_HOST: " " }, EVENT); expect(h.constructSpy).toHaveBeenCalledWith("phc_test", { host: "https://us.i.posthog.com", flushAt: 1, @@ -113,15 +124,24 @@ describe("recordMcpToolCall", () => { }); }); - it("never throws when the PostHog client fails to initialize", () => { + it("never throws when the PostHog client fails to initialize", async () => { h.state.throwOnConstruct = true; - expect(() => recordMcpToolCall({ POSTHOG_API_KEY: "phc_test" }, EVENT)).not.toThrow(); + await expect(recordMcpToolCall({ POSTHOG_API_KEY: "phc_test" }, EVENT)).resolves.toBeUndefined(); expect(h.captureSpy).not.toHaveBeenCalled(); }); - it("never throws when capture itself fails", () => { + it("never throws when capture itself fails", async () => { h.state.throwOnCapture = true; - expect(() => recordMcpToolCall({ POSTHOG_API_KEY: "phc_test" }, EVENT)).not.toThrow(); + await expect(recordMcpToolCall({ POSTHOG_API_KEY: "phc_test" }, EVENT)).resolves.toBeUndefined(); + expect(h.captureSpy).toHaveBeenCalledTimes(1); + // capture() threw, so flush() is never reached — same catch branch as the constructor failure above. + expect(h.flushSpy).not.toHaveBeenCalled(); + }); + + it("never throws when flush itself fails (#7233) — the event was captured/queued regardless", async () => { + h.state.throwOnFlush = true; + await expect(recordMcpToolCall({ POSTHOG_API_KEY: "phc_test" }, EVENT)).resolves.toBeUndefined(); expect(h.captureSpy).toHaveBeenCalledTimes(1); + expect(h.flushSpy).toHaveBeenCalledTimes(1); }); });