From 44204e0424a1799b29248fc86cc9353a974806c4 Mon Sep 17 00:00:00 2001 From: Mikael Hugo Date: Sat, 2 May 2026 11:46:37 +0200 Subject: [PATCH] chore(sf): add optional token telemetry --- src/headless.ts | 50 +++++--- src/tests/headless-token-telemetry.test.ts | 132 +++++++++++++++++++++ 2 files changed, 167 insertions(+), 15 deletions(-) create mode 100644 src/tests/headless-token-telemetry.test.ts diff --git a/src/headless.ts b/src/headless.ts index 99bd17b2a..9ba22e807 100644 --- a/src/headless.ts +++ b/src/headless.ts @@ -11,15 +11,22 @@ */ import type { ChildProcess } from "node:child_process"; -import { existsSync, mkdirSync, readFileSync, renameSync, unlinkSync, writeFileSync } from "node:fs"; +import { + existsSync, + mkdirSync, + readFileSync, + renameSync, + unlinkSync, + writeFileSync, +} from "node:fs"; import { join, resolve } from "node:path"; import type { SessionInfo } from "@singularity-forge/pi-coding-agent"; import { RpcClient, SessionManager } from "@singularity-forge/pi-coding-agent"; +import { error, formatStructuredError } from "./errors.js"; import { AnswerInjector, loadAndValidateAnswerFile, } from "./headless-answers.js"; -import { error, formatStructuredError } from "./errors.js"; import { bootstrapProject, buildAutoBootstrapContext, @@ -448,7 +455,9 @@ export async function runHeadless(options: HeadlessOptions): Promise { } } // No sentinel or read failed — treat as normal restart - process.stderr.write("[headless] Reload: sentinel not found, starting fresh\n"); + process.stderr.write( + "[headless] Reload: sentinel not found, starting fresh\n", + ); } // Crash/error — check if we should restart @@ -587,7 +596,10 @@ async function runHeadlessOnce( formatStructuredError( error("Failed to load context for new-milestone", { operation: "loadContext", - file: options.context === "-" ? "stdin" : resolve(options.context ?? ""), + file: + options.context === "-" + ? "stdin" + : resolve(options.context ?? ""), guidance: 'Use --context-text "..." for inline context, or verify the file path', cause: err, @@ -1227,6 +1239,11 @@ async function runHeadlessOnce( inputTokens: tokens?.input ?? 0, outputTokens: tokens?.output ?? 0, }; + if (process.env.PI_TOKEN_TELEMETRY === "1") { + process.stderr.write( + `[PI_TOKEN] input=${tokens?.input ?? 0} output=${tokens?.output ?? 0} cache_read=${tokens?.cacheRead ?? 0} cache_write=${tokens?.cacheWrite ?? 0} cost=$${Number(cumCost.costUsd ?? 0).toFixed(4)}\n`, + ); + } } } @@ -1242,7 +1259,9 @@ async function runHeadlessOnce( if (isNewMilestone && options.auto && ame?.type === "text_delta") { const deltaText = String(ame?.delta ?? ame?.text ?? ""); if (deltaText) { - milestoneDetectionBuffer = (milestoneDetectionBuffer + deltaText).slice(-200); + milestoneDetectionBuffer = ( + milestoneDetectionBuffer + deltaText + ).slice(-200); milestoneReady ||= isMilestoneReadyText(milestoneDetectionBuffer); } } @@ -1578,15 +1597,15 @@ async function runHeadlessOnce( const result = resolveResumeSession(sessions, options.resumeSession); if (result.error) { process.stderr.write( - formatStructuredError( - error(result.error, { - operation: "resolveResumeSession", - guidance: - "Use the full session ID, or run 'sf sessions' to list and select interactively", - }), - "[headless]", - ), - ); + formatStructuredError( + error(result.error, { + operation: "resolveResumeSession", + guidance: + "Use the full session ID, or run 'sf sessions' to list and select interactively", + }), + "[headless]", + ), + ); await client.stop(); if (timeoutTimer) clearTimeout(timeoutTimer); process.exit(1); @@ -1599,7 +1618,8 @@ async function runHeadlessOnce( error(`Session switch to '${matched.id}' was cancelled`, { operation: "switchSession", file: matched.path, - guidance: "Check extension logs or disable the cancelling extension", + guidance: + "Check extension logs or disable the cancelling extension", }), "[headless]", ), diff --git a/src/tests/headless-token-telemetry.test.ts b/src/tests/headless-token-telemetry.test.ts new file mode 100644 index 000000000..00d7dd9eb --- /dev/null +++ b/src/tests/headless-token-telemetry.test.ts @@ -0,0 +1,132 @@ +/** + * Tests for PI_TOKEN_TELEMETRY=1 per-call token data written to stderr. + * + * Uses extracted logic mirrors to avoid importing modules with native + * dependencies (same pattern as headless-events.test.ts). + */ + +import assert from "node:assert/strict"; +import { test } from "vitest"; + +// Extracted cost_update handler logic mirrors headless.ts. + +interface CostUpdateEvent { + type: "cost_update"; + cumulativeCost?: { + costUsd?: number; + }; + tokens?: { + input?: number; + output?: number; + cacheRead?: number; + cacheWrite?: number; + }; +} + +function handleCostUpdateForTelemetry( + event: CostUpdateEvent, + env: { PI_TOKEN_TELEMETRY?: string }, + stderrWrites: string[], +): void { + const cumCost = event.cumulativeCost; + if (!cumCost) return; + const tokens = event.tokens; + if (env.PI_TOKEN_TELEMETRY === "1") { + stderrWrites.push( + `[PI_TOKEN] input=${tokens?.input ?? 0} output=${tokens?.output ?? 0} cache_read=${tokens?.cacheRead ?? 0} cache_write=${tokens?.cacheWrite ?? 0} cost=$${Number(cumCost.costUsd ?? 0).toFixed(4)}\n`, + ); + } +} + +// PI_TOKEN_TELEMETRY behavior contracts. + +test("PI_TOKEN_TELEMETRY=1 writes formatted token line to stderr", () => { + const writes: string[] = []; + const event: CostUpdateEvent = { + type: "cost_update", + cumulativeCost: { costUsd: 0.042 }, + tokens: { + input: 5000, + output: 1200, + cacheRead: 3000, + cacheWrite: 150, + }, + }; + + handleCostUpdateForTelemetry(event, { PI_TOKEN_TELEMETRY: "1" }, writes); + + assert.equal(writes.length, 1); + assert.ok(writes[0].startsWith("[PI_TOKEN] ")); + assert.ok(writes[0].includes("input=5000")); + assert.ok(writes[0].includes("output=1200")); + assert.ok(writes[0].includes("cache_read=3000")); + assert.ok(writes[0].includes("cache_write=150")); + assert.ok(writes[0].includes("cost=$0.0420")); +}); + +test("PI_TOKEN_TELEMETRY unset does not write to stderr", () => { + const writes: string[] = []; + const event: CostUpdateEvent = { + type: "cost_update", + cumulativeCost: { costUsd: 0.1 }, + tokens: { input: 100, output: 50 }, + }; + + handleCostUpdateForTelemetry(event, {}, writes); + + assert.equal(writes.length, 0); +}); + +test("PI_TOKEN_TELEMETRY=0 does not write to stderr", () => { + const writes: string[] = []; + const event: CostUpdateEvent = { + type: "cost_update", + cumulativeCost: { costUsd: 0.1 }, + tokens: { input: 100, output: 50 }, + }; + + handleCostUpdateForTelemetry(event, { PI_TOKEN_TELEMETRY: "0" }, writes); + + assert.equal(writes.length, 0); +}); + +test("cost_update without cumulativeCost does not write to stderr", () => { + const writes: string[] = []; + const event: CostUpdateEvent = { + type: "cost_update", + }; + + handleCostUpdateForTelemetry(event, { PI_TOKEN_TELEMETRY: "1" }, writes); + + assert.equal(writes.length, 0); +}); + +test("cost_update with missing tokens defaults to zeros", () => { + const writes: string[] = []; + const event: CostUpdateEvent = { + type: "cost_update", + cumulativeCost: { costUsd: 0.005 }, + }; + + handleCostUpdateForTelemetry(event, { PI_TOKEN_TELEMETRY: "1" }, writes); + + assert.equal(writes.length, 1); + assert.ok(writes[0].includes("input=0")); + assert.ok(writes[0].includes("output=0")); + assert.ok(writes[0].includes("cache_read=0")); + assert.ok(writes[0].includes("cache_write=0")); + assert.ok(writes[0].includes("cost=$0.0050")); +}); + +test("cost_update rounds cost to 4 decimal places", () => { + const writes: string[] = []; + const event: CostUpdateEvent = { + type: "cost_update", + cumulativeCost: { costUsd: 0.1234567 }, + tokens: { input: 1, output: 1 }, + }; + + handleCostUpdateForTelemetry(event, { PI_TOKEN_TELEMETRY: "1" }, writes); + + assert.ok(writes[0].includes("cost=$0.1235")); +});