From 0c9d986ddfab884cf6178c048083bdd0dd9470c9 Mon Sep 17 00:00:00 2001 From: Hubert Gancarczyk Date: Thu, 11 Jun 2026 09:31:48 +0200 Subject: [PATCH] feat: record AI client metadata for telemetry --- package-lock.json | 5 +- packages/argent-mcp/package.json | 1 + packages/argent-mcp/src/mcp-server.ts | 21 ++++- packages/telemetry/src/ai-identity.ts | 74 +++++++++++++++ packages/telemetry/src/events.ts | 7 +- packages/telemetry/src/index.ts | 8 ++ packages/telemetry/src/registry-listener.ts | 8 +- packages/telemetry/src/sanitize.ts | 12 +++ packages/telemetry/test/ai-identity.test.ts | 64 +++++++++++++ .../telemetry/test/registry-listener.test.ts | 40 ++++++++ packages/telemetry/test/sanitize.test.ts | 47 ++++++++++ packages/tool-server/src/http.ts | 57 ++++++++--- packages/tool-server/src/index.ts | 2 + .../tool-server/test/http-tools-meta.test.ts | 94 +++++++++++++++++++ 14 files changed, 420 insertions(+), 20 deletions(-) create mode 100644 packages/telemetry/src/ai-identity.ts create mode 100644 packages/telemetry/test/ai-identity.test.ts diff --git a/package-lock.json b/package-lock.json index e13e8db1..3fcda933 100644 --- a/package-lock.json +++ b/package-lock.json @@ -3755,7 +3755,9 @@ "dependencies": { "@argent/registry": "file:../registry", "@argent/telemetry": "file:../telemetry", - "@argent/tools-client": "file:../argent-tools-client" + "@argent/tools-client": "file:../argent-tools-client", + "@clack/prompts": "^1.1.0", + "picocolors": "^1.1.1" }, "devDependencies": { "@types/node": "^25.9.0", @@ -3851,6 +3853,7 @@ "name": "@argent/mcp", "version": "0.10.0", "dependencies": { + "@argent/telemetry": "file:../telemetry", "@argent/tools-client": "file:../argent-tools-client", "@modelcontextprotocol/sdk": "^1.29.0" }, diff --git a/packages/argent-mcp/package.json b/packages/argent-mcp/package.json index f36c2ae8..d059f172 100644 --- a/packages/argent-mcp/package.json +++ b/packages/argent-mcp/package.json @@ -12,6 +12,7 @@ "typecheck:tests": "tsc --noEmit -p tsconfig.test.json" }, "dependencies": { + "@argent/telemetry": "file:../telemetry", "@argent/tools-client": "file:../argent-tools-client", "@modelcontextprotocol/sdk": "^1.29.0" }, diff --git a/packages/argent-mcp/src/mcp-server.ts b/packages/argent-mcp/src/mcp-server.ts index 7db8d1ae..1dc00ef9 100644 --- a/packages/argent-mcp/src/mcp-server.ts +++ b/packages/argent-mcp/src/mcp-server.ts @@ -11,6 +11,7 @@ import { type ToolMeta, type ToolsServerPaths, } from "@argent/tools-client"; +import { canonicalizeAiClient, AI_CLIENT_NAME_PATTERN } from "@argent/telemetry"; import { toMcpContent, flowRunToMcpContent, @@ -106,6 +107,24 @@ export async function startMcpServer(options: StartMcpServerOptions): Promise { + const rawName = server.getClientVersion()?.name?.trim() || undefined; + const aiClient = canonicalizeAiClient(rawName); + if (aiClient) return { "X-Argent-AI-Client": aiClient }; + if (rawName) { + return { + "X-Argent-AI-Client": "other", + ...(AI_CLIENT_NAME_PATTERN.test(rawName) ? { "X-Argent-AI-Client-Name": rawName } : {}), + }; + } + return {}; + } + let reconnectPromise: Promise | null = null; async function reconnect(): Promise { @@ -162,7 +181,7 @@ export async function startMcpServer(options: StartMcpServerOptions): Promise `${TOOLS_URL}/tools/${name}`, reconnect, { init: { method: "POST", - headers: { "Content-Type": "application/json", ...authHeader() }, + headers: { "Content-Type": "application/json", ...authHeader(), ...aiClientHeaders() }, body: JSON.stringify(args ?? {}), }, fetchTimeoutMs: meta?.longRunning ? null : FETCH_TIMEOUT_MS, diff --git a/packages/telemetry/src/ai-identity.ts b/packages/telemetry/src/ai-identity.ts new file mode 100644 index 00000000..dd7ebcfa --- /dev/null +++ b/packages/telemetry/src/ai-identity.ts @@ -0,0 +1,74 @@ +// Coarse identity of the AI coding tool driving the MCP server. We record only +// a canonical client slug (which tool) — never prompts, model output, or args. +// +// The single signal is the MCP `initialize` handshake `clientInfo.name` (ground +// truth of what is actually connecting), read at runtime via +// `Server.getClientVersion()`. Anything we can't map is reported as `other` with +// a sanitized free-form name. + +export const AI_CLIENTS = [ + "codex", + "claude_code", + "cursor", + "gemini", + "vscode", + "windsurf", + "zed", + "opencode", + "copilot", + "other", +] as const; + +export type AiClient = (typeof AI_CLIENTS)[number]; + +/** Bounded free-form client name, used for the long tail of unrecognized tools. */ +export const AI_CLIENT_NAME_PATTERN = /^[A-Za-z0-9 ._-]{1,80}$/; + +export type AiTelemetryProps = { + ai_client?: AiClient; + /** Raw, sanitized clientInfo.name — only carried when `ai_client` is `other`. */ + ai_client_name?: string; +}; + +// Runtime MCP `clientInfo.name` → canonical slug. Patterns are tested against the +// trimmed, lower-cased name. Each is verified against the tool's source; we match +// the CLIENT identity precisely so decoys are excluded +const RUNTIME_CLIENT_PATTERNS: ReadonlyArray = [ + [/^codex-mcp-client\b/, "codex"], + [/^claude-code\b/, "claude_code"], + [/^cursor\b/, "cursor"], + [/^gemini-cli-mcp-client\b/, "gemini"], + [/^visual studio code\b/, "vscode"], + [/^code - oss\b/, "vscode"], + [/^windsurf\b/, "windsurf"], + [/^zed\b/, "zed"], + [/^opencode\b/, "opencode"], + [/^github-copilot-developer\b/, "copilot"], +]; + +/** + * Pick out only the AI-client telemetry keys from a wider metadata object, + * omitting any that are absent so events never carry `undefined` values. Shared + * by every emitter so the spread shape stays identical across call sites. + */ +export function aiTelemetryFromMeta(meta: AiTelemetryProps): AiTelemetryProps { + return { + ...(meta.ai_client ? { ai_client: meta.ai_client } : {}), + ...(meta.ai_client_name ? { ai_client_name: meta.ai_client_name } : {}), + }; +} + +/** + * Normalize a runtime MCP `clientInfo.name` to an {@link AiClient}. Returns + * `undefined` for anything unrecognized (callers may then fall back to `other` + + * a sanitized free-form name). + */ +export function canonicalizeAiClient(value: string | undefined | null): AiClient | undefined { + if (typeof value !== "string") return undefined; + const lower = value.trim().toLowerCase(); + if (!lower) return undefined; + for (const [pattern, slug] of RUNTIME_CLIENT_PATTERNS) { + if (pattern.test(lower)) return slug; + } + return undefined; +} diff --git a/packages/telemetry/src/events.ts b/packages/telemetry/src/events.ts index e8eda6d4..2c52b301 100644 --- a/packages/telemetry/src/events.ts +++ b/packages/telemetry/src/events.ts @@ -2,6 +2,7 @@ // same surface at runtime. import type { FailureSignal } from "@argent/registry"; +import type { AiTelemetryProps } from "./ai-identity.js"; // Installation events @@ -98,20 +99,20 @@ export interface InstallationCliUninstallCompleteProps extends FailureTelemetryP // Tool usage events -export interface ToolInvokeProps { +export interface ToolInvokeProps extends AiTelemetryProps { tool: string; tool_invocation_id: string; platform?: "ios" | "android"; } -export interface ToolCompleteProps { +export interface ToolCompleteProps extends AiTelemetryProps { tool: string; tool_invocation_id: string; platform?: "ios" | "android"; duration_ms: number; } -export interface ToolFailProps extends FailureTelemetryProps { +export interface ToolFailProps extends FailureTelemetryProps, AiTelemetryProps { tool: string; tool_invocation_id?: string; platform?: "ios" | "android"; diff --git a/packages/telemetry/src/index.ts b/packages/telemetry/src/index.ts index 8f06d4d3..8c7a8cc2 100644 --- a/packages/telemetry/src/index.ts +++ b/packages/telemetry/src/index.ts @@ -27,6 +27,14 @@ export { EVENT_NAMES } from "./events.js"; export { isDebugEnabled } from "./debug.js"; export { getConsentState } from "./consent.js"; export { getSessionId } from "./base-props.js"; +export { + AI_CLIENTS, + AI_CLIENT_NAME_PATTERN, + canonicalizeAiClient, + aiTelemetryFromMeta, + type AiClient, + type AiTelemetryProps, +} from "./ai-identity.js"; const SHORT_FLUSH_TIMEOUT_MS = 1_500; diff --git a/packages/telemetry/src/registry-listener.ts b/packages/telemetry/src/registry-listener.ts index d517132a..2355e896 100644 --- a/packages/telemetry/src/registry-listener.ts +++ b/packages/telemetry/src/registry-listener.ts @@ -1,9 +1,10 @@ import { FAILURE_CODES, getFailureSignalOrFallback, type Registry } from "@argent/registry"; import { track } from "./index.js"; +import { aiTelemetryFromMeta, type AiTelemetryProps } from "./ai-identity.js"; // HTTP captures request-only metadata here so registry lifecycle events can -// include platform context without carrying raw params. -export interface InvocationMeta { +// include platform context (and the coarse AI client) without carrying raw params. +export interface InvocationMeta extends AiTelemetryProps { platform?: "ios" | "android"; } @@ -33,6 +34,7 @@ export function attachRegistryTelemetry(registry: Registry): AttachHandle { tool: toolId, tool_invocation_id: toolInvocationId, ...(meta.platform ? { platform: meta.platform } : {}), + ...aiTelemetryFromMeta(meta), }); }; @@ -43,6 +45,7 @@ export function attachRegistryTelemetry(registry: Registry): AttachHandle { tool_invocation_id: toolInvocationId, ...(meta.platform ? { platform: meta.platform } : {}), duration_ms: durationMs, + ...aiTelemetryFromMeta(meta), }); }; @@ -65,6 +68,7 @@ export function attachRegistryTelemetry(registry: Registry): AttachHandle { ...(meta.platform ? { platform: meta.platform } : {}), duration_ms: durationMs, ...signal, + ...aiTelemetryFromMeta(meta), }); }; diff --git a/packages/telemetry/src/sanitize.ts b/packages/telemetry/src/sanitize.ts index f3405a47..1e491c25 100644 --- a/packages/telemetry/src/sanitize.ts +++ b/packages/telemetry/src/sanitize.ts @@ -8,6 +8,7 @@ import { NETWORK_FAILURES, } from "@argent/registry"; import type { EventName } from "./events.js"; +import { AI_CLIENTS, AI_CLIENT_NAME_PATTERN } from "./ai-identity.js"; // Per-event property allowlist and validators. Unknown keys and invalid values // are dropped before anything reaches PostHog. @@ -82,6 +83,14 @@ const FAILURE_SIGNAL_NAME = oneOf(FAILURE_SIGNAL_NAMES); const FAILURE_SPAWN_CODE = oneOf(FAILURE_SPAWN_CODES); const NETWORK_FAILURE = oneOf(NETWORK_FAILURES); +const AI_CLIENT = oneOf(AI_CLIENTS); +const AI_CLIENT_NAME = matches(AI_CLIENT_NAME_PATTERN, 80); + +const AI_TELEMETRY = { + ai_client: AI_CLIENT, + ai_client_name: AI_CLIENT_NAME, +}; + const FAILURE_SIGNAL = { error_code: ERROR_CODE, failure_stage: FAILURE_STAGE, @@ -166,12 +175,14 @@ export const ALLOWED: Record> = { tool: TOOL_NAME, tool_invocation_id: UUID, platform: PLATFORM, + ...AI_TELEMETRY, }, "tool:complete": { tool: TOOL_NAME, tool_invocation_id: UUID, platform: PLATFORM, duration_ms: DURATION_MS, + ...AI_TELEMETRY, }, "tool:fail": { tool: TOOL_NAME, @@ -179,6 +190,7 @@ export const ALLOWED: Record> = { platform: PLATFORM, duration_ms: DURATION_MS, ...FAILURE_SIGNAL, + ...AI_TELEMETRY, }, "cli:run_fail": { tool: TOOL_NAME, diff --git a/packages/telemetry/test/ai-identity.test.ts b/packages/telemetry/test/ai-identity.test.ts new file mode 100644 index 00000000..59d685d5 --- /dev/null +++ b/packages/telemetry/test/ai-identity.test.ts @@ -0,0 +1,64 @@ +import { describe, expect, it } from "vitest"; +import { AI_CLIENT_NAME_PATTERN, canonicalizeAiClient } from "../src/ai-identity.js"; + +describe("canonicalizeAiClient", () => { + describe("runtime MCP clientInfo.name (verified from each tool's source)", () => { + const RUNTIME_NAMES: Array<[string, string]> = [ + ["codex-mcp-client", "codex"], + ["claude-code", "claude_code"], + ["cursor-vscode", "cursor"], + ["gemini-cli-mcp-client", "gemini"], + ["gemini-cli-mcp-client-myserver", "gemini"], // older suffixed form + ["Visual Studio Code", "vscode"], + ["Visual Studio Code - Insiders", "vscode"], + ["Code - OSS", "vscode"], + ["windsurf", "windsurf"], + ["Zed", "zed"], + ["opencode", "opencode"], + ["github-copilot-developer", "copilot"], // observed in real Copilot MCP logs + ]; + + it.each(RUNTIME_NAMES)("maps %s → %s", (name, slug) => { + expect(canonicalizeAiClient(name)).toBe(slug); + }); + + it("does not mistake Codex's server-mode identity for the client", () => { + // We are always the server; we only ever receive `codex-mcp-client`. The + // server-mode string must not be matched by the client patterns. + expect(canonicalizeAiClient("codex-mcp-server")).toBeUndefined(); + }); + }); + + it("returns undefined for unrecognized clients (caller decides on `other`)", () => { + // `other` is a caller decision (raw clientInfo.name we couldn't map), not a + // value canonicalize should ever produce from an unknown string. + expect(canonicalizeAiClient("other")).toBeUndefined(); + expect(canonicalizeAiClient("some-unknown-mcp-client")).toBeUndefined(); + // Hermes sends a generic name and is no longer specially attributed. + expect(canonicalizeAiClient("mcp")).toBeUndefined(); + }); + + it("returns undefined for empty / non-string input", () => { + expect(canonicalizeAiClient(undefined)).toBeUndefined(); + expect(canonicalizeAiClient(null)).toBeUndefined(); + expect(canonicalizeAiClient("")).toBeUndefined(); + expect(canonicalizeAiClient(" ")).toBeUndefined(); + }); + + it("tolerates surrounding whitespace", () => { + expect(canonicalizeAiClient(" codex-mcp-client ")).toBe("codex"); + }); +}); + +describe("AI_CLIENT_NAME_PATTERN", () => { + it("accepts real clientInfo.name values", () => { + for (const name of ["codex-mcp-client", "Visual Studio Code - Insiders", "Zed", "opencode"]) { + expect(AI_CLIENT_NAME_PATTERN.test(name)).toBe(true); + } + }); + + it("rejects path-like or oversized values", () => { + expect(AI_CLIENT_NAME_PATTERN.test("/Users/alice/secret-client")).toBe(false); + expect(AI_CLIENT_NAME_PATTERN.test("x".repeat(81))).toBe(false); + }); +}); diff --git a/packages/telemetry/test/registry-listener.test.ts b/packages/telemetry/test/registry-listener.test.ts index 8e630394..9d38b450 100644 --- a/packages/telemetry/test/registry-listener.test.ts +++ b/packages/telemetry/test/registry-listener.test.ts @@ -50,6 +50,46 @@ describe("attachRegistryTelemetry", () => { handle.detach(); }); + it("threads the AI client through invoke / complete / fail events", () => { + const trackSpy = vi.spyOn(telemetry, "track"); + const registry = new Registry(); + const handle = attachRegistryTelemetry(registry); + + handle.recordInvocation(INVOCATION_ID_1, { platform: "ios", ai_client: "codex" }); + registry.events.emit("toolInvoked", "gesture-tap", INVOCATION_ID_1); + registry.events.emit("toolCompleted", "gesture-tap", INVOCATION_ID_1, 10); + + handle.recordInvocation(INVOCATION_ID_2, { + ai_client: "other", + ai_client_name: "some-new-tool", + }); + registry.events.emit("toolInvoked", "screenshot", INVOCATION_ID_2); + registry.events.emit("toolFailed", "screenshot", INVOCATION_ID_2, new Error("boom"), 5); + + expect(trackSpy.mock.calls[0]![1]).toMatchObject({ ai_client: "codex" }); + expect(trackSpy.mock.calls[1]![1]).toMatchObject({ ai_client: "codex" }); + expect(trackSpy.mock.calls[3]![1]).toMatchObject({ + ai_client: "other", + ai_client_name: "some-new-tool", + }); + + handle.detach(); + }); + + it("omits AI keys entirely when no client metadata was recorded", () => { + const trackSpy = vi.spyOn(telemetry, "track"); + const registry = new Registry(); + const handle = attachRegistryTelemetry(registry); + + handle.recordInvocation(INVOCATION_ID_1, { platform: "ios" }); + registry.events.emit("toolInvoked", "gesture-tap", INVOCATION_ID_1); + + expect(trackSpy.mock.calls[0]![1]).not.toHaveProperty("ai_client"); + expect(trackSpy.mock.calls[0]![1]).not.toHaveProperty("ai_client_name"); + + handle.detach(); + }); + it("emits tool:fail with tool metadata and real duration", () => { const trackSpy = vi.spyOn(telemetry, "track"); const registry = new Registry(); diff --git a/packages/telemetry/test/sanitize.test.ts b/packages/telemetry/test/sanitize.test.ts index c9f3a149..052e321a 100644 --- a/packages/telemetry/test/sanitize.test.ts +++ b/packages/telemetry/test/sanitize.test.ts @@ -59,6 +59,53 @@ describe("sanitize", () => { }); } }); + + it("accepts coarse AI client metadata on tool events", () => { + expect( + sanitize("tool:invoke", { + tool: "gesture-tap", + tool_invocation_id: "11111111-1111-4111-8111-111111111111", + ai_client: "codex", + ai_client_name: "codex-mcp-client", + }) + ).toEqual({ + tool: "gesture-tap", + tool_invocation_id: "11111111-1111-4111-8111-111111111111", + ai_client: "codex", + ai_client_name: "codex-mcp-client", + }); + }); + + it("drops unregistered AI client slugs and path-leaking client names", () => { + expect( + sanitize("tool:invoke", { + tool: "gesture-tap", + tool_invocation_id: "11111111-1111-4111-8111-111111111111", + ai_client: "my-private-client", + ai_client_name: "/Users/alice/secret-client", + }) + ).toEqual({ + tool: "gesture-tap", + tool_invocation_id: "11111111-1111-4111-8111-111111111111", + }); + }); + + it("carries AI client metadata on tool:fail alongside the failure signal", () => { + expect( + sanitize("tool:fail", { + tool: "screenshot", + tool_invocation_id: "11111111-1111-4111-8111-111111111111", + duration_ms: 12, + error_code: FAILURE_CODES.REGISTRY_TOOL_FAILURE_UNCLASSIFIED, + ai_client: "other", + ai_client_name: "some-new-tool", + }) + ).toMatchObject({ + ai_client: "other", + ai_client_name: "some-new-tool", + error_code: FAILURE_CODES.REGISTRY_TOOL_FAILURE_UNCLASSIFIED, + }); + }); }); describe("matches validator", () => { diff --git a/packages/tool-server/src/http.ts b/packages/tool-server/src/http.ts index 86f92dcd..283fa005 100644 --- a/packages/tool-server/src/http.ts +++ b/packages/tool-server/src/http.ts @@ -1,6 +1,7 @@ import express, { Request, Response } from "express"; import { randomUUID } from "node:crypto"; import { FAILURE_CODES, type FailureSignal, type Registry } from "@argent/registry"; +import { AI_CLIENTS, AI_CLIENT_NAME_PATTERN, type AiTelemetryProps } from "@argent/telemetry"; import { ToolNotFoundError } from "@argent/registry"; import { createIdleTimer } from "./utils/idle-timer"; import { DependencyMissingError, ensureDeps } from "./utils/check-deps"; @@ -63,13 +64,13 @@ function extractDeviceArg(data: unknown): string | null { return null; } -type InvocationMeta = { platform?: "ios" | "android" }; +type InvocationMeta = { platform?: "ios" | "android" } & AiTelemetryProps; // Only coarse platform context is retained for failure telemetry. The raw // device id (UDID / serial) is used transiently to infer platform and never // stored or forwarded. -type HttpFailureMeta = { platform?: "ios" | "android" }; +type HttpFailureMeta = { platform?: "ios" | "android" } & AiTelemetryProps; -function inferPlatform(deviceId: string | null): HttpFailureMeta["platform"] | null { +function inferPlatform(deviceId: string | null): "ios" | "android" | null { if (!deviceId) return null; try { return resolveDevice(deviceId).platform; @@ -78,17 +79,45 @@ function inferPlatform(deviceId: string | null): HttpFailureMeta["platform"] | n } } -function extractInvocationMeta(hasCapability: boolean, data: unknown): InvocationMeta | null { - if (!hasCapability || !data || typeof data !== "object") return null; - const record = data as Record; - const deviceArg = extractDeviceArg(record); - if (deviceArg) { - return { platform: resolveDevice(deviceArg).platform }; +function firstHeader(value: string | string[] | undefined): string | undefined { + return Array.isArray(value) ? value[0] : value; +} + +// The MCP server forwards the coarse AI-client identity as request headers (it +// lives in a different process). We re-validate here against the same allowlist / +// pattern the sanitizer enforces, so a misbehaving client can't inject arbitrary +// values into telemetry. The free-form name is only retained when the client is +// `other` — mirroring the producer and the documented `AiTelemetryProps` contract, +// so a recognized (or absent) client can never carry a stray name. +function extractAiTelemetryMeta(req: Request): AiTelemetryProps { + const meta: AiTelemetryProps = {}; + const client = firstHeader(req.headers["x-argent-ai-client"]); + if (client && (AI_CLIENTS as readonly string[]).includes(client)) { + meta.ai_client = client as AiTelemetryProps["ai_client"]; } - if (typeof record.avdName === "string") { - return { platform: "android" }; + const clientName = firstHeader(req.headers["x-argent-ai-client-name"]); + if (meta.ai_client === "other" && clientName && AI_CLIENT_NAME_PATTERN.test(clientName)) { + meta.ai_client_name = clientName; } - return null; + return meta; +} + +function extractInvocationMeta( + hasCapability: boolean, + data: unknown, + aiMeta: AiTelemetryProps +): InvocationMeta | null { + const meta: InvocationMeta = { ...aiMeta }; + if (hasCapability && data && typeof data === "object") { + const record = data as Record; + const deviceArg = extractDeviceArg(record); + if (deviceArg) { + meta.platform = resolveDevice(deviceArg).platform; + } else if (typeof record.avdName === "string") { + meta.platform = "android"; + } + } + return Object.keys(meta).length > 0 ? meta : null; } // ── HTTP app ──────────────────────────────────────────────────────── @@ -295,6 +324,7 @@ export function createHttpApp(registry: Registry, options?: HttpAppOptions): Htt async (req: Request, res: Response) => { const name = req.params.name!; const requestStartedAt = performance.now(); + const aiMeta = extractAiTelemetryMeta(req); const emitHttpFailure = ( signal: FailureSignal, @@ -307,6 +337,7 @@ export function createHttpApp(registry: Registry, options?: HttpAppOptions): Htt name, { ...(platform ? { platform } : {}), + ...aiMeta, }, signal, performance.now() - requestStartedAt @@ -424,7 +455,7 @@ export function createHttpApp(registry: Registry, options?: HttpAppOptions): Htt const toolInvocationId = randomUUID(); let releaseInvocationMeta: (() => void) | undefined; if (options?.recordInvocation) { - const invocationMeta = extractInvocationMeta(Boolean(def.capability), parsedData); + const invocationMeta = extractInvocationMeta(Boolean(def.capability), parsedData, aiMeta); if (invocationMeta) { releaseInvocationMeta = options.recordInvocation(toolInvocationId, invocationMeta); } diff --git a/packages/tool-server/src/index.ts b/packages/tool-server/src/index.ts index 1d2fd81b..1be7d100 100644 --- a/packages/tool-server/src/index.ts +++ b/packages/tool-server/src/index.ts @@ -4,6 +4,7 @@ import { attachRegistryTelemetry, track as telemetryTrack, shutdown as telemetryShutdown, + aiTelemetryFromMeta, } from "@argent/telemetry"; import { createHttpApp } from "./http"; import { createRegistry } from "./utils/setup-registry"; @@ -159,6 +160,7 @@ export function start(): void { ...(meta.platform ? { platform: meta.platform } : {}), duration_ms: durationMs, ...signal, + ...aiTelemetryFromMeta(meta), }); }, }); diff --git a/packages/tool-server/test/http-tools-meta.test.ts b/packages/tool-server/test/http-tools-meta.test.ts index 601a8509..04693570 100644 --- a/packages/tool-server/test/http-tools-meta.test.ts +++ b/packages/tool-server/test/http-tools-meta.test.ts @@ -168,4 +168,98 @@ describe("GET /tools progressive-loading metadata", () => { expect(recordInvocation).toHaveBeenCalledWith(expect.any(String), { platform: "android" }); }); + + it("records the AI client from request headers alongside platform", async () => { + let seenMeta: Record | undefined; + const recordInvocation = vi.fn((_id: string, meta: Record) => { + seenMeta = meta; + return vi.fn(); + }); + handle.dispose(); + handle = createHttpApp(stubRegistry(), { recordInvocation }); + + await request(handle.app) + .post("/tools/device-tool") + .set("X-Argent-AI-Client", "codex") + .send({ udid: "11111111-1111-1111-1111-111111111111" }) + .expect(200); + + expect(seenMeta).toEqual({ platform: "ios", ai_client: "codex" }); + }); + + it("captures an unknown tool by name when ai_client is `other`", async () => { + let seenMeta: Record | undefined; + const recordInvocation = vi.fn((_id: string, meta: Record) => { + seenMeta = meta; + return vi.fn(); + }); + handle.dispose(); + handle = createHttpApp(stubRegistry(), { recordInvocation }); + + // A non-device tool with no platform context still records because AI + // metadata is present on the request. + await request(handle.app) + .post("/tools/plain-tool") + .set("X-Argent-AI-Client", "other") + .set("X-Argent-AI-Client-Name", "some-new-tool") + .send({}) + .expect(200); + + expect(seenMeta).toEqual({ ai_client: "other", ai_client_name: "some-new-tool" }); + }); + + it("drops unregistered AI client slugs and path-leaking client names", async () => { + const recordInvocation = vi.fn(() => vi.fn()); + handle.dispose(); + handle = createHttpApp(stubRegistry(), { recordInvocation }); + + // Unknown slug + unsafe name → nothing usable → no metadata → not recorded + // for a non-device tool. + await request(handle.app) + .post("/tools/plain-tool") + .set("X-Argent-AI-Client", "evil-client") + .set("X-Argent-AI-Client-Name", "/Users/alice/secret") + .send({}) + .expect(200); + + expect(recordInvocation).not.toHaveBeenCalled(); + }); + + it("drops the free-form client name unless ai_client is `other`", async () => { + let seenMeta: Record | undefined; + const recordInvocation = vi.fn((_id: string, meta: Record) => { + seenMeta = meta; + return vi.fn(); + }); + handle.dispose(); + handle = createHttpApp(stubRegistry(), { recordInvocation }); + + // A recognized client must never carry a free-form name (the invariant is + // that the name is only meaningful for the `other` long tail). The name is + // a perfectly valid pattern here — it is dropped purely on the coupling. + await request(handle.app) + .post("/tools/device-tool") + .set("X-Argent-AI-Client", "codex") + .set("X-Argent-AI-Client-Name", "claude-code") + .send({ udid: "11111111-1111-1111-1111-111111111111" }) + .expect(200); + + expect(seenMeta).toEqual({ platform: "ios", ai_client: "codex" }); + }); + + it("drops a client name sent with no ai_client header", async () => { + const recordInvocation = vi.fn(() => vi.fn()); + handle.dispose(); + handle = createHttpApp(stubRegistry(), { recordInvocation }); + + // Name-only, no slug → nothing usable → no metadata → not recorded for a + // non-device tool. + await request(handle.app) + .post("/tools/plain-tool") + .set("X-Argent-AI-Client-Name", "some-new-tool") + .send({}) + .expect(200); + + expect(recordInvocation).not.toHaveBeenCalled(); + }); });