From f155accf4bbd6e72eb91931e94e8013ff605126e Mon Sep 17 00:00:00 2001 From: Prasanna A P <106952318+Prasanna721@users.noreply.github.com> Date: Wed, 29 Jul 2026 21:07:34 -0700 Subject: [PATCH] track stateless mcp tool usage --- apps/mcp/package.json | 1 + apps/mcp/src/server/analytics.test.ts | 171 ++++++++++++++++++ apps/mcp/src/server/analytics.ts | 216 +++++++++++++++++++++++ apps/mcp/src/server/client/index.test.ts | 1 + apps/mcp/src/server/client/index.ts | 13 +- apps/mcp/src/server/index.ts | 5 +- apps/mcp/src/server/server.ts | 14 +- apps/mcp/src/server/tools/types.ts | 2 +- apps/mcp/src/server/types.ts | 2 + bun.lock | 3 + 10 files changed, 420 insertions(+), 8 deletions(-) create mode 100644 apps/mcp/src/server/analytics.test.ts create mode 100644 apps/mcp/src/server/analytics.ts diff --git a/apps/mcp/package.json b/apps/mcp/package.json index 498e14ad..2be89db5 100644 --- a/apps/mcp/package.json +++ b/apps/mcp/package.json @@ -35,6 +35,7 @@ "clsx": "^2.1.1", "hono": "^4.11.1", "jose": "^6.2.0", + "posthog-node": "^5.18.0", "react": "^19.2.4", "react-dom": "^19.2.4", "supermemory": "^4.0.0", diff --git a/apps/mcp/src/server/analytics.test.ts b/apps/mcp/src/server/analytics.test.ts new file mode 100644 index 00000000..3e449455 --- /dev/null +++ b/apps/mcp/src/server/analytics.test.ts @@ -0,0 +1,171 @@ +import type { McpServer, ServerContext } from "@modelcontextprotocol/server" +import { describe, expect, it, vi } from "vitest" +import { z } from "zod" +import { + createTrackedToolServer, + posthogEventForToolExecution, + type McpToolAnalytics, +} from "./analytics" + +function testServer() { + let callback: ((...args: unknown[]) => unknown) | undefined + const registerTool = vi.fn( + ( + _name: string, + _config: unknown, + handler: (...args: unknown[]) => unknown, + ) => { + callback = handler + return {} + }, + ) + + return { + server: { registerTool } as unknown as McpServer, + invoke(...args: unknown[]) { + if (!callback) throw new Error("Tool was not registered") + return callback(...args) + }, + } +} + +const context = { + mcpReq: { envelope: {} }, +} as unknown as ServerContext + +describe("MCP tool analytics", () => { + it("records sanitized completion metadata without tool content", async () => { + const harness = testServer() + const record = vi.fn() + const analytics: McpToolAnalytics = { record } + const server = createTrackedToolServer(harness.server, analytics, () => ({ + name: "claude", + version: "1.2.3", + })) + + server.registerTool( + "search_memory", + { + inputSchema: z.object({ query: z.string(), containerTag: z.string() }), + }, + async () => ({ + content: [{ type: "text" as const, text: "secret result" }], + }), + ) + + await harness.invoke( + { query: "private query", containerTag: "private-workspace" }, + context, + ) + + expect(record).toHaveBeenCalledOnce() + expect(record).toHaveBeenCalledWith( + expect.objectContaining({ + toolName: "search_memory", + surface: "model_tool", + outcome: "success", + workspaceExplicit: true, + client: { name: "claude", version: "1.2.3" }, + }), + ) + expect(JSON.stringify(record.mock.calls[0])).not.toContain("private query") + expect(JSON.stringify(record.mock.calls[0])).not.toContain( + "private-workspace", + ) + expect(JSON.stringify(record.mock.calls[0])).not.toContain("secret result") + }) + + it("treats returned MCP errors as failed executions", async () => { + const harness = testServer() + const record = vi.fn() + const server = createTrackedToolServer( + harness.server, + { record }, + () => null, + ) + + server.registerTool( + "save-memory", + { inputSchema: z.object({}) }, + async () => ({ + content: [{ type: "text" as const, text: "failed" }], + isError: true, + }), + ) + + await harness.invoke({}, context) + + expect(record).toHaveBeenCalledWith( + expect.objectContaining({ + surface: "app_action", + outcome: "error", + errorType: "tool_result", + }), + ) + }) + + it("records thrown error categories and preserves the rejection", async () => { + const harness = testServer() + const record = vi.fn() + const server = createTrackedToolServer( + harness.server, + { record }, + () => null, + ) + + server.registerTool( + "fetch-graph-data", + { inputSchema: z.object({}) }, + async () => { + throw new TypeError("sensitive failure") + }, + ) + + await expect(harness.invoke({}, context)).rejects.toThrow( + "sensitive failure", + ) + expect(record).toHaveBeenCalledWith( + expect.objectContaining({ + surface: "app_internal", + outcome: "error", + errorType: "TypeError", + }), + ) + expect(JSON.stringify(record.mock.calls[0])).not.toContain( + "sensitive failure", + ) + }) + + it("uses the existing user identity and company group", () => { + const event = posthogEventForToolExecution( + { + userId: "user_123", + organizationId: "org_123", + oauthClientId: "client_123", + }, + { + toolName: "guided-save", + surface: "app_launcher", + outcome: "success", + durationMs: 42, + workspaceExplicit: false, + }, + ) + + expect(event).toEqual({ + distinctId: "user_123", + event: "mcp_tool_executed", + groups: { company: "org_123" }, + properties: { + app: "mcp", + tool_name: "guided-save", + outcome: "success", + duration_ms: 42, + mcp_runtime: "stateless", + mcp_surface: "app_launcher", + workspace_explicit: false, + oauth_client_id: "client_123", + }, + }) + }) +}) diff --git a/apps/mcp/src/server/analytics.ts b/apps/mcp/src/server/analytics.ts new file mode 100644 index 00000000..5f1b87e2 --- /dev/null +++ b/apps/mcp/src/server/analytics.ts @@ -0,0 +1,216 @@ +import type { McpServer, ServerContext } from "@modelcontextprotocol/server" +import { PostHog } from "posthog-node" +import type { ActorContext, ServerEnv } from "./types" + +const DEFAULT_POSTHOG_HOST = "https://us.i.posthog.com" + +export type McpToolSurface = + | "model_tool" + | "app_launcher" + | "app_action" + | "app_internal" + +export type McpToolOutcome = "success" | "error" + +export interface McpToolExecution { + toolName: string + surface: McpToolSurface + outcome: McpToolOutcome + durationMs: number + workspaceExplicit: boolean + client?: { name: string; version?: string } + errorType?: string +} + +export interface McpToolAnalytics { + record(execution: McpToolExecution): void +} + +export type WaitUntil = (promise: Promise) => void + +type ClientInfoResolver = ( + context: ServerContext, +) => { name: string; version?: string } | null + +const TOOL_SURFACES: Record = { + search_memory: "model_tool", + listDocuments: "model_tool", + getDocument: "model_tool", + listMemories: "model_tool", + listSpaces: "model_tool", + whoAmI: "model_tool", + add_memory: "model_tool", + "select-workspace": "app_launcher", + "memory-graph": "app_launcher", + "guided-save": "app_launcher", + "upload-file": "app_launcher", + "set-active-tag": "app_action", + "save-memory": "app_action", + "upload-file-submit": "app_action", + "fetch-graph-data": "app_internal", +} + +let posthogConfig: + | { + apiKey: string + host: string + client: PostHog + } + | undefined + +function posthogClient(apiKey: string, host: string): PostHog { + if (posthogConfig?.apiKey === apiKey && posthogConfig.host === host) { + return posthogConfig.client + } + + const client = new PostHog(apiKey, { + host, + flushAt: 1, + flushInterval: 0, + }) + posthogConfig = { apiKey, host, client } + return client +} + +export function posthogEventForToolExecution( + actor: Pick, + execution: McpToolExecution, +) { + return { + distinctId: actor.userId, + event: "mcp_tool_executed", + groups: { company: actor.organizationId }, + properties: { + app: "mcp", + tool_name: execution.toolName, + outcome: execution.outcome, + duration_ms: execution.durationMs, + mcp_runtime: "stateless", + mcp_surface: execution.surface, + workspace_explicit: execution.workspaceExplicit, + ...(execution.client + ? { + mcp_client_name: execution.client.name, + ...(execution.client.version + ? { mcp_client_version: execution.client.version } + : {}), + } + : {}), + ...(actor.oauthClientId ? { oauth_client_id: actor.oauthClientId } : {}), + ...(execution.errorType ? { error_type: execution.errorType } : {}), + }, + } +} + +export function createPosthogAnalytics( + env: ServerEnv, + actor: ActorContext, + waitUntil: WaitUntil, +): McpToolAnalytics { + const apiKey = env.POSTHOG_API_KEY + if (!apiKey) return { record: () => undefined } + + const client = posthogClient(apiKey, env.POSTHOG_HOST || DEFAULT_POSTHOG_HOST) + + return { + record(execution) { + try { + const capture = client + .captureImmediate(posthogEventForToolExecution(actor, execution)) + .catch((error) => console.error("PostHog MCP tracking error:", error)) + waitUntil(capture) + } catch (error) { + console.error("PostHog MCP tracking error:", error) + } + }, + } +} + +function workspaceWasExplicit(value: unknown): boolean { + if (!value || typeof value !== "object") return false + const containerTag = Reflect.get(value, "containerTag") + return typeof containerTag === "string" && containerTag.trim().length > 0 +} + +function isErrorResult(value: unknown): boolean { + return ( + !!value && + typeof value === "object" && + Reflect.get(value, "isError") === true + ) +} + +function thrownErrorType(error: unknown): string { + if (error instanceof Error && error.name) return error.name + if (error && typeof error === "object") { + const status = Reflect.get(error, "status") + if (typeof status === "number") return `http_${status}` + } + return "unknown" +} + +function safeRecord(analytics: McpToolAnalytics, execution: McpToolExecution) { + try { + analytics.record(execution) + } catch (error) { + console.error("MCP analytics recording error:", error) + } +} + +export function createTrackedToolServer( + server: McpServer, + analytics: McpToolAnalytics, + getClientInfo: ClientInfoResolver, +): Pick { + const registerTool = (( + name: string, + config: unknown, + handler: (...args: unknown[]) => unknown, + ) => { + const trackedHandler = async (...callbackArgs: unknown[]) => { + const startedAt = performance.now() + const input = callbackArgs.length > 1 ? callbackArgs[0] : undefined + const context = callbackArgs.at(-1) as ServerContext + + const finish = (outcome: McpToolOutcome, errorType?: string) => { + let client: ReturnType = null + try { + client = getClientInfo(context) + } catch { + // Client metadata is optional and must never affect a tool call. + } + + safeRecord(analytics, { + toolName: name, + surface: TOOL_SURFACES[name] ?? "model_tool", + outcome, + durationMs: Math.max(0, Math.round(performance.now() - startedAt)), + workspaceExplicit: workspaceWasExplicit(input), + ...(client ? { client } : {}), + ...(errorType ? { errorType } : {}), + }) + } + + try { + const result = await handler(...callbackArgs) + if (isErrorResult(result)) { + finish("error", "tool_result") + } else { + finish("success") + } + return result + } catch (error) { + finish("error", thrownErrorType(error)) + throw error + } + } + + return Reflect.apply(server.registerTool, server, [ + name, + config, + trackedHandler, + ]) + }) as McpServer["registerTool"] + + return { registerTool } +} diff --git a/apps/mcp/src/server/client/index.test.ts b/apps/mcp/src/server/client/index.test.ts index c7088e1c..2a84f605 100644 --- a/apps/mcp/src/server/client/index.test.ts +++ b/apps/mcp/src/server/client/index.test.ts @@ -50,6 +50,7 @@ describe("SupermemoryClient memory listing", () => { expect(init.headers).toMatchObject({ Authorization: "Bearer oauth-token", "Content-Type": "application/json", + "x-sm-source": "supermemory-mcp", }) expect(JSON.parse(init.body as string)).toEqual({ containerTags: ["snowcone_grande"], diff --git a/apps/mcp/src/server/client/index.ts b/apps/mcp/src/server/client/index.ts index e66d53dc..e23fe5cc 100644 --- a/apps/mcp/src/server/client/index.ts +++ b/apps/mcp/src/server/client/index.ts @@ -13,6 +13,7 @@ import type { const MAX_CHARS = 200000 export const DEFAULT_PROJECT_ID = "sm_project_default" const FETCH_TIMEOUT_MS = 30_000 +const MCP_SOURCE = "supermemory-mcp" export type { ContainerTag, @@ -134,6 +135,7 @@ export class SupermemoryClient { apiKey: bearerToken, baseURL: apiUrl, timeout: FETCH_TIMEOUT_MS, + defaultHeaders: { "x-sm-source": MCP_SOURCE }, }) this.hasExplicitContainerTag = Boolean(containerTag) this.containerTag = containerTag || DEFAULT_PROJECT_ID @@ -146,7 +148,7 @@ export class SupermemoryClient { const result = await this.client.add({ content, containerTag: this.containerTag, - metadata: { sm_source: "supermemory-mcp" }, + metadata: { sm_source: MCP_SOURCE }, }) return { id: result.id, @@ -319,6 +321,7 @@ export class SupermemoryClient { headers: { Authorization: `Bearer ${this.bearerToken}`, "Content-Type": "application/json", + "x-sm-source": MCP_SOURCE, }, signal, }) @@ -352,6 +355,7 @@ export class SupermemoryClient { headers: { Authorization: `Bearer ${this.bearerToken}`, "Content-Type": "application/json", + "x-sm-source": MCP_SOURCE, }, body: JSON.stringify({ page, @@ -411,6 +415,7 @@ export class SupermemoryClient { headers: { Authorization: `Bearer ${this.bearerToken}`, "Content-Type": "application/json", + "x-sm-source": MCP_SOURCE, }, body: JSON.stringify({ containerTags: [this.containerTag], @@ -449,15 +454,13 @@ export class SupermemoryClient { if (containerTag) { formData.append("containerTags", containerTag) } - formData.append( - "metadata", - JSON.stringify({ sm_source: "supermemory-mcp" }), - ) + formData.append("metadata", JSON.stringify({ sm_source: MCP_SOURCE })) const response = await fetch(`${this.apiUrl}/v3/documents/file`, { method: "POST", headers: { Authorization: `Bearer ${this.bearerToken}`, + "x-sm-source": MCP_SOURCE, }, body: formData, }) diff --git a/apps/mcp/src/server/index.ts b/apps/mcp/src/server/index.ts index 9247a745..69ac6870 100644 --- a/apps/mcp/src/server/index.ts +++ b/apps/mcp/src/server/index.ts @@ -183,7 +183,10 @@ async function handleMcpRequest( ? new Request(new URL(rewritePath, c.req.url).toString(), c.req.raw) : c.req.raw const handler = createMcpHandler( - () => createSupermemoryServer(c.env, actor), + () => + createSupermemoryServer(c.env, actor, (promise) => + c.executionCtx.waitUntil(promise), + ), { route: "/mcp", legacy: "stateless", diff --git a/apps/mcp/src/server/server.ts b/apps/mcp/src/server/server.ts index 4a4f4d47..93fa90e6 100644 --- a/apps/mcp/src/server/server.ts +++ b/apps/mcp/src/server/server.ts @@ -3,6 +3,11 @@ import { McpServer, type ServerContext, } from "@modelcontextprotocol/server" +import { + createPosthogAnalytics, + createTrackedToolServer, + type WaitUntil, +} from "./analytics" import { fetchSession } from "./auth" import { SupermemoryClient } from "./client" import { registerContextPrompt } from "./prompts/context" @@ -41,6 +46,7 @@ function clientInfoFromContext(context: ServerContext): ClientInfo | null { export function createSupermemoryServer( env: ServerEnv, actor: ActorContext, + waitUntil: WaitUntil, ): McpServer { const server = new McpServer({ name: "supermemory", @@ -58,9 +64,15 @@ export function createSupermemoryServer( workspaceState.setActiveContainerTag(containerTag) const resolveContainerTag = (explicit?: string) => resolveWorkspaceContainerTag(explicit, getActiveContainerTag) + const analytics = createPosthogAnalytics(env, actor, waitUntil) + const toolServer = createTrackedToolServer( + server, + analytics, + clientInfoFromContext, + ) registerAllTools({ - server, + server: toolServer, actor, getClient, getSession: () => fetchSession(actor.bearerToken, apiUrl), diff --git a/apps/mcp/src/server/tools/types.ts b/apps/mcp/src/server/tools/types.ts index 9974cfe3..3ccb4cb7 100644 --- a/apps/mcp/src/server/tools/types.ts +++ b/apps/mcp/src/server/tools/types.ts @@ -6,7 +6,7 @@ import type { ActorContext } from "../types" // Dependencies passed to every tool's register() function. // Keep this surface small — tools should read this rather than reach into the agent. export interface ToolDeps { - server: McpServer + server: Pick actor: ActorContext getClient: (containerTag?: string) => SupermemoryClient getSession: () => Promise diff --git a/apps/mcp/src/server/types.ts b/apps/mcp/src/server/types.ts index 4698f3dd..f10df537 100644 --- a/apps/mcp/src/server/types.ts +++ b/apps/mcp/src/server/types.ts @@ -12,4 +12,6 @@ export interface ServerEnv { API_URL?: string MCP_RESOURCE?: string ALLOWED_MCP_ORIGIN_HOSTNAMES?: string + POSTHOG_API_KEY?: string + POSTHOG_HOST?: string } diff --git a/bun.lock b/bun.lock index 3b4f8ba7..855f742c 100644 --- a/bun.lock +++ b/bun.lock @@ -102,6 +102,7 @@ "clsx": "^2.1.1", "hono": "^4.11.1", "jose": "^6.2.0", + "posthog-node": "^5.18.0", "react": "^19.2.4", "react-dom": "^19.2.4", "supermemory": "^4.0.0", @@ -4148,6 +4149,8 @@ "posthog-js": ["posthog-js@1.359.0", "", { "dependencies": { "@opentelemetry/api": "^1.9.0", "@opentelemetry/api-logs": "^0.208.0", "@opentelemetry/exporter-logs-otlp-http": "^0.208.0", "@opentelemetry/resources": "^2.2.0", "@opentelemetry/sdk-logs": "^0.208.0", "@posthog/core": "1.23.2", "@posthog/types": "1.359.0", "core-js": "^3.38.1", "dompurify": "^3.3.1", "fflate": "^0.4.8", "preact": "^10.28.2", "query-selector-shadow-dom": "^1.0.1", "web-vitals": "^5.1.0" } }, "sha512-W3yrLKgJfc9qm7Z9e7LblPdCuSJsfB8xzg36GSklyFETqb5b8KtmF2cx/tiXMGEDNmKPPJ+oex/N9Q+vdTjPug=="], + "posthog-node": ["posthog-node@5.28.0", "", { "dependencies": { "@posthog/core": "1.23.2" }, "peerDependencies": { "rxjs": "^7.0.0" }, "optionalPeers": ["rxjs"] }, "sha512-EETYV0zA+7BLQmXzY+vGyDMoQK8uHf8f/1utbRjKncI41gPkw+4piGP7l4UT5Luld+4vQpJPOR1q1YrbXm7XjQ=="], + "powershell-utils": ["powershell-utils@0.1.0", "", {}, "sha512-dM0jVuXJPsDN6DvRpea484tCUaMiXWjuCn++HGTqUWzGDjv5tZkEZldAJ/UMlqRYGFrD/etByo4/xOuC/snX2A=="], "preact": ["preact@10.28.4", "", {}, "sha512-uKFfOHWuSNpRFVTnljsCluEFq57OKT+0QdOiQo8XWnQ/pSvg7OpX5eNOejELXJMWy+BwM2nobz0FkvzmnpCNsQ=="],