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.
22 kB · 791 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792import { randomUUID } from "node:crypto";import { lstatSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";import { arch, platform } from "node:os";import { join } from "node:path";import type { AssistantMessage, Usage } from "@earendil-works/pi-ai";import { detectInstallMethod, VERSION } from "../config.js";import { writeFileAtomicSync } from "../utils/atomic-file.js";import type { AgentSession, AgentSessionEvent } from "./agent-session.js";import type { AgentExecutionMode } from "./agent-session-config.js";import type { AuthCredential, AuthStatus } from "./auth-storage.js";import { platformFidelity } from "./platform-fidelity.js";import type { SettingsManager } from "./settings-manager.js";import { isBuiltinSlashCommandName, resolveBuiltinSlashCommandName } from "./slash-commands.js";
const DEFAULT_TELEMETRY_ENDPOINT = "https://api.primeintellect.ai/api/v1/agent-analytics/events";const TELEMETRY_STATE_FILE = "telemetry.json";const TELEMETRY_STATE_VERSION = 1;const DEFAULT_BATCH_SIZE = 10;const DEFAULT_FLUSH_INTERVAL_MS = 10_000;const DEFAULT_REQUEST_TIMEOUT_MS = 1_500;
type TelemetryPrimitive = string | number | boolean | null;type TelemetryProperties = Record<string, TelemetryPrimitive>;
export type TelemetryEventName = | "agent started" | "onboarding completed" | "agent command used" | "agent run completed" | "agent session ended";
export type TelemetryExecutionMode = AgentExecutionMode | "unknown";export type TelemetryOnboardingOutcome = "success" | "error" | "aborted";export type TelemetryAuthCategory = | "oauth" | "api_key" | "runtime_api_key" | "environment" | "prime_cli" | "models_json" | "fallback" | "stale" | "stored" | "none";
export interface TelemetryEvent { id: string; name: TelemetryEventName; timestamp: string; properties: TelemetryProperties;}
export interface TelemetryBatch { installation_id: string; events: TelemetryEvent[];}
export interface TelemetrySink { capture(name: TelemetryEventName, properties: TelemetryProperties): void; flush(): Promise<void>;}
interface TelemetryState { version: number; installationId: string;}
interface TelemetryClientOptions { agentDir: string; endpoint?: string; fetch?: typeof fetch; now?: () => number; randomId?: () => string; batchSize?: number; flushIntervalMs?: number; requestTimeoutMs?: number;}
interface InstallAgentTelemetryOptions { agentDir: string; settingsManager: SettingsManager; executionMode?: AgentExecutionMode; sink?: TelemetrySink; now?: () => number; randomId?: () => string;}
export interface CaptureOnboardingCompletedOptions { agentDir: string; settingsManager: SettingsManager; durationMs: number; outcome: TelemetryOnboardingOutcome; provider?: string; authSource?: AuthStatus["source"]; storedCredentialType?: AuthCredential["type"]; sink?: TelemetrySink; now?: () => number; randomId?: () => string;}
export interface CaptureAgentCommandUsedOptions { agentDir: string; settingsManager: SettingsManager; commandName: string; sink?: TelemetrySink; now?: () => number; randomId?: () => string;}
interface UsageTotals extends Usage { modelCallCount: number;}
interface ActiveRun { startedAt: number; agentEnded: boolean; firstTurnStartedAt?: number; firstModelEventMs?: number; visibleTtftMs?: number; currentTurnStartedAt?: number; modelLatencyMs: number; maxModelLatencyMs: number; turnCount: number; toolCallCount: number; toolErrorCount: number; compactionCount: number; retryCount: number; usage: UsageTotals; lastAssistant?: AssistantMessage;}
interface SessionTotals { startedAt: number; runCount: number; successfulRunCount: number; failedRunCount: number; abortedRunCount: number; promptCount: number; toolCallCount: number; compactionCount: number; usage: UsageTotals;}
const EMPTY_USAGE_TOTALS: UsageTotals = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0, }, modelCallCount: 0,};
function newUsageTotals(): UsageTotals { return structuredClone(EMPTY_USAGE_TOTALS);}
function addUsage(target: UsageTotals, usage: Usage): void { target.input += usage.input; target.output += usage.output; target.cacheRead += usage.cacheRead; target.cacheWrite += usage.cacheWrite; target.totalTokens += usage.totalTokens; target.cost.input += usage.cost.input; target.cost.output += usage.cost.output; target.cost.cacheRead += usage.cost.cacheRead; target.cost.cacheWrite += usage.cost.cacheWrite; target.cost.total += usage.cost.total; target.modelCallCount++;}
function mergeUsage(target: UsageTotals, usage: UsageTotals): void { target.input += usage.input; target.output += usage.output; target.cacheRead += usage.cacheRead; target.cacheWrite += usage.cacheWrite; target.totalTokens += usage.totalTokens; target.cost.input += usage.cost.input; target.cost.output += usage.cost.output; target.cost.cacheRead += usage.cost.cacheRead; target.cost.cacheWrite += usage.cost.cacheWrite; target.cost.total += usage.cost.total; target.modelCallCount += usage.modelCallCount;}
function parseBooleanOverride(value: string | undefined): boolean | undefined { if (value === undefined) { return undefined; } const normalized = value.trim().toLowerCase(); if (["1", "true", "yes", "on"].includes(normalized)) { return true; } if (["0", "false", "no", "off"].includes(normalized)) { return false; } return undefined;}
export function isTelemetryEnabled(settingsManager: SettingsManager): boolean { if (parseBooleanOverride(process.env.PI_OFFLINE) === true) { return false; } if (parseBooleanOverride(process.env.DO_NOT_TRACK) === true) { return false; } const override = parseBooleanOverride(process.env.PRIME_AGENT_TELEMETRY); if (override !== undefined) { return override; } return settingsManager.getTelemetryEnabled();}
function isInstallationId(value: unknown): value is string { return ( typeof value === "string" && /^[0-9a-f]{8}-[0-9a-f]{4}-[1-8][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i.test(value) );}
function readInstallationId(path: string): string | undefined { try { const parsed = JSON.parse(readFileSync(path, "utf8")) as Partial<TelemetryState>; return parsed.version === TELEMETRY_STATE_VERSION && isInstallationId(parsed.installationId) ? parsed.installationId : undefined; } catch { return undefined; }}
function writeTelemetryStateAtomically(path: string, state: TelemetryState): void { writeFileAtomicSync(path, JSON.stringify(state, null, 2), { mode: 0o600 });}
export function getOrCreateTelemetryInstallationId(agentDir: string, randomId: () => string = randomUUID): string { const path = join(agentDir, TELEMETRY_STATE_FILE); let replaceInvalidState = false; try { const stats = lstatSync(path); if (!stats.isFile()) { throw new Error("Telemetry state path must be a regular file"); } const existing = readInstallationId(path); if (existing) { return existing; } replaceInvalidState = true; } catch (error) { const code = typeof error === "object" && error !== null && "code" in error ? String((error as { code?: unknown }).code) : undefined; if (code !== "ENOENT") { throw error; } }
const installationId = randomId(); if (!isInstallationId(installationId)) { throw new Error("Telemetry installation ID generator returned an invalid UUID"); } mkdirSync(agentDir, { recursive: true }); const state: TelemetryState = { version: TELEMETRY_STATE_VERSION, installationId, }; try { if (replaceInvalidState) { writeTelemetryStateAtomically(path, state); } else { writeFileSync(path, JSON.stringify(state, null, 2), { encoding: "utf8", flag: "wx", mode: 0o600, }); } return installationId; } catch (error) { const code = typeof error === "object" && error !== null && "code" in error ? String((error as { code?: unknown }).code) : undefined; if (code !== "EEXIST") { throw error; } return readInstallationId(path) ?? installationId; }}
export class TelemetryClient implements TelemetrySink { private readonly endpoint: string; private readonly fetchImpl: typeof fetch; private readonly now: () => number; private readonly randomId: () => string; private readonly batchSize: number; private readonly flushIntervalMs: number; private readonly requestTimeoutMs: number; private installationId?: string; private queue: TelemetryEvent[] = []; private flushTimer?: ReturnType<typeof setTimeout>; private flushInFlight?: Promise<void>; private disabled = false;
constructor(private readonly options: TelemetryClientOptions) { this.endpoint = options.endpoint ?? process.env.PRIME_AGENT_TELEMETRY_ENDPOINT ?? DEFAULT_TELEMETRY_ENDPOINT; this.fetchImpl = options.fetch ?? fetch; this.now = options.now ?? Date.now; this.randomId = options.randomId ?? randomUUID; this.batchSize = options.batchSize ?? DEFAULT_BATCH_SIZE; this.flushIntervalMs = options.flushIntervalMs ?? DEFAULT_FLUSH_INTERVAL_MS; this.requestTimeoutMs = options.requestTimeoutMs ?? DEFAULT_REQUEST_TIMEOUT_MS; }
capture(name: TelemetryEventName, properties: TelemetryProperties): void { if (this.disabled) { return; } try { this.installationId ??= getOrCreateTelemetryInstallationId(this.options.agentDir, this.randomId); this.queue.push({ id: this.randomId(), name, timestamp: new Date(this.now()).toISOString(), properties, }); } catch { this.disabled = true; this.queue = []; return; } if (this.queue.length >= this.batchSize) { void this.flush(); return; } this.scheduleFlush(); }
private scheduleFlush(): void { if (this.flushTimer) { return; } this.flushTimer = setTimeout(() => { this.flushTimer = undefined; void this.flush(); }, this.flushIntervalMs); this.flushTimer.unref?.(); }
async flush(): Promise<void> { if (this.flushTimer) { clearTimeout(this.flushTimer); this.flushTimer = undefined; } if (this.flushInFlight) { await this.flushInFlight; if (this.queue.length > 0) { await this.flush(); } return; } if (this.queue.length === 0 || !this.installationId) { return; }
const drain = this.drainQueue(); this.flushInFlight = drain; try { await drain; } finally { if (this.flushInFlight === drain) { this.flushInFlight = undefined; } } if (this.queue.length > 0) { await this.flush(); } }
private async drainQueue(): Promise<void> { while (this.queue.length > 0 && this.installationId) { const events = this.queue.splice(0, this.batchSize); const batch: TelemetryBatch = { installation_id: this.installationId, events, }; await this.send(batch); } }
private async send(batch: TelemetryBatch): Promise<void> { try { await this.fetchImpl(this.endpoint, { method: "POST", headers: { "content-type": "application/json", "user-agent": `prime-agent/${VERSION}`, }, body: JSON.stringify(batch), signal: AbortSignal.timeout(this.requestTimeoutMs), }); } catch { // Product analytics is best-effort and must never affect the agent. } }}
function modelCategory(model: string): string { const normalized = model.toLowerCase(); const categories = [ "claude", "gpt", "o1", "o3", "o4", "gemini", "glm", "kimi", "qwen", "deepseek", "llama", "mistral", ]; return categories.find((category) => normalized.includes(category)) ?? "custom";}
export function telemetryProviderCategory(provider: string | undefined): string { if (!provider) { return "unknown"; } const normalized = provider.toLowerCase(); const categories = [ "anthropic", "openai", "google", "prime", "openrouter", "bedrock", "vertex", "mistral", "groq", "xai", ]; return categories.find((category) => normalized.includes(category)) ?? "custom";}
export function telemetryAuthCategory( source: AuthStatus["source"], storedCredentialType?: AuthCredential["type"],): TelemetryAuthCategory { switch (source) { case "stored": // A pasted MCP static token is a stored credential like any other; // the category taxonomy stays model-provider-shaped. if (storedCredentialType === "mcp_static_token") return "stored"; return storedCredentialType ?? "stored"; case "runtime": return "runtime_api_key"; case "environment": return "environment"; case "prime_cli": return "prime_cli"; case "models_json_key": case "models_json_command": return "models_json"; case "fallback": return "fallback"; case "stale": return "stale"; default: return "none"; }}
function baseProperties(executionMode: TelemetryExecutionMode): TelemetryProperties { const fidelity = platformFidelity(); return { version: VERSION, os_family: platform(), architecture: arch(), install_method: detectInstallMethod(), execution_mode: executionMode, libc: fidelity.libc, libc_version: fidelity.libc_version, cpu_baseline: fidelity.cpu_baseline, os_release: fidelity.os_release, os_product_version: fidelity.os_product_version, };}
export async function captureOnboardingCompleted(options: CaptureOnboardingCompletedOptions): Promise<void> { if (!isTelemetryEnabled(options.settingsManager)) { return; }
const now = options.now ?? Date.now; const randomId = options.randomId ?? randomUUID; const sink = options.sink ?? new TelemetryClient({ agentDir: options.agentDir, now, randomId, }); sink.capture("onboarding completed", { ...baseProperties("interactive"), duration_ms: Math.max(0, options.durationMs), outcome: options.outcome, auth_category: telemetryAuthCategory(options.authSource, options.storedCredentialType), provider_category: telemetryProviderCategory(options.provider), }); await sink.flush();}
export async function captureAgentCommandUsed(options: CaptureAgentCommandUsedOptions): Promise<void> { if (!isBuiltinSlashCommandName(options.commandName) || !isTelemetryEnabled(options.settingsManager)) { return; } const commandName = resolveBuiltinSlashCommandName(options.commandName);
const now = options.now ?? Date.now; const randomId = options.randomId ?? randomUUID; const sink = options.sink ?? new TelemetryClient({ agentDir: options.agentDir, now, randomId, }); sink.capture("agent command used", { ...baseProperties("interactive"), command_name: commandName, }); await sink.flush();}
function errorCategory(message: AssistantMessage | undefined): string | null { if (!message || message.stopReason !== "error") { return null; } const error = message.errorMessage?.toLowerCase() ?? ""; if (/\b401\b|\b403\b|auth|api.?key|credential|unauthori[sz]ed|forbidden/.test(error)) { return "authentication"; } if (/\b429\b|rate.?limit|quota/.test(error)) { return "rate_limit"; } if (/timeout|timed out/.test(error)) { return "timeout"; } if (/context|token.*limit|too long|maximum.*length/.test(error)) { return "context_limit"; } if (/network|socket|connection|fetch/.test(error)) { return "network"; } if (/\b5\d\d\b|overload|unavailable/.test(error)) { return "provider_unavailable"; } return "other";}
function runOutcome(message: AssistantMessage | undefined): "success" | "error" | "aborted" { if (message?.stopReason === "aborted") { return "aborted"; } if (!message || message.stopReason === "error") { return "error"; } return "success";}
function assistantMessage(event: AgentSessionEvent): AssistantMessage | undefined { if (event.type !== "message_end" || event.message.role !== "assistant") { return undefined; } return event.message;}
function createActiveRun(now: () => number): ActiveRun { return { startedAt: now(), agentEnded: false, modelLatencyMs: 0, maxModelLatencyMs: 0, turnCount: 0, toolCallCount: 0, toolErrorCount: 0, compactionCount: 0, retryCount: 0, usage: newUsageTotals(), };}
export function installAgentTelemetry(session: AgentSession, options: InstallAgentTelemetryOptions): void { if (!isTelemetryEnabled(options.settingsManager)) { return; }
const now = options.now ?? Date.now; const randomId = options.randomId ?? randomUUID; const sink = options.sink ?? new TelemetryClient({ agentDir: options.agentDir, now, randomId, }); const sessionId = randomId(); const sessionTotals: SessionTotals = { startedAt: now(), runCount: 0, successfulRunCount: 0, failedRunCount: 0, abortedRunCount: 0, promptCount: 0, toolCallCount: 0, compactionCount: 0, usage: newUsageTotals(), }; let activeRun: ActiveRun | undefined; let turnActionActive = false;
const commonProperties = (): TelemetryProperties => ({ ...baseProperties(options.executionMode ?? "unknown"), session_id: sessionId, });
const finalizeRun = (): void => { const run = activeRun; if (!run) { return; } activeRun = undefined; const outcome = runOutcome(run.lastAssistant); sessionTotals.runCount++; sessionTotals.toolCallCount += run.toolCallCount; sessionTotals.compactionCount += run.compactionCount; if (outcome === "success") { sessionTotals.successfulRunCount++; } else if (outcome === "aborted") { sessionTotals.abortedRunCount++; } else { sessionTotals.failedRunCount++; } mergeUsage(sessionTotals.usage, run.usage); const lastAssistant = run.lastAssistant; sink.capture("agent run completed", { ...commonProperties(), outcome, duration_ms: Math.max(0, now() - run.startedAt), visible_ttft_ms: run.visibleTtftMs ?? null, first_model_event_ms: run.firstModelEventMs ?? null, model_latency_ms: run.modelLatencyMs, max_model_latency_ms: run.maxModelLatencyMs, model_call_count: run.usage.modelCallCount, turn_count: run.turnCount, tool_call_count: run.toolCallCount, tool_error_count: run.toolErrorCount, input_tokens: run.usage.input, output_tokens: run.usage.output, cache_read_tokens: run.usage.cacheRead, cache_write_tokens: run.usage.cacheWrite, total_tokens: run.usage.totalTokens, compaction_count: run.compactionCount, retry_count: run.retryCount, provider_category: telemetryProviderCategory(lastAssistant?.provider), model_category: lastAssistant ? modelCategory(lastAssistant.model) : "unknown", error_category: errorCategory(lastAssistant), }); };
sink.capture("agent started", commonProperties());
const unsubscribe = session.subscribe((event) => { switch (event.type) { case "session_action_update": turnActionActive = event.actions.active?.kind === "turn"; if (!turnActionActive && activeRun?.agentEnded) { finalizeRun(); } break; case "agent_start": if (activeRun?.agentEnded && !turnActionActive) { finalizeRun(); } activeRun ??= createActiveRun(now); activeRun.agentEnded = false; break; case "message_start": if (event.message.role === "user") { sessionTotals.promptCount++; } break; case "turn_start": if (activeRun) { activeRun.currentTurnStartedAt = now(); activeRun.firstTurnStartedAt ??= activeRun.currentTurnStartedAt; activeRun.turnCount++; } break; case "message_update": if (!activeRun) { break; } if (activeRun.firstModelEventMs === undefined && activeRun.firstTurnStartedAt !== undefined) { activeRun.firstModelEventMs = Math.max(0, now() - activeRun.firstTurnStartedAt); } if ( activeRun.visibleTtftMs === undefined && activeRun.firstTurnStartedAt !== undefined && event.assistantMessageEvent.type === "text_delta" && event.assistantMessageEvent.delta.length > 0 ) { activeRun.visibleTtftMs = Math.max(0, now() - activeRun.firstTurnStartedAt); } break; case "message_end": { const message = assistantMessage(event); if (!activeRun || !message) { break; } activeRun.lastAssistant = message; addUsage(activeRun.usage, message.usage); if (activeRun.currentTurnStartedAt !== undefined) { const latency = Math.max(0, now() - activeRun.currentTurnStartedAt); activeRun.modelLatencyMs += latency; activeRun.maxModelLatencyMs = Math.max(activeRun.maxModelLatencyMs, latency); activeRun.currentTurnStartedAt = undefined; } break; } case "tool_execution_end": if (activeRun) { activeRun.toolCallCount++; if (event.isError) { activeRun.toolErrorCount++; } } break; case "compaction_end": if (activeRun && event.result && !event.aborted) { activeRun.compactionCount++; } break; case "auto_retry_start": if (activeRun) { activeRun.retryCount++; } break; case "agent_end": if (!activeRun) { break; } activeRun.agentEnded = true; if (!turnActionActive) { finalizeRun(); } break; } });
session.registerDisposeCallback(async () => { unsubscribe(); finalizeRun(); sink.capture("agent session ended", { ...commonProperties(), duration_ms: Math.max(0, now() - sessionTotals.startedAt), prompt_count: sessionTotals.promptCount, run_count: sessionTotals.runCount, successful_run_count: sessionTotals.successfulRunCount, failed_run_count: sessionTotals.failedRunCount, aborted_run_count: sessionTotals.abortedRunCount, tool_call_count: sessionTotals.toolCallCount, compaction_count: sessionTotals.compactionCount, model_call_count: sessionTotals.usage.modelCallCount, input_tokens: sessionTotals.usage.input, output_tokens: sessionTotals.usage.output, cache_read_tokens: sessionTotals.usage.cacheRead, cache_write_tokens: sessionTotals.usage.cacheWrite, total_tokens: sessionTotals.usage.totalTokens, }); await sink.flush(); });}