diff --git a/integrations/pi-agent-plugin/src/attribution.test.ts b/integrations/pi-agent-plugin/src/attribution.test.ts new file mode 100644 index 000000000..79d938c62 --- /dev/null +++ b/integrations/pi-agent-plugin/src/attribution.test.ts @@ -0,0 +1,53 @@ +import { describe, expect, it } from "vitest"; +import { applySurfaceHeaders, PLATFORM_APPLICATION, PLATFORM_SOURCE } from "./attribution.ts"; + +function client(headers: Record = {}) { + return { headers: { Authorization: "Token k", ...headers } } as never; +} + +describe("applySurfaceHeaders", () => { + it("stamps the shared client so every path is attributed, not just commands", () => { + const mem0 = client(); + applySurfaceHeaders(mem0); + const headers = (mem0 as unknown as { headers: Record }).headers; + + expect(headers["X-Mem0-Source"]).toBe(PLATFORM_SOURCE); + expect(headers["X-Application"]).toBe(PLATFORM_APPLICATION); + expect(headers["X-Mem0-Client"]).toMatch(/^mem0-pi-agent\//); + expect(headers.Authorization).toBe("Token k"); + }); + + it("defers to a surface an outer wrapper already declared", () => { + const mem0 = client({ "X-Mem0-Source": "OPENCLAW", "X-Application": "vscode" }); + applySurfaceHeaders(mem0); + const headers = (mem0 as unknown as { headers: Record }).headers; + + expect(headers["X-Mem0-Source"]).toBe("OPENCLAW"); + expect(headers["X-Application"]).toBe("vscode"); + }); + + it("appends to the client stack rather than replacing it", () => { + const mem0 = client({ "X-Mem0-Client": "openclaw/2.1.0" }); + applySurfaceHeaders(mem0); + const headers = (mem0 as unknown as { headers: Record }).headers; + + expect(headers["X-Mem0-Client"]).toMatch(/^openclaw\/2\.1\.0, mem0-pi-agent\//); + }); + + it("treats a blank header as absent", () => { + const mem0 = client({ "X-Mem0-Source": " " }); + applySurfaceHeaders(mem0); + const headers = (mem0 as unknown as { headers: Record }).headers; + + expect(headers["X-Mem0-Source"]).toBe(PLATFORM_SOURCE); + }); + + it("bounds the stack so a long chain cannot grow the header without limit", () => { + const mem0 = client({ "X-Mem0-Client": "a/1, b/1, c/1, d/1, e/1" }); + applySurfaceHeaders(mem0); + const headers = (mem0 as unknown as { headers: Record }).headers; + + expect(headers["X-Mem0-Client"].split(",").length).toBeLessThanOrEqual(4); + expect(headers["X-Mem0-Client"].length).toBeLessThanOrEqual(200); + }); +}); diff --git a/integrations/pi-agent-plugin/src/attribution.ts b/integrations/pi-agent-plugin/src/attribution.ts new file mode 100644 index 000000000..7ab3ea6d6 --- /dev/null +++ b/integrations/pi-agent-plugin/src/attribution.ts @@ -0,0 +1,50 @@ +import type MemoryClient from "mem0ai"; +import * as fs from "node:fs"; + +/** Surface identity for this plugin, as the platform's EventSource knows it. */ +export const PLATFORM_SOURCE = "PI_AGENT"; + +/** Host app the plugin runs inside. Allowlisted server-side. */ +export const PLATFORM_APPLICATION = "pi"; + +const PLUGIN_VERSION = (() => { + try { + return JSON.parse( + fs.readFileSync(new URL("../package.json", import.meta.url), "utf-8"), + ).version as string; + } catch { + return "unknown"; + } +})(); + +const MAX_STACK_ENTRIES = 4; +const MAX_HEADER_CHARS = 200; + +/** + * Stamp surface identity onto the shared client, once, at construction. + * + * Tagging individual call sites was not enough: automatic recall, capture, the + * memory tools and deletion all go through this same client, so everything + * except the explicit slash commands reached the platform as generic SDK + * traffic. Every request method in the SDK sends `this.headers`, so setting + * them here covers all of them. + * + * X-Mem0-Source and X-Application are set-once, so a wrapper that already named + * a surface keeps it. X-Mem0-Client is append-only, so the platform sees the + * whole chain rather than only the last speaker. + */ +export function applySurfaceHeaders(client: MemoryClient): void { + const headers = client.headers as Record; + if (!headers["X-Mem0-Source"]?.trim()) headers["X-Mem0-Source"] = PLATFORM_SOURCE; + if (!headers["X-Application"]?.trim()) headers["X-Application"] = PLATFORM_APPLICATION; + + const existing = (headers["X-Mem0-Client"] ?? "") + .split(",") + .map((part) => part.trim()) + .filter(Boolean); + existing.push(`mem0-pi-agent/${PLUGIN_VERSION}`); + headers["X-Mem0-Client"] = existing + .slice(0, MAX_STACK_ENTRIES) + .join(", ") + .slice(0, MAX_HEADER_CHARS); +} diff --git a/integrations/pi-agent-plugin/src/commands.ts b/integrations/pi-agent-plugin/src/commands.ts index a733aea0a..e969478a1 100644 --- a/integrations/pi-agent-plugin/src/commands.ts +++ b/integrations/pi-agent-plugin/src/commands.ts @@ -6,10 +6,12 @@ import { DEFAULT_CUSTOM_CATEGORIES } from "./types.ts"; import { resolveSearchFilters, resolveAddParams } from "./memory/scoping.ts"; import { formatMemoryList, formatMemoryCompact, groupByCategory } from "./memory/formatting.ts"; import { captureCommandEvent } from "./telemetry.ts"; +import { PLATFORM_SOURCE } from "./attribution.ts"; -// Surface attribution on the wire. Previously a PostHog property only, so -// the platform saw these calls as generic SDK traffic. -const PLATFORM_SOURCE = "PI_AGENT"; +// Wire identity is set once on the shared client in entry.ts, which covers +// every path including recall, capture, tools and deletion. It stays in the +// body of the two calls below as well: body `source` is what the backend reads +// when the header is absent. const SEARCH_TOP_K = 10; diff --git a/integrations/pi-agent-plugin/src/entry.ts b/integrations/pi-agent-plugin/src/entry.ts index d49d78ad0..256c05aa4 100644 --- a/integrations/pi-agent-plugin/src/entry.ts +++ b/integrations/pi-agent-plugin/src/entry.ts @@ -10,6 +10,7 @@ import { captureEvent } from "./telemetry.ts"; import * as os from "node:os"; import type { ScopeContext } from "./types.ts"; import { createMemoryLifecycle } from "../../agent-plugin-core/typescript/src/lifecycle.ts"; +import { applySurfaceHeaders } from "./attribution.ts"; export { buildRecallContext } from "../../agent-plugin-core/typescript/src/lifecycle.ts"; @@ -29,6 +30,11 @@ export default function mem0Extension(pi: ExtensionAPI): void { } const mem0 = new MemoryClient({ apiKey: config.apiKey }); + // Every path below shares this client: automatic recall, capture, the memory + // tools and deletion as well as the slash commands. Attribution belongs here + // rather than on individual calls, or everything except the commands reports + // as generic SDK traffic. + applySurfaceHeaders(mem0); const scopeCtx: ScopeContext = { userId: resolveUserId(config.userId),