diff --git a/apps/server/src/integrations/mcp/mcp-auth.helpers.ts b/apps/server/src/integrations/mcp/mcp-auth.helpers.ts index 80e34d4f..e1709666 100644 --- a/apps/server/src/integrations/mcp/mcp-auth.helpers.ts +++ b/apps/server/src/integrations/mcp/mcp-auth.helpers.ts @@ -149,6 +149,15 @@ export type DocmostMcpConfig = ( has?: (uri: string) => boolean; evict?: (uri: string) => void; }; + // Dependency-neutral metrics sink injected by McpService (mirror of the + // package's onMetric). The package emits generic (name, value, labels) + // samples; McpService maps them onto the prom-client registry. Undefined + // when metrics are disabled → the package no-ops. + onMetric?: ( + name: string, + value: number, + labels?: Record, + ) => void; }; export interface ResolvedMcpAuth { diff --git a/apps/server/src/integrations/mcp/mcp.service.ts b/apps/server/src/integrations/mcp/mcp.service.ts index d6942963..11ae9d64 100644 --- a/apps/server/src/integrations/mcp/mcp.service.ts +++ b/apps/server/src/integrations/mcp/mcp.service.ts @@ -30,6 +30,11 @@ import { ResolvedMcpAuth, } from './mcp-auth.helpers'; import { SandboxStore } from '../sandbox/sandbox.store'; +import { + isMetricsEnabled, + observeMcpTool, + incConnectTimeout, +} from '../metrics/metrics.registry'; // Minimal shape of the embedded MCP HTTP handler exported by @docmost/mcp/http. interface McpHttpHandler { @@ -331,7 +336,28 @@ export class McpService implements OnModuleDestroy { // can store blobs in the shared in-RAM store regardless of which // credential variant resolved. The sink (put/has/evict + uri↔id // mapping) is owned by SandboxStore.asSink(). - return { ...resolved.config, sandbox: this.sandboxStore.asSink() }; + // Route the package's dependency-neutral metric samples onto the + // prom-client registry. When metrics are disabled, onMetric is + // undefined → the package's tool-timer/timeout hooks no-op with + // zero overhead. labels?.tool is guarded defensively (the tool + // wrapper always sets it). + return { + ...resolved.config, + sandbox: this.sandboxStore.asSink(), + onMetric: isMetricsEnabled() + ? ( + name: string, + value: number, + labels?: Record, + ) => { + if (name === 'mcp_tool_duration_seconds') { + observeMcpTool(labels?.tool ?? 'other', value); + } else if (name === 'collab_connect_timeouts_total') { + incConnectTimeout(); + } + } + : undefined, + }; }, { identify: (req: IncomingMessage) => { diff --git a/packages/mcp/src/client.ts b/packages/mcp/src/client.ts index 1ddc3432..d370a2b6 100644 --- a/packages/mcp/src/client.ts +++ b/packages/mcp/src/client.ts @@ -137,6 +137,15 @@ export type DocmostMcpConfig = { apiUrl: string } & ( has?: (uri: string) => boolean; evict?: (uri: string) => void; }; + // Dependency-neutral metrics sink. When present, the client emits generic + // (name, value, labels) samples; the HOST maps those names onto its own + // metrics registry (the package never depends on prom-client or the server). + // Absent in standalone/stdio mode → the client is a complete no-op here. + onMetric?: ( + name: string, + value: number, + labels?: Record, + ) => void; }; // Canonical UUID shape (versions 1–8, matching the `uuid` package's `validate` @@ -173,6 +182,11 @@ export class DocmostClient { // op's image blobs if the final doc put throws. Null when the sink omits them. private sandboxHas: ((uri: string) => boolean) | null = null; private sandboxEvict: ((uri: string) => void) | null = null; + // Optional dependency-neutral metrics sink (see DocmostMcpConfig.onMetric). + // Null on the legacy positional form and whenever the host omits it → no-op. + private onMetricFn: + | ((name: string, value: number, labels?: Record) => void) + | null = null; // In-flight login dedup: when the token expires, the 401 interceptor, // ensureAuthenticated, getCollabTokenWithReauth and the two multipart retries // can all call login() at once. Memoizing a single promise collapses that @@ -222,6 +236,8 @@ export class DocmostClient { this.sandboxHas = config.sandbox.has ?? null; this.sandboxEvict = config.sandbox.evict ?? null; } + // Legacy positional form carries no onMetric → null (complete no-op). + this.onMetricFn = config.onMetric ?? null; this.client = axios.create({ baseURL: this.apiUrl, // Default request timeout so a hung connection cannot wedge a per-page @@ -447,6 +463,10 @@ export class DocmostClient { }; connectTimer = setTimeout(() => { + // Only the actual 25s collab connect timeout fires here — the agent's + // collab connection to the server never became ready. This is the + // connect-vs-unload signal; the other finish() paths must NOT emit it. + this.onMetricFn?.("collab_connect_timeouts_total", 1); finish(new Error("Connection timeout to collaboration server")); }, CONNECT_TIMEOUT_MS); diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index 771c3e6f..de2cbcb3 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -64,6 +64,33 @@ const jsonContent = (data: any) => ({ * REST + the collaboration WebSocket using the provided service-account * credentials and auto-re-authenticates. */ +/** + * Wrap a tool handler so its wall-clock duration is reported through the host's + * dependency-neutral sink as `mcp_tool_duration_seconds` (labelled by tool + * name). Pure and side-effect-free apart from the optional `onMetric` call: + * - preserves the handler's exact return value (awaited); + * - observes in a `finally`, so it records on BOTH success and throw, then + * rethrows the original error unchanged (never swallowed); + * - with no `onMetric` (standalone/stdio) it is a transparent pass-through. + * Exported so the timing contract can be unit-tested without a live transport. + */ +export function timeToolHandler( + name: string, + handler: (...args: any[]) => any, + onMetric?: (name: string, value: number, labels?: Record) => void, +): (...args: any[]) => Promise { + return async (...handlerArgs: any[]) => { + const start = performance.now(); + try { + return await handler(...handlerArgs); + } finally { + onMetric?.("mcp_tool_duration_seconds", (performance.now() - start) / 1000, { + tool: name, + }); + } + }; +} + export function createDocmostMcpServer(config: DocmostMcpConfig): McpServer { // Pass the whole config union through: the client branches internally on // credentials vs. getToken, so both the external /mcp (creds) and the @@ -78,6 +105,26 @@ export function createDocmostMcpServer(config: DocmostMcpConfig): McpServer { { instructions: SERVER_INSTRUCTIONS }, ); + // Single choke point for MCP tool timing. Both `registerShared` (below) and + // the inline `server.registerTool(...)` calls funnel through this one method, + // so monkeypatching it HERE — before any tool is registered and before + // `registerShared` captures a reference to it — times every tool with no + // per-tool boilerplate. The wrapped handler records wall-clock duration and, + // in a `finally`, feeds the host's dependency-neutral sink + // `config.onMetric("mcp_tool_duration_seconds", seconds, { tool })`. The tool + // name is the registration name (bounded cardinality). When no onMetric is + // provided (standalone/stdio) the wrapper is a pure pass-through: it still + // returns the original result and rethrows the original error unchanged. + const originalRegisterTool = server.registerTool.bind(server) as ( + ...args: any[] + ) => any; + (server as any).registerTool = (...args: any[]) => { + const name = args[0] as string; + const handler = args[args.length - 1]; + const timedHandler = timeToolHandler(name, handler, config.onMetric); + return originalRegisterTool(...args.slice(0, -1), timedHandler); + }; + // Register a tool from the shared, zod-agnostic spec registry. The spec owns // the canonical name + model-facing description + (optional) schema builder; // only the execute body is supplied per call. buildShape is invoked with THIS diff --git a/packages/mcp/test/unit/tool-timing.test.mjs b/packages/mcp/test/unit/tool-timing.test.mjs new file mode 100644 index 00000000..543bb04f --- /dev/null +++ b/packages/mcp/test/unit/tool-timing.test.mjs @@ -0,0 +1,67 @@ +import { test } from "node:test"; +import assert from "node:assert/strict"; + +import { timeToolHandler } from "../../build/index.js"; + +// #402 — the registerTool wrapper times every tool through a single choke point +// and feeds the host's dependency-neutral onMetric sink. These assert the +// timing contract without a live transport (the factory monkeypatches +// server.registerTool with exactly this helper). + +test("times a tool and preserves the handler's return value", async () => { + const calls = []; + const onMetric = (name, value, labels) => calls.push({ name, value, labels }); + + const handler = async (arg) => ({ ok: true, echo: arg }); + const wrapped = timeToolHandler("get_page", handler, onMetric); + + const result = await wrapped("hello"); + // Return value passes through untouched. + assert.deepEqual(result, { ok: true, echo: "hello" }); + + // Exactly one sample, correct name/labels, numeric non-negative duration. + assert.equal(calls.length, 1); + assert.equal(calls[0].name, "mcp_tool_duration_seconds"); + assert.deepEqual(calls[0].labels, { tool: "get_page" }); + assert.equal(typeof calls[0].value, "number"); + assert.ok(calls[0].value >= 0, "duration must be non-negative seconds"); +}); + +test("standalone (no onMetric) is a transparent pass-through", async () => { + const handler = async (a, b) => a + b; + const wrapped = timeToolHandler("noop_tool", handler, undefined); + + // No throw, result unchanged, and nothing to observe. + const result = await wrapped(2, 3); + assert.equal(result, 5); +}); + +test("still observes on throw, then rethrows the original error", async () => { + const calls = []; + const onMetric = (name, value, labels) => calls.push({ name, value, labels }); + + const boom = new Error("handler failed"); + const handler = async () => { + throw boom; + }; + const wrapped = timeToolHandler("boom_tool", handler, onMetric); + + await assert.rejects(() => wrapped(), (err) => err === boom); + + // Observed exactly once despite the throw. + assert.equal(calls.length, 1); + assert.equal(calls[0].name, "mcp_tool_duration_seconds"); + assert.deepEqual(calls[0].labels, { tool: "boom_tool" }); +}); + +test("standalone still rethrows the original error (no swallow, no observe)", async () => { + const boom = new Error("standalone failure"); + const wrapped = timeToolHandler( + "boom_standalone", + async () => { + throw boom; + }, + undefined, + ); + await assert.rejects(() => wrapped(), (err) => err === boom); +});