fix(pi-agent): attribute the shared client, not two command call sites
Review finding from @kartik-mem0 on this PR. Only the explicit slash commands set a source. Automatic recall at entry.ts:80, the capture path, the memory tools and deletion all go through the same client constructed at entry.ts:31 with no attribution, so everything except the commands still reached the platform as generic SDK traffic. That is most of the plugin's traffic. Identity is now stamped once on the shared client. Deliberately by mutating client.headers rather than through the SDK's MEM0_SOURCE environment support: this package pins mem0ai ^3.0.7, the installed build has no such support, and setting an environment variable it does not read would have looked like a fix and changed nothing. Every request method in the published client sends this.headers, so this covers all of them and keeps working when the SDK gains the env path. Set-once and append-only are preserved, so a wrapper that already named a surface keeps it and the client stack accumulates rather than being replaced. PLATFORM_SOURCE moves into the new module and commands.ts imports it; the body source stays on those two calls because that is what the backend reads when the header is absent. Five tests for the header contract, plus the existing suite: 94 pi-agent tests pass and the tsup DTS build is clean. Python side 307 passed, 8 skipped. Claude-Session: https://claude.ai/code/session_01C7tEmH86HAr7GoAAKCEHZb
This commit is contained in:
@@ -0,0 +1,53 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { applySurfaceHeaders, PLATFORM_APPLICATION, PLATFORM_SOURCE } from "./attribution.ts";
|
||||
|
||||
function client(headers: Record<string, string> = {}) {
|
||||
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<string, string> }).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<string, string> }).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<string, string> }).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<string, string> }).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<string, string> }).headers;
|
||||
|
||||
expect(headers["X-Mem0-Client"].split(",").length).toBeLessThanOrEqual(4);
|
||||
expect(headers["X-Mem0-Client"].length).toBeLessThanOrEqual(200);
|
||||
});
|
||||
});
|
||||
@@ -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<string, string>;
|
||||
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);
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user