Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
16 kB · 524 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525import { lstatSync, mkdtempSync, readFileSync, statSync, symlinkSync, writeFileSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import type { AssistantMessage } from "@earendil-works/pi-ai";import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";import type { AgentSession, AgentSessionEvent } from "../src/core/agent-session.js";import { SettingsManager } from "../src/core/settings-manager.js";import { captureAgentCommandUsed, captureOnboardingCompleted, getOrCreateTelemetryInstallationId, installAgentTelemetry, isTelemetryEnabled, TelemetryClient, type TelemetryEvent, type TelemetryEventName, type TelemetrySink, telemetryAuthCategory,} from "../src/core/telemetry.js";
function uuidGenerator(): () => string { let counter = 0; return () => { counter++; return `00000000-0000-4000-8000-${String(counter).padStart(12, "0")}`; };}
function assistantMessage(overrides: Partial<AssistantMessage> = {}): AssistantMessage { return { role: "assistant", content: [{ type: "text", text: "private response" }], api: "openai-responses", provider: "openai", model: "gpt-5", usage: { input: 100, output: 20, cacheRead: 50, cacheWrite: 0, totalTokens: 170, cost: { input: 0.001, output: 0.002, cacheRead: 0.0001, cacheWrite: 0, total: 0.0031, }, }, stopReason: "stop", timestamp: 1, ...overrides, };}
class FakeTelemetrySink implements TelemetrySink { readonly events: Array<{ name: TelemetryEventName; properties: Record<string, string | number | boolean | null> }> = []; flushCount = 0;
capture(name: TelemetryEventName, properties: Record<string, string | number | boolean | null>): void { this.events.push({ name, properties }); }
async flush(): Promise<void> { this.flushCount++; }}
class FakeAgentSession { private listener?: (event: AgentSessionEvent) => void; private disposeCallback?: () => void | Promise<void>;
subscribe(listener: (event: AgentSessionEvent) => void): () => void { this.listener = listener; return () => { this.listener = undefined; }; }
registerDisposeCallback(callback: () => void | Promise<void>): void { this.disposeCallback = callback; }
emit(event: AgentSessionEvent): void { this.listener?.(event); }
dispose(): void { void this.disposeCallback?.(); }
async disposeAsync(): Promise<void> { await this.disposeCallback?.(); }}
afterEach(() => { vi.unstubAllEnvs();});
describe("telemetry identity and transport", () => { it("creates a private stable installation ID", () => { const agentDir = mkdtempSync(join(tmpdir(), "prime-agent-telemetry-")); const randomId = uuidGenerator();
const first = getOrCreateTelemetryInstallationId(agentDir, randomId); const second = getOrCreateTelemetryInstallationId(agentDir, randomId);
expect(second).toBe(first); const path = join(agentDir, "telemetry.json"); expect(JSON.parse(readFileSync(path, "utf8"))).toEqual({ version: 1, installationId: first, }); expect(statSync(path).mode & 0o777).toBe(0o600); });
it("replaces invalid persisted installation state", () => { const agentDir = mkdtempSync(join(tmpdir(), "prime-agent-telemetry-")); const path = join(agentDir, "telemetry.json"); writeFileSync(path, "not-json"); const randomId = uuidGenerator();
const installationId = getOrCreateTelemetryInstallationId(agentDir, randomId);
expect(JSON.parse(readFileSync(path, "utf8"))).toEqual({ version: 1, installationId, }); expect(getOrCreateTelemetryInstallationId(agentDir, randomId)).toBe(installationId); });
it("does not follow a telemetry state symlink", async () => { const agentDir = mkdtempSync(join(tmpdir(), "prime-agent-telemetry-")); const targetPath = join(agentDir, "target.json"); const telemetryPath = join(agentDir, "telemetry.json"); writeFileSync(targetPath, "do not overwrite"); symlinkSync(targetPath, telemetryPath); const client = new TelemetryClient({ agentDir, randomId: uuidGenerator() });
expect(() => client.capture("agent started", {})).not.toThrow(); await expect(client.flush()).resolves.toBeUndefined();
expect(readFileSync(targetPath, "utf8")).toBe("do not overwrite"); expect(lstatSync(telemetryPath).isSymbolicLink()).toBe(true); });
it("batches events through the configured Prime endpoint", async () => { const agentDir = mkdtempSync(join(tmpdir(), "prime-agent-telemetry-")); const requests: Array<{ url: string; init: RequestInit }> = []; const fetchMock: typeof fetch = async (input, init) => { requests.push({ url: String(input), init: init ?? {} }); return new Response(null, { status: 204 }); }; const client = new TelemetryClient({ agentDir, endpoint: "https://api.example.test/api/v1/agent-analytics/events", fetch: fetchMock, randomId: uuidGenerator(), now: () => Date.UTC(2026, 6, 23), });
client.capture("agent started", { version: "1.2.3", os_family: "darwin", }); await client.flush();
expect(requests).toHaveLength(1); expect(requests[0].url).toBe("https://api.example.test/api/v1/agent-analytics/events"); const body = JSON.parse(String(requests[0].init.body)) as { installation_id: string; events: TelemetryEvent[]; }; expect(body.installation_id).toMatch(/^[0-9a-f-]{36}$/); expect(body.events).toEqual([ expect.objectContaining({ name: "agent started", timestamp: "2026-07-23T00:00:00.000Z", properties: { version: "1.2.3", os_family: "darwin", }, }), ]); });
it("never throws when the analytics endpoint fails", async () => { const client = new TelemetryClient({ agentDir: mkdtempSync(join(tmpdir(), "prime-agent-telemetry-")), fetch: async () => { throw new Error("network failed"); }, randomId: uuidGenerator(), });
client.capture("agent started", {}); await expect(client.flush()).resolves.toBeUndefined(); });
it("drains every queued batch before flush resolves", async () => { const batchSizes: number[] = []; const client = new TelemetryClient({ agentDir: mkdtempSync(join(tmpdir(), "prime-agent-telemetry-")), fetch: async (_input, init) => { const body = JSON.parse(String(init?.body)) as { events: TelemetryEvent[] }; batchSizes.push(body.events.length); return new Response(null, { status: 204 }); }, randomId: uuidGenerator(), batchSize: 2, });
for (let index = 0; index < 5; index++) { client.capture("agent started", { index }); } await client.flush();
expect(batchSizes.reduce((total, size) => total + size, 0)).toBe(5); expect(batchSizes.every((size) => size <= 2)).toBe(true); });
it("never throws when the local telemetry state cannot be written", async () => { const parent = mkdtempSync(join(tmpdir(), "prime-agent-telemetry-")); const agentDir = join(parent, "not-a-directory"); writeFileSync(agentDir, "occupied"); const client = new TelemetryClient({ agentDir, randomId: uuidGenerator() });
expect(() => client.capture("agent started", {})).not.toThrow(); await expect(client.flush()).resolves.toBeUndefined(); });});
describe("telemetry controls", () => { it("honors settings and environment opt-outs", () => { const settings = SettingsManager.inMemory({ telemetry: { enabled: true } });
vi.stubEnv("DO_NOT_TRACK", "0"); expect(isTelemetryEnabled(settings)).toBe(true);
vi.stubEnv("DO_NOT_TRACK", "1"); expect(isTelemetryEnabled(settings)).toBe(false);
vi.stubEnv("DO_NOT_TRACK", "0"); vi.stubEnv("PRIME_AGENT_TELEMETRY", "0"); expect(isTelemetryEnabled(settings)).toBe(false);
vi.stubEnv("PRIME_AGENT_TELEMETRY", "1"); vi.stubEnv("PI_OFFLINE", "true"); expect(isTelemetryEnabled(settings)).toBe(false); });
it("normalizes malformed telemetry settings before updating them", async () => { const settings = SettingsManager.inMemory({ telemetry: true as never }); const disabledSettings = SettingsManager.inMemory({ telemetry: false as never });
expect(disabledSettings.getTelemetryEnabled()).toBe(false); expect(() => settings.setTelemetryNoticeShown(true)).not.toThrow(); expect(() => settings.setTelemetryEnabled(false)).not.toThrow(); await settings.flush();
expect(settings.getGlobalSettings().telemetry).toEqual({ noticeShown: true, enabled: false, }); });});
describe("agent telemetry aggregation", () => { beforeEach(() => { vi.stubEnv("DO_NOT_TRACK", "0"); });
it("captures only allowlisted built-in command names", async () => { vi.stubEnv("PRIME_AGENT_TELEMETRY", "1"); const sink = new FakeTelemetrySink();
await captureAgentCommandUsed({ agentDir: "/not-used", settingsManager: SettingsManager.inMemory(), commandName: "model", sink, }); await captureAgentCommandUsed({ agentDir: "/not-used", settingsManager: SettingsManager.inMemory(), commandName: "private-extension-command", sink, });
expect(sink.events).toContainEqual({ name: "agent command used", properties: expect.objectContaining({ execution_mode: "interactive", command_name: "model", }), }); expect(sink.events).toHaveLength(1); });
it("captures onboarding completion with categorized auth and provider data", async () => { vi.stubEnv("PRIME_AGENT_TELEMETRY", "1"); const sink = new FakeTelemetrySink();
await captureOnboardingCompleted({ agentDir: "/not-used", settingsManager: SettingsManager.inMemory(), durationMs: 250, outcome: "success", provider: "prime", authSource: "stored", storedCredentialType: "oauth", sink, });
expect(sink.events).toContainEqual({ name: "onboarding completed", properties: expect.objectContaining({ execution_mode: "interactive", duration_ms: 250, outcome: "success", auth_category: "oauth", provider_category: "prime", }), }); expect(sink.flushCount).toBe(1); expect(telemetryAuthCategory("models_json_command")).toBe("models_json"); });
it("emits aggregate metrics without message or tool content", () => { vi.stubEnv("PRIME_AGENT_TELEMETRY", "1"); let timestamp = 1_000; const now = () => timestamp; const randomId = uuidGenerator(); const sink = new FakeTelemetrySink(); const fakeSession = new FakeAgentSession();
installAgentTelemetry(fakeSession as unknown as AgentSession, { agentDir: "/not-used", settingsManager: SettingsManager.inMemory(), executionMode: "interactive", sink, now, randomId, });
const assistant = assistantMessage(); fakeSession.emit({ type: "session_action_update", actions: { queuedCount: 0, steering: [], followUps: [], active: { kind: "turn", phase: "running" } }, }); fakeSession.emit({ type: "agent_start" }); fakeSession.emit({ type: "message_start", message: { role: "user", content: "private prompt", timestamp, }, }); timestamp = 1_010; fakeSession.emit({ type: "turn_start" }); timestamp = 1_035; fakeSession.emit({ type: "message_update", message: assistant, assistantMessageEvent: { type: "text_delta", contentIndex: 0, delta: "private streamed text", partial: assistant, }, }); timestamp = 1_050; fakeSession.emit({ type: "tool_execution_end", toolCallId: "private-tool-id", toolName: "bash", result: { private: "tool output" }, isError: false, }); timestamp = 1_100; fakeSession.emit({ type: "message_end", message: assistant }); fakeSession.emit({ type: "turn_end", message: assistant, toolResults: [] }); timestamp = 1_125; fakeSession.emit({ type: "agent_end", messages: [assistant] }); fakeSession.emit({ type: "session_action_update", actions: { queuedCount: 0, steering: [], followUps: [] } });
const run = sink.events.find((event) => event.name === "agent run completed"); expect(run?.properties).toMatchObject({ execution_mode: "interactive", outcome: "success", duration_ms: 125, visible_ttft_ms: 25, first_model_event_ms: 25, model_latency_ms: 90, turn_count: 1, tool_call_count: 1, tool_error_count: 0, input_tokens: 100, output_tokens: 20, cache_read_tokens: 50, total_tokens: 170, retry_count: 0, provider_category: "openai", model_category: "gpt", }); expect(JSON.stringify(run)).not.toContain("private");
timestamp = 1_200; fakeSession.dispose(); const ended = sink.events.find((event) => event.name === "agent session ended"); expect(ended?.properties).toMatchObject({ duration_ms: 200, prompt_count: 1, run_count: 1, successful_run_count: 1, total_tokens: 170, }); expect(sink.flushCount).toBe(1); });
it("waits for post-run compaction before finalizing run metrics", () => { vi.stubEnv("PRIME_AGENT_TELEMETRY", "1"); const sink = new FakeTelemetrySink(); const fakeSession = new FakeAgentSession();
installAgentTelemetry(fakeSession as unknown as AgentSession, { agentDir: "/not-used", settingsManager: SettingsManager.inMemory(), sink, });
const assistant = assistantMessage(); fakeSession.emit({ type: "session_action_update", actions: { queuedCount: 0, steering: [], followUps: [], active: { kind: "turn", phase: "running" } }, }); fakeSession.emit({ type: "agent_start" }); fakeSession.emit({ type: "message_end", message: assistant }); fakeSession.emit({ type: "agent_end", messages: [assistant] }); expect(sink.events.find((event) => event.name === "agent run completed")).toBeUndefined();
fakeSession.emit({ type: "compaction_end", reason: "threshold", result: { summary: "private", firstKeptEntryId: "private", tokensBefore: 100 }, aborted: false, willRetry: false, }); fakeSession.emit({ type: "session_action_update", actions: { queuedCount: 0, steering: [], followUps: [] } });
expect(sink.events.find((event) => event.name === "agent run completed")?.properties.compaction_count).toBe(1); });
it("keeps automatic retries in one completed run", () => { vi.stubEnv("PRIME_AGENT_TELEMETRY", "1"); const sink = new FakeTelemetrySink(); const fakeSession = new FakeAgentSession();
installAgentTelemetry(fakeSession as unknown as AgentSession, { agentDir: "/not-used", settingsManager: SettingsManager.inMemory(), sink, });
const failed = assistantMessage({ stopReason: "error", errorMessage: "rate limit" }); const succeeded = assistantMessage(); fakeSession.emit({ type: "session_action_update", actions: { queuedCount: 0, steering: [], followUps: [], active: { kind: "turn", phase: "running" } }, }); fakeSession.emit({ type: "agent_start" }); fakeSession.emit({ type: "message_end", message: failed }); fakeSession.emit({ type: "agent_end", messages: [failed] }); fakeSession.emit({ type: "auto_retry_start", attempt: 1, maxAttempts: 3, delayMs: 100, errorMessage: "rate limit", }); fakeSession.emit({ type: "agent_start" }); fakeSession.emit({ type: "message_end", message: succeeded }); fakeSession.emit({ type: "agent_end", messages: [succeeded] }); fakeSession.emit({ type: "session_action_update", actions: { queuedCount: 0, steering: [], followUps: [] } });
const runs = sink.events.filter((event) => event.name === "agent run completed"); expect(runs).toHaveLength(1); expect(runs[0].properties).toMatchObject({ outcome: "success", retry_count: 1, model_call_count: 2 }); });
it("awaits the final telemetry flush during async session disposal", async () => { vi.stubEnv("PRIME_AGENT_TELEMETRY", "1"); const sink = new FakeTelemetrySink(); const fakeSession = new FakeAgentSession(); let releaseFlush: () => void = () => {}; const flushGate = new Promise<void>((resolve) => { releaseFlush = resolve; }); let flushStarted = false; vi.spyOn(sink, "flush").mockImplementation(async () => { flushStarted = true; await flushGate; });
installAgentTelemetry(fakeSession as unknown as AgentSession, { agentDir: "/not-used", settingsManager: SettingsManager.inMemory(), sink, });
let disposed = false; const disposal = fakeSession.disposeAsync().then(() => { disposed = true; }); await vi.waitFor(() => expect(flushStarted).toBe(true)); expect(disposed).toBe(false); releaseFlush(); await disposal; expect(disposed).toBe(true); });});