From 33a470d3a9210826139f3de5d5734f68ea21a172 Mon Sep 17 00:00:00 2001 From: ozymandiashh <234437643+ozymandiashh@users.noreply.github.com> Date: Sat, 23 May 2026 03:01:17 +0300 Subject: [PATCH 1/2] Capture Antigravity CLI usage via statusline --- README.md | 2 + docs/providers/antigravity.md | 17 +- src/antigravity-statusline.ts | 163 ++++++++++ src/main.ts | 44 +++ src/parser.ts | 8 +- src/providers/antigravity.ts | 426 ++++++++++++++++++++++++--- tests/antigravity-statusline.test.ts | 85 ++++++ tests/providers/antigravity.test.ts | 217 +++++++++++++- 8 files changed, 913 insertions(+), 49 deletions(-) create mode 100644 src/antigravity-statusline.ts create mode 100644 tests/antigravity-statusline.test.ts diff --git a/README.md b/README.md index 64ddd5d2..98254ea0 100644 --- a/README.md +++ b/README.md @@ -136,6 +136,8 @@ The `--provider` flag filters any command to a single provider: `codeburn report **Gemini CLI** stores sessions as single JSON files. Each session embeds real token counts (input, output, cached, thoughts) per message, so no estimation is needed. Gemini reports input tokens inclusive of cached; CodeBurn subtracts cached from input before pricing to avoid double charging. +**Antigravity CLI** exposes exact usage through a short-lived local process while `agy` is running. Install the optional live hook with `codeburn antigravity-hook install` to capture short CLI sessions even when the menubar's 30-second refresh misses that window. The hook stores sanitized usage totals only, not prompts or local working-directory paths. Remove it with `codeburn antigravity-hook uninstall`; if `--force` replaced an existing statusLine command, uninstall restores that previous command. + **Mistral Vibe** stores sessions as folders under `~/.vibe/logs/session/` (or `$VIBE_HOME/logs/session/`). CodeBurn reads cumulative prompt/completion totals and model pricing from `meta.json`, then reads `messages.jsonl` for the first user prompt and assistant tool calls. Subagent sessions under `agents/` are counted as separate Vibe sessions. **Kiro** stores conversations as `.chat` JSON files. Token counts are estimated from content length. The underlying model is not exposed, so sessions are labeled `kiro-auto` and costed at Sonnet rates. diff --git a/docs/providers/antigravity.md b/docs/providers/antigravity.md index ca6d23f7..adc96e5b 100644 --- a/docs/providers/antigravity.md +++ b/docs/providers/antigravity.md @@ -23,6 +23,13 @@ builds can expose `--extension_server_port` and forms are supported. If the language server is not running, the parser falls back to the cached results file. +For Antigravity CLI (`agy`), CodeBurn can also install an opt-in status line +hook with `codeburn antigravity-hook install`. The hook records the CLI's +sanitized `context_window.current_usage` payload while `agy` is still alive, +without prompts or local working-directory paths. It also attempts a best-effort +RPC snapshot for full response metadata. Remove it with +`codeburn antigravity-hook uninstall`; if `--force` replaced an existing +statusLine command, uninstall restores that previous command. ## Storage format @@ -34,11 +41,19 @@ Custom file cache at `$CODEBURN_CACHE_DIR/antigravity-results.json` (defaults to ## Deduplication -Per `:`. +Per `:` for RPC data. The status line fallback collapses +repeated identical usage snapshots, ignores singleton intermediate snapshots +when a later stabilized usage total is observed for the same conversation, and +uses positive deltas for monotonic snapshots so cumulative counters are not +double-counted. ## Quirks - **Antigravity is the only provider that requires a live process.** A user who closes Antigravity loses the most-recent data until next launch (the cache covers older runs). +- **Antigravity CLI has a shorter capture window than the desktop app.** `agy` + exposes its language server only while the CLI session is active. The status + line hook closes that gap for future sessions; older CLI `.pb` files still + cannot be priced exactly unless an RPC snapshot was captured. - The 16 MB cap on RPC responses is necessary because individual cascades can balloon. Raising it risks OOM on the user's machine. - Token types are split across `inputTokens`, `responseOutputTokens`, and `thinkingOutputTokens`. Thinking is billed at output rate. diff --git a/src/antigravity-statusline.ts b/src/antigravity-statusline.ts new file mode 100644 index 00000000..a66cf422 --- /dev/null +++ b/src/antigravity-statusline.ts @@ -0,0 +1,163 @@ +import { mkdir, open, readFile, rename, unlink } from 'fs/promises' +import { randomBytes } from 'crypto' +import { dirname, join } from 'path' +import { homedir } from 'os' + +import { + recordAntigravityStatusLinePayload, + snapshotAntigravityStatusLinePayload, +} from './providers/antigravity.js' + +type Settings = Record & { + statusLine?: { + type?: string + command?: string + padding?: number + } +} + +type StatusLineSettings = NonNullable + +function isObject(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value) +} + +function isCodeBurnHook(command: unknown): boolean { + return typeof command === 'string' && command.includes('agy-statusline-hook') +} + +function shellQuote(value: string): string { + if (process.platform === 'win32') return `"${value.replace(/(["\\])/g, '\\$1')}"` + return `'${value.replace(/'/g, `'\\''`)}'` +} + +function hookCommand(): string { + const script = process.argv[1] || 'codeburn' + if (script === 'codeburn') return 'codeburn agy-statusline-hook' + return `${shellQuote(process.execPath)} ${shellQuote(script)} agy-statusline-hook` +} + +function settingsPath(): string { + return process.env['CODEBURN_ANTIGRAVITY_SETTINGS_PATH'] + ?? join(homedir(), '.gemini', 'antigravity-cli', 'settings.json') +} + +function codeburnCacheDir(): string { + return process.env['CODEBURN_CACHE_DIR'] ?? join(homedir(), '.cache', 'codeburn') +} + +function previousStatusLinePath(): string { + return join(codeburnCacheDir(), 'antigravity-statusline-previous.json') +} + +async function readSettings(): Promise { + try { + const raw = await readFile(settingsPath(), 'utf-8') + const parsed = JSON.parse(raw) + return isObject(parsed) ? parsed as Settings : {} + } catch { + return {} + } +} + +async function writeJsonAtomic(path: string, value: unknown): Promise { + await mkdir(dirname(path), { recursive: true }) + const tempPath = `${path}.${randomBytes(8).toString('hex')}.tmp` + const handle = await open(tempPath, 'w', 0o600) + try { + await handle.writeFile(`${JSON.stringify(value, null, 2)}\n`, { encoding: 'utf-8' }) + await handle.sync() + } finally { + await handle.close() + } + try { + await rename(tempPath, path) + } catch (err) { + try { await unlink(tempPath) } catch { /* cleanup */ } + throw err + } +} + +async function writeSettings(settings: Settings): Promise { + await writeJsonAtomic(settingsPath(), settings) +} + +async function savePreviousStatusLine(statusLine: StatusLineSettings): Promise { + await writeJsonAtomic(previousStatusLinePath(), { + savedAt: new Date().toISOString(), + statusLine, + }) +} + +async function readPreviousStatusLine(): Promise { + try { + const raw = await readFile(previousStatusLinePath(), 'utf-8') + const parsed = JSON.parse(raw) + if (!isObject(parsed) || !isObject(parsed.statusLine)) return null + return parsed.statusLine as StatusLineSettings + } catch { + return null + } +} + +async function clearPreviousStatusLine(): Promise { + try { + await unlink(previousStatusLinePath()) + } catch { /* no previous hook backup */ } +} + +export async function installAntigravityStatusLineHook(force = false): Promise<'installed' | 'already-installed'> { + const settings = await readSettings() + const existing = settings.statusLine + if (existing && !isCodeBurnHook(existing.command) && !force) { + throw new Error( + 'Antigravity CLI already has a custom statusLine command. Re-run with --force to replace it.' + ) + } + + if (isCodeBurnHook(existing?.command)) return 'already-installed' + if (existing && !isCodeBurnHook(existing.command)) await savePreviousStatusLine(existing) + + settings.statusLine = { + type: 'command', + command: hookCommand(), + padding: 0, + } + await writeSettings(settings) + return 'installed' +} + +export async function uninstallAntigravityStatusLineHook(): Promise<'removed' | 'restored' | 'not-installed'> { + const settings = await readSettings() + if (!isCodeBurnHook(settings.statusLine?.command)) return 'not-installed' + + const previous = await readPreviousStatusLine() + if (previous) settings.statusLine = previous + else delete settings.statusLine + + await writeSettings(settings) + await clearPreviousStatusLine() + return previous ? 'restored' : 'removed' +} + +function readStdin(): Promise { + return new Promise((resolve, reject) => { + let input = '' + process.stdin.setEncoding('utf8') + process.stdin.on('data', chunk => { input += chunk }) + process.stdin.on('end', () => resolve(input)) + process.stdin.on('error', reject) + }) +} + +export async function runAgyStatusLineHook(): Promise { + try { + const input = await readStdin() + const payload = input.trim() ? JSON.parse(input) : null + await recordAntigravityStatusLinePayload(payload) + await snapshotAntigravityStatusLinePayload(payload) + } catch { + // Status line hooks run inside the user's terminal UI. Never surface parser + // or transient RPC failures there; the next status line update can retry. + } +} diff --git a/src/main.ts b/src/main.ts index 48071c8e..754a24cd 100644 --- a/src/main.ts +++ b/src/main.ts @@ -17,6 +17,11 @@ import { formatDateRangeLabel, parseDateRangeFlags, getDateRange, toPeriod, type import { runOptimize, scanAndDetect } from './optimize.js' import { renderCompare } from './compare.js' import { getAllProviders } from './providers/index.js' +import { + installAntigravityStatusLineHook, + runAgyStatusLineHook, + uninstallAntigravityStatusLineHook, +} from './antigravity-statusline.js' import { clearPlan, readConfig, readPlan, readPlans, saveConfig, savePlan, getConfigFilePath, type Plan, type PlanId, type PlanProvider } from './config.js' import { clampResetDay, getPlanUsageOrNull, getPlanUsages, type PlanUsage } from './plan-usage.js' import { getPresetPlan, isPlanId, isPlanProvider, PLAN_IDS, PLAN_PROVIDERS, planDisplayName } from './plans.js' @@ -861,6 +866,45 @@ program } }) +program + .command('antigravity-hook') + .description('Install or remove exact Antigravity CLI usage capture') + .argument('', 'install or uninstall') + .option('--force', 'Replace an existing custom Antigravity CLI statusLine command') + .action(async (action: string, opts: { force?: boolean }) => { + try { + if (action === 'install') { + const result = await installAntigravityStatusLineHook(!!opts.force) + console.log(result === 'already-installed' + ? '\n Antigravity CLI usage capture is already installed.\n' + : '\n Antigravity CLI usage capture installed.\n') + return + } + if (action === 'uninstall') { + const result = await uninstallAntigravityStatusLineHook() + console.log(result === 'not-installed' + ? '\n Antigravity CLI usage capture is not installed.\n' + : result === 'restored' + ? '\n Antigravity CLI usage capture removed; previous statusLine restored.\n' + : '\n Antigravity CLI usage capture removed.\n') + return + } + console.error('\n Usage: codeburn antigravity-hook \n') + process.exit(1) + } catch (err) { + const message = err instanceof Error ? err.message : String(err) + console.error(`\n Antigravity hook failed: ${message}\n`) + process.exit(1) + } + }) + +program + .command('agy-statusline-hook', { hidden: true }) + .description('Internal Antigravity CLI statusLine hook') + .action(async () => { + await runAgyStatusLineHook() + }) + program .command('currency [code]') .description('Set display currency (e.g. codeburn currency GBP)') diff --git a/src/parser.ts b/src/parser.ts index cda0a688..91a9e524 100644 --- a/src/parser.ts +++ b/src/parser.ts @@ -4,7 +4,7 @@ import { readSessionLines } from './fs-utils.js' import { calculateCost, getShortModelName } from './models.js' import { discoverAllSessions, getProvider } from './providers/index.js' import { flushCodexCache } from './codex-cache.js' -import { antigravityCascadeIdFromPath, flushAntigravityCache } from './providers/antigravity.js' +import { antigravityCascadeIdFromPath, flushAntigravityCache, shouldReparseAntigravitySource } from './providers/antigravity.js' import { isSqliteBusyError } from './sqlite.js' import { type CachedCall, @@ -1766,10 +1766,10 @@ function getOrCreateProviderSection(cache: SessionCache, provider: string): Prov return section } -function cachedFileNeedsProviderReparse(providerName: string, cached: CachedFile): boolean { +function cachedFileNeedsProviderReparse(providerName: string, sourcePath: string, cached: CachedFile): boolean { // Antigravity data comes from the live server, not from the conversation file. // A 0-turn cache entry may just mean the server was unavailable last run. - if (providerName === 'antigravity' && cached.turns.length === 0) return true + if (providerName === 'antigravity') return shouldReparseAntigravitySource(sourcePath, cached.turns.length) if (providerName !== 'gemini') return false @@ -1815,7 +1815,7 @@ async function parseProviderSources( const cached = section.files[source.path] const action = reconcileFile(fp, cached) - if (action.action === 'unchanged' && cached && !cachedFileNeedsProviderReparse(providerName, cached)) { + if (action.action === 'unchanged' && cached && !cachedFileNeedsProviderReparse(providerName, source.path, cached)) { unchangedSources.push({ source, cached }) } else { changedSources.push({ source, fp }) diff --git a/src/providers/antigravity.ts b/src/providers/antigravity.ts index 742978ad..5392d2aa 100644 --- a/src/providers/antigravity.ts +++ b/src/providers/antigravity.ts @@ -1,4 +1,4 @@ -import { readdir, readFile, mkdir, stat, open, rename, unlink } from 'fs/promises' +import { appendFile, readdir, readFile, mkdir, stat, open, rename, unlink } from 'fs/promises' import { execFile } from 'child_process' import { randomBytes } from 'crypto' import { basename, join } from 'path' @@ -82,6 +82,38 @@ type GeneratorMetadataResponse = { } } +type StatusLineCurrentUsage = { + input_tokens?: number + output_tokens?: number + cache_creation_input_tokens?: number + cache_read_input_tokens?: number +} + +type StatusLinePayload = { + conversation_id?: string + session_id?: string + model?: string | { + id?: string + display_name?: string + } + context_window?: { + current_usage?: StatusLineCurrentUsage | null + } +} + +type StatusLineEvent = { + at: string + conversationId: string + sessionId?: string + model: string + usage: { + inputTokens: number + outputTokens: number + cacheCreationInputTokens: number + cacheReadInputTokens: number + } +} + type CachedCascade = { mtimeMs: number sizeBytes: number @@ -116,6 +148,10 @@ function getCachePath(): string { return join(getCacheDir(), 'antigravity-results.json') } +export function getAntigravityStatusLineEventsPath(): string { + return join(getCacheDir(), 'antigravity-statusline.jsonl') +} + function execFileText(command: string, args: string[], timeout = 3000): Promise { return new Promise((resolve, reject) => { execFile(command, args, { encoding: 'utf-8', timeout, maxBuffer: 1024 * 1024 }, (err, stdout) => { @@ -402,18 +438,118 @@ function normalizePricingModel(model: string): string { return PRICING_ALIASES[stripped] ?? stripped } +function parseFiniteToken(value: unknown): number { + return typeof value === 'number' && Number.isFinite(value) && value > 0 + ? Math.floor(value) + : 0 +} + +function usageSignature(event: StatusLineEvent): string { + const u = event.usage + return [ + event.model, + u.inputTokens, + u.outputTokens, + u.cacheCreationInputTokens, + u.cacheReadInputTokens, + ].join(':') +} + +function usageHasTokens(usage: StatusLineEvent['usage']): boolean { + return ( + usage.inputTokens > 0 || + usage.outputTokens > 0 || + usage.cacheCreationInputTokens > 0 || + usage.cacheReadInputTokens > 0 + ) +} + +function usageIsMonotonic(current: StatusLineEvent['usage'], previous: StatusLineEvent['usage']): boolean { + return ( + current.inputTokens >= previous.inputTokens && + current.outputTokens >= previous.outputTokens && + current.cacheCreationInputTokens >= previous.cacheCreationInputTokens && + current.cacheReadInputTokens >= previous.cacheReadInputTokens + ) +} + +function usageDelta(current: StatusLineEvent['usage'], previous: StatusLineEvent['usage']): StatusLineEvent['usage'] { + return { + inputTokens: current.inputTokens - previous.inputTokens, + outputTokens: current.outputTokens - previous.outputTokens, + cacheCreationInputTokens: current.cacheCreationInputTokens - previous.cacheCreationInputTokens, + cacheReadInputTokens: current.cacheReadInputTokens - previous.cacheReadInputTokens, + } +} + export function antigravityCascadeIdFromPath(path: string): string { return basename(path).replace(/\.(pb|db)$/i, '') } +function buildCallsFromGeneratorMetadata( + cascadeId: string, + metadata: GeneratorMetadata[], + modelMap: ModelMap, +): ParsedProviderCall[] { + const results: ParsedProviderCall[] = [] + + for (let i = 0; i < metadata.length; i++) { + const entry = metadata[i]! + const usage = entry.chatModel?.usage + if (!usage) continue + + const inputTokens = parseInt(usage.inputTokens ?? '0', 10) + const outputTokens = parseInt(usage.outputTokens ?? '0', 10) + const thinkingTokens = parseInt(usage.thinkingOutputTokens ?? '0', 10) + const responseTokens = parseInt(usage.responseOutputTokens ?? '0', 10) + + if (inputTokens === 0 && outputTokens === 0) continue + + const responseId = usage.responseId || String(i) + const dedupKey = `antigravity:${cascadeId}:${responseId}` + + const model = modelMap[usage.model] ?? usage.model + const pricingModel = normalizePricingModel(model) + const timestamp = entry.chatModel?.chatStartMetadata?.createdAt ?? '' + const costUSD = calculateCost(pricingModel, inputTokens, responseTokens + thinkingTokens, 0, 0, 0) + + results.push({ + provider: 'antigravity', + model, + inputTokens, + outputTokens: responseTokens, + cacheCreationInputTokens: 0, + cacheReadInputTokens: 0, + cachedInputTokens: 0, + reasoningTokens: thinkingTokens, + webSearchRequests: 0, + costUSD, + tools: [], + bashCommands: [], + timestamp, + speed: 'standard', + deduplicationKey: dedupKey, + userMessage: '', + sessionId: cascadeId, + }) + } + + return results +} + function isConversationFile(file: string, extensions: readonly string[]): boolean { const lowerFile = file.toLowerCase() return extensions.some(ext => lowerFile.endsWith(ext)) } +export function isAntigravityStatusLineEventsPath(path: string): boolean { + return path === getAntigravityStatusLineEventsPath() +} + export async function discoverAntigravitySessionSources( roots: readonly AntigravityConversationRoot[] = CONVERSATION_ROOTS, ): Promise { + const includeStatusLineEvents = roots === CONVERSATION_ROOTS const sources: SessionSource[] = [] for (const root of roots) { let files: string[] @@ -435,12 +571,254 @@ export async function discoverAntigravitySessionSources( }) } } + + if (includeStatusLineEvents) { + const statusLinePath = getAntigravityStatusLineEventsPath() + const statusLineStat = await stat(statusLinePath).catch(() => null) + if (statusLineStat?.isFile()) { + sources.push({ + path: statusLinePath, + project: 'antigravity-cli', + provider: 'antigravity', + }) + } + } + return sources } +function parseStatusLinePayload(input: unknown): StatusLineEvent | null { + if (!input || typeof input !== 'object') return null + const payload = input as StatusLinePayload + if (typeof payload.conversation_id !== 'string' || payload.conversation_id.length === 0) return null + const usage = payload.context_window?.current_usage + if (!usage) return null + + const event: StatusLineEvent = { + at: new Date().toISOString(), + conversationId: payload.conversation_id, + sessionId: typeof payload.session_id === 'string' ? payload.session_id : undefined, + model: typeof payload.model === 'string' + ? payload.model + : payload.model?.id ?? payload.model?.display_name ?? 'unknown', + usage: { + inputTokens: parseFiniteToken(usage.input_tokens), + outputTokens: parseFiniteToken(usage.output_tokens), + cacheCreationInputTokens: parseFiniteToken(usage.cache_creation_input_tokens), + cacheReadInputTokens: parseFiniteToken(usage.cache_read_input_tokens), + }, + } + + const u = event.usage + if (u.inputTokens === 0 && u.outputTokens === 0 && u.cacheCreationInputTokens === 0 && u.cacheReadInputTokens === 0) { + return null + } + if (event.model === 'unknown') return null + return event +} + +export async function recordAntigravityStatusLinePayload(input: unknown): Promise { + const event = parseStatusLinePayload(input) + if (!event) return false + + const path = getAntigravityStatusLineEventsPath() + await mkdir(getCacheDir(), { recursive: true }) + await appendFile(path, `${JSON.stringify(event)}\n`, { encoding: 'utf-8', mode: 0o600 }) + return true +} + +function parseStatusLineEvent(input: unknown): StatusLineEvent | null { + if (!input || typeof input !== 'object') return null + const event = input as StatusLineEvent + if (typeof event.at !== 'string' || Number.isNaN(new Date(event.at).getTime())) return null + if (typeof event.conversationId !== 'string' || event.conversationId.length === 0) return null + if (typeof event.model !== 'string' || event.model.length === 0) return null + if (!event.usage || typeof event.usage !== 'object') return null + + const usage = { + inputTokens: parseFiniteToken(event.usage.inputTokens), + outputTokens: parseFiniteToken(event.usage.outputTokens), + cacheCreationInputTokens: parseFiniteToken(event.usage.cacheCreationInputTokens), + cacheReadInputTokens: parseFiniteToken(event.usage.cacheReadInputTokens), + } + + if ( + usage.inputTokens === 0 && + usage.outputTokens === 0 && + usage.cacheCreationInputTokens === 0 && + usage.cacheReadInputTokens === 0 + ) return null + + return { + at: event.at, + conversationId: event.conversationId, + sessionId: typeof event.sessionId === 'string' ? event.sessionId : undefined, + model: event.model, + usage, + } +} + +function hasRpcCacheForConversation(seenKeys: Set, conversationId: string): boolean { + const prefix = `antigravity:${conversationId}:` + for (const key of seenKeys) { + if (key.startsWith(prefix)) return true + } + return false +} + +async function parseStatusLineCalls(source: SessionSource, seenKeys: Set): Promise { + const raw = await readFile(source.path, 'utf-8').catch(() => '') + const runsByConversation = new Map>() + + for (const line of raw.split(/\r?\n/)) { + if (!line.trim()) continue + let parsed: unknown + try { + parsed = JSON.parse(line) + } catch { + continue + } + + const event = parseStatusLineEvent(parsed) + if (!event) continue + if (hasRpcCacheForConversation(seenKeys, event.conversationId)) continue + + const signature = usageSignature(event) + const runs = runsByConversation.get(event.conversationId) ?? [] + const lastRun = runs.at(-1) + if (lastRun?.signature === signature) { + lastRun.count += 1 + lastRun.event = event + } else { + runs.push({ event, signature, count: 1 }) + runsByConversation.set(event.conversationId, runs) + } + } + + const results: ParsedProviderCall[] = [] + + for (const runs of runsByConversation.values()) { + let turnIndex = 0 + let previousSnapshotUsage: StatusLineEvent['usage'] | null = null + for (let i = 0; i < runs.length; i++) { + const run = runs[i]! + const isLastRun = i === runs.length - 1 + if (run.count === 1 && !isLastRun) continue + + const event = run.event + const signature = run.signature + const billableUsage = previousSnapshotUsage && usageIsMonotonic(event.usage, previousSnapshotUsage) + ? usageDelta(event.usage, previousSnapshotUsage) + : event.usage + previousSnapshotUsage = event.usage + if (!usageHasTokens(billableUsage)) continue + + const dedupKey = `antigravity-statusline:${event.conversationId}:${turnIndex}:${signature}` + turnIndex += 1 + if (seenKeys.has(dedupKey)) continue + + const u = billableUsage + const costUSD = calculateCost( + normalizePricingModel(event.model), + u.inputTokens, + u.outputTokens, + u.cacheCreationInputTokens, + u.cacheReadInputTokens, + 0, + ) + + results.push({ + provider: 'antigravity', + model: event.model, + inputTokens: u.inputTokens, + outputTokens: u.outputTokens, + cacheCreationInputTokens: u.cacheCreationInputTokens, + cacheReadInputTokens: u.cacheReadInputTokens, + cachedInputTokens: 0, + // StatusLine current_usage exposes aggregate output tokens, not a + // separate thinking/response split. Preserve the exact total instead + // of inventing a breakdown. + reasoningTokens: 0, + webSearchRequests: 0, + costUSD, + tools: [], + bashCommands: [], + timestamp: event.at, + speed: 'standard', + deduplicationKey: dedupKey, + userMessage: '', + sessionId: event.conversationId, + project: source.project, + }) + } + } + + return results +} + +export function shouldReparseAntigravitySource(path: string, cachedTurnCount: number): boolean { + if (cachedTurnCount === 0) return true + return isAntigravityStatusLineEventsPath(path) +} + +async function findCascadeSource(cascadeId: string): Promise { + const sources = await discoverAntigravitySessionSources() + return sources.find(source => + source.path.replace(/\\/g, '/').toLowerCase().includes('/.gemini/antigravity-cli/') && + antigravityCascadeIdFromPath(source.path) === cascadeId + ) ?? null +} + +export async function snapshotAntigravityStatusLinePayload(input: unknown): Promise { + const event = parseStatusLinePayload(input) + if (!event) return false + + const cascadeId = event.conversationId + const source = await findCascadeSource(cascadeId) + if (!source) return false + + const s = await stat(source.path).catch(() => null) + if (!s) return false + + const cache = await loadCache() + const cached = cache.cascades[cascadeId] + if (cached && cached.mtimeMs === s.mtimeMs && cached.sizeBytes === s.size && cached.calls.length > 0) { + return true + } + + const server = await detectServer('antigravity-cli') + if (!server) return false + + let metadata: GeneratorMetadata[] + try { + const modelMap = await getModelMap(server) + metadata = extractAntigravityGeneratorMetadata( + await rpc(server, 'GetCascadeTrajectoryGeneratorMetadata', { cascadeId }), + ) + cache.cascades[cascadeId] = { + mtimeMs: s.mtimeMs, + sizeBytes: s.size, + calls: buildCallsFromGeneratorMetadata(cascadeId, metadata, modelMap), + } + cacheDirty = true + await flushCache() + return cache.cascades[cascadeId]!.calls.length > 0 + } catch { + return false + } +} + function createParser(source: SessionSource, seenKeys: Set): SessionParser { return { async *parse(): AsyncGenerator { + if (isAntigravityStatusLineEventsPath(source.path)) { + for (const call of await parseStatusLineCalls(source, seenKeys)) { + seenKeys.add(call.deduplicationKey) + yield call + } + return + } + const cascadeId = antigravityCascadeIdFromPath(source.path) const cache = await loadCache() @@ -487,48 +865,7 @@ function createParser(source: SessionSource, seenKeys: Set): SessionPars return } - const results: ParsedProviderCall[] = [] - - for (let i = 0; i < metadata.length; i++) { - const entry = metadata[i]! - const usage = entry.chatModel?.usage - if (!usage) continue - - const inputTokens = parseInt(usage.inputTokens ?? '0', 10) - const outputTokens = parseInt(usage.outputTokens ?? '0', 10) - const thinkingTokens = parseInt(usage.thinkingOutputTokens ?? '0', 10) - const responseTokens = parseInt(usage.responseOutputTokens ?? '0', 10) - - if (inputTokens === 0 && outputTokens === 0) continue - - const responseId = usage.responseId || String(i) - const dedupKey = `antigravity:${cascadeId}:${responseId}` - - const model = modelMap[usage.model] ?? usage.model - const pricingModel = normalizePricingModel(model) - const timestamp = entry.chatModel?.chatStartMetadata?.createdAt ?? '' - const costUSD = calculateCost(pricingModel, inputTokens, responseTokens + thinkingTokens, 0, 0, 0) - - results.push({ - provider: 'antigravity', - model, - inputTokens, - outputTokens: responseTokens, - cacheCreationInputTokens: 0, - cacheReadInputTokens: 0, - cachedInputTokens: 0, - reasoningTokens: thinkingTokens, - webSearchRequests: 0, - costUSD, - tools: [], - bashCommands: [], - timestamp, - speed: 'standard', - deduplicationKey: dedupKey, - userMessage: '', - sessionId: cascadeId, - }) - } + const results = buildCallsFromGeneratorMetadata(cascadeId, metadata, modelMap) cache.cascades[cascadeId] = { mtimeMs: s.mtimeMs, @@ -557,6 +894,9 @@ const modelDisplayNames: Record = { 'gemini-3.5-flash-high': 'Gemini 3.5 Flash', 'gemini-3.5-flash-medium': 'Gemini 3.5 Flash', 'gemini-3.5-flash-low': 'Gemini 3.5 Flash', + 'Gemini 3.5 Flash (High)': 'Gemini 3.5 Flash', + 'Gemini 3.5 Flash (Medium)': 'Gemini 3.5 Flash', + 'Gemini 3.5 Flash (Low)': 'Gemini 3.5 Flash', 'gemini-3.1-flash-image': 'Gemini 3.1 Flash', 'gemini-3.1-flash-lite': 'Gemini 3.1 Flash Lite', 'claude-opus-4-6-thinking': 'Opus 4.6', diff --git a/tests/antigravity-statusline.test.ts b/tests/antigravity-statusline.test.ts new file mode 100644 index 00000000..0429a327 --- /dev/null +++ b/tests/antigravity-statusline.test.ts @@ -0,0 +1,85 @@ +import { mkdtemp, readFile, rm, writeFile } from 'fs/promises' +import { tmpdir } from 'os' +import { join } from 'path' +import { describe, expect, it } from 'vitest' + +import { + installAntigravityStatusLineHook, + uninstallAntigravityStatusLineHook, +} from '../src/antigravity-statusline.js' + +describe('Antigravity CLI statusLine hook installer', () => { + async function withTempSettings(run: (dir: string, settingsPath: string) => Promise) { + const dir = await mkdtemp(join(tmpdir(), 'codeburn-agy-hook-')) + const settingsPath = join(dir, 'settings.json') + const oldSettingsPath = process.env['CODEBURN_ANTIGRAVITY_SETTINGS_PATH'] + const oldCacheDir = process.env['CODEBURN_CACHE_DIR'] + process.env['CODEBURN_ANTIGRAVITY_SETTINGS_PATH'] = settingsPath + process.env['CODEBURN_CACHE_DIR'] = join(dir, 'cache') + + try { + await run(dir, settingsPath) + } finally { + if (oldSettingsPath === undefined) delete process.env['CODEBURN_ANTIGRAVITY_SETTINGS_PATH'] + else process.env['CODEBURN_ANTIGRAVITY_SETTINGS_PATH'] = oldSettingsPath + if (oldCacheDir === undefined) delete process.env['CODEBURN_CACHE_DIR'] + else process.env['CODEBURN_CACHE_DIR'] = oldCacheDir + await rm(dir, { recursive: true, force: true }) + } + } + + it('backs up and restores an existing custom statusLine when forced', async () => { + await withTempSettings(async (dir, settingsPath) => { + const customStatusLine = { + type: 'command', + command: 'custom-statusline', + padding: 1, + } + await writeFile(settingsPath, `${JSON.stringify({ statusLine: customStatusLine }, null, 2)}\n`) + + await expect(installAntigravityStatusLineHook(false)).rejects.toThrow('already has a custom statusLine') + expect(await installAntigravityStatusLineHook(true)).toBe('installed') + + const installed = JSON.parse(await readFile(settingsPath, 'utf-8')) + expect(installed.statusLine.command).toContain('agy-statusline-hook') + + const backupPath = join(dir, 'cache', 'antigravity-statusline-previous.json') + const backup = JSON.parse(await readFile(backupPath, 'utf-8')) + expect(backup.statusLine).toEqual(customStatusLine) + + expect(await uninstallAntigravityStatusLineHook()).toBe('restored') + const restored = JSON.parse(await readFile(settingsPath, 'utf-8')) + expect(restored.statusLine).toEqual(customStatusLine) + }) + }) + + it('installs CodeBurn statusLine when no statusLine exists', async () => { + await withTempSettings(async (_dir, settingsPath) => { + expect(await installAntigravityStatusLineHook(false)).toBe('installed') + expect(await installAntigravityStatusLineHook(false)).toBe('already-installed') + + const settings = JSON.parse(await readFile(settingsPath, 'utf-8')) + expect(settings.statusLine).toMatchObject({ + type: 'command', + padding: 0, + }) + expect(settings.statusLine.command).toContain('agy-statusline-hook') + }) + }) + + it('removes CodeBurn statusLine when there is no previous hook backup', async () => { + await withTempSettings(async (_dir, settingsPath) => { + await writeFile(settingsPath, JSON.stringify({ + statusLine: { + type: 'command', + command: 'codeburn agy-statusline-hook', + padding: 0, + }, + })) + + expect(await uninstallAntigravityStatusLineHook()).toBe('removed') + const settings = JSON.parse(await readFile(settingsPath, 'utf-8')) + expect(settings).not.toHaveProperty('statusLine') + }) + }) +}) diff --git a/tests/providers/antigravity.test.ts b/tests/providers/antigravity.test.ts index 1d8c5026..c1e1be47 100644 --- a/tests/providers/antigravity.test.ts +++ b/tests/providers/antigravity.test.ts @@ -1,4 +1,4 @@ -import { mkdtemp, mkdir, rm, writeFile } from 'fs/promises' +import { mkdtemp, mkdir, readFile, rm, writeFile } from 'fs/promises' import { tmpdir } from 'os' import { join } from 'path' import { describe, expect, it } from 'vitest' @@ -11,8 +11,11 @@ import { extractAntigravityAppDataDirFromLine, extractAntigravityGeneratorMetadata, extractAntigravityModelMap, + getAntigravityStatusLineEventsPath, parseAntigravityServerInfo, parseAntigravityServerInfoFromLine, + recordAntigravityStatusLinePayload, + shouldReparseAntigravitySource, } from '../../src/providers/antigravity.js' describe('antigravity provider helpers', () => { @@ -207,5 +210,217 @@ describe('antigravity provider helpers', () => { expect(provider.modelDisplayName('gemini-3.5-flash-high')).toBe('Gemini 3.5 Flash') expect(provider.modelDisplayName('gemini-3.5-flash-medium')).toBe('Gemini 3.5 Flash') expect(provider.modelDisplayName('gemini-3.5-flash-low')).toBe('Gemini 3.5 Flash') + expect(provider.modelDisplayName('Gemini 3.5 Flash (High)')).toBe('Gemini 3.5 Flash') + }) + + it('captures exact Antigravity CLI statusLine usage as fallback calls', async () => { + const dir = await mkdtemp(join(tmpdir(), 'codeburn-antigravity-statusline-')) + const oldCacheDir = process.env['CODEBURN_CACHE_DIR'] + process.env['CODEBURN_CACHE_DIR'] = dir + + try { + const payload = { + conversation_id: 'ce061468-2e2b-4c6f-bf4f-e072bd5fa986', + session_id: 'session-1', + cwd: '/workspace/project', + model: { + id: 'Gemini 3.5 Flash (High)', + display_name: 'Gemini 3.5 Flash (High)', + }, + context_window: { + current_usage: { + input_tokens: 28407, + output_tokens: 137, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + }, + }, + } + + expect(await recordAntigravityStatusLinePayload(payload)).toBe(true) + expect(await recordAntigravityStatusLinePayload(payload)).toBe(true) + + const recorded = await readFile(getAntigravityStatusLineEventsPath(), 'utf-8') + expect(recorded).not.toContain('/workspace/project') + expect(JSON.parse(recorded.split(/\r?\n/)[0]!)).not.toHaveProperty('cwd') + + const source = { + path: getAntigravityStatusLineEventsPath(), + project: 'antigravity-cli', + provider: 'antigravity', + } + + const parser = createAntigravityProvider().createSessionParser(source, new Set()) + const calls = [] + for await (const call of parser.parse()) calls.push(call) + + expect(calls).toHaveLength(1) + expect(calls[0]).toMatchObject({ + provider: 'antigravity', + model: 'Gemini 3.5 Flash (High)', + inputTokens: 28407, + outputTokens: 137, + cacheCreationInputTokens: 0, + cacheReadInputTokens: 0, + cachedInputTokens: 0, + sessionId: 'ce061468-2e2b-4c6f-bf4f-e072bd5fa986', + project: 'antigravity-cli', + }) + expect(calls[0]!.projectPath).toBeUndefined() + expect(calls[0]!.costUSD).toBeGreaterThan(0) + } finally { + if (oldCacheDir === undefined) delete process.env['CODEBURN_CACHE_DIR'] + else process.env['CODEBURN_CACHE_DIR'] = oldCacheDir + await rm(dir, { recursive: true, force: true }) + } + }) + + it('skips statusLine fallback calls when RPC cache already covered the conversation', async () => { + const dir = await mkdtemp(join(tmpdir(), 'codeburn-antigravity-statusline-rpc-dedup-')) + const oldCacheDir = process.env['CODEBURN_CACHE_DIR'] + process.env['CODEBURN_CACHE_DIR'] = dir + + try { + expect(await recordAntigravityStatusLinePayload({ + conversation_id: 'rpc-covered-conversation', + session_id: 'session-1', + model: 'Gemini 3.5 Flash (High)', + context_window: { + current_usage: { + input_tokens: 1000, + output_tokens: 100, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + }, + }, + })).toBe(true) + + const parser = createAntigravityProvider().createSessionParser({ + path: getAntigravityStatusLineEventsPath(), + project: 'antigravity-cli', + provider: 'antigravity', + }, new Set(['antigravity:rpc-covered-conversation:0'])) + + const calls = [] + for await (const call of parser.parse()) calls.push(call) + + expect(calls).toEqual([]) + } finally { + if (oldCacheDir === undefined) delete process.env['CODEBURN_CACHE_DIR'] + else process.env['CODEBURN_CACHE_DIR'] = oldCacheDir + await rm(dir, { recursive: true, force: true }) + } + }) + + it('skips singleton statusLine snapshots and deltas monotonic usage', async () => { + const dir = await mkdtemp(join(tmpdir(), 'codeburn-antigravity-statusline-runs-')) + const oldCacheDir = process.env['CODEBURN_CACHE_DIR'] + process.env['CODEBURN_CACHE_DIR'] = dir + + const basePayload = { + conversation_id: 'statusline-runs', + session_id: 'session-1', + model: 'Gemini 3.5 Flash (High)', + } + + const withUsage = ( + input_tokens: number, + output_tokens: number, + cache_read_input_tokens = 0, + ) => ({ + ...basePayload, + context_window: { + current_usage: { + input_tokens, + output_tokens, + cache_creation_input_tokens: 0, + cache_read_input_tokens, + }, + }, + }) + + try { + expect(await recordAntigravityStatusLinePayload(withUsage(100, 10))).toBe(true) + expect(await recordAntigravityStatusLinePayload(withUsage(200, 20))).toBe(true) + expect(await recordAntigravityStatusLinePayload(withUsage(200, 20))).toBe(true) + expect(await recordAntigravityStatusLinePayload(withUsage(300, 30, 50))).toBe(true) + + const parser = createAntigravityProvider().createSessionParser({ + path: getAntigravityStatusLineEventsPath(), + project: 'antigravity-cli', + provider: 'antigravity', + }, new Set()) + + const calls = [] + for await (const call of parser.parse()) calls.push(call) + + expect(calls).toHaveLength(2) + expect(calls.map(call => [call.inputTokens, call.outputTokens, call.cacheReadInputTokens])).toEqual([ + [200, 20, 0], + [100, 10, 50], + ]) + expect(calls.map(call => call.cachedInputTokens)).toEqual([0, 0]) + } finally { + if (oldCacheDir === undefined) delete process.env['CODEBURN_CACHE_DIR'] + else process.env['CODEBURN_CACHE_DIR'] = oldCacheDir + await rm(dir, { recursive: true, force: true }) + } + }) + + it('treats non-monotonic statusLine usage as a new request snapshot', async () => { + const dir = await mkdtemp(join(tmpdir(), 'codeburn-antigravity-statusline-reset-')) + const oldCacheDir = process.env['CODEBURN_CACHE_DIR'] + process.env['CODEBURN_CACHE_DIR'] = dir + + const payload = ( + input_tokens: number, + output_tokens: number, + cache_read_input_tokens = 0, + ) => ({ + conversation_id: 'statusline-reset', + session_id: 'session-1', + model: 'Gemini 3.5 Flash (High)', + context_window: { + current_usage: { + input_tokens, + output_tokens, + cache_creation_input_tokens: 0, + cache_read_input_tokens, + }, + }, + }) + + try { + expect(await recordAntigravityStatusLinePayload(payload(1000, 100))).toBe(true) + expect(await recordAntigravityStatusLinePayload(payload(1000, 100))).toBe(true) + expect(await recordAntigravityStatusLinePayload(payload(200, 30, 500))).toBe(true) + + const parser = createAntigravityProvider().createSessionParser({ + path: getAntigravityStatusLineEventsPath(), + project: 'antigravity-cli', + provider: 'antigravity', + }, new Set()) + + const calls = [] + for await (const call of parser.parse()) calls.push(call) + + expect(calls).toHaveLength(2) + expect(calls.map(call => [call.inputTokens, call.outputTokens, call.cacheReadInputTokens])).toEqual([ + [1000, 100, 0], + [200, 30, 500], + ]) + } finally { + if (oldCacheDir === undefined) delete process.env['CODEBURN_CACHE_DIR'] + else process.env['CODEBURN_CACHE_DIR'] = oldCacheDir + await rm(dir, { recursive: true, force: true }) + } + }) + + it('always reparses append-only statusLine sources but not unchanged cached cascades', () => { + const statusLinePath = getAntigravityStatusLineEventsPath() + + expect(shouldReparseAntigravitySource(statusLinePath, 1)).toBe(true) + expect(shouldReparseAntigravitySource('/tmp/antigravity/conversation.pb', 0)).toBe(true) + expect(shouldReparseAntigravitySource('/tmp/antigravity/conversation.pb', 1)).toBe(false) }) }) From 88f962d15cf7426865e99773b5247b3010ddff81 Mon Sep 17 00:00:00 2001 From: iamtoruk Date: Sun, 24 May 2026 01:27:28 -0700 Subject: [PATCH 2/2] Harden statusLine hook: cap stdin at 1MB, fix JSONL file permissions - readStdin() now rejects input exceeding 1MB to prevent OOM from a buggy or runaway caller. - Switch from appendFile (mode only on creation) to open(path, 'a', 0o600) so the JSONL file is always created with private permissions. - Set cache directory mode to 0o700. --- src/antigravity-statusline.ts | 9 ++++++++- src/providers/antigravity.ts | 11 ++++++++--- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/src/antigravity-statusline.ts b/src/antigravity-statusline.ts index a66cf422..e23314bf 100644 --- a/src/antigravity-statusline.ts +++ b/src/antigravity-statusline.ts @@ -140,11 +140,18 @@ export async function uninstallAntigravityStatusLineHook(): Promise<'removed' | return previous ? 'restored' : 'removed' } +const MAX_STDIN_BYTES = 1024 * 1024 + function readStdin(): Promise { return new Promise((resolve, reject) => { let input = '' + let bytes = 0 process.stdin.setEncoding('utf8') - process.stdin.on('data', chunk => { input += chunk }) + process.stdin.on('data', chunk => { + bytes += Buffer.byteLength(chunk, 'utf8') + if (bytes > MAX_STDIN_BYTES) { process.stdin.destroy(); reject(new Error('stdin too large')); return } + input += chunk + }) process.stdin.on('end', () => resolve(input)) process.stdin.on('error', reject) }) diff --git a/src/providers/antigravity.ts b/src/providers/antigravity.ts index 5392d2aa..031af280 100644 --- a/src/providers/antigravity.ts +++ b/src/providers/antigravity.ts @@ -1,4 +1,4 @@ -import { appendFile, readdir, readFile, mkdir, stat, open, rename, unlink } from 'fs/promises' +import { readdir, readFile, mkdir, stat, open, rename, unlink } from 'fs/promises' import { execFile } from 'child_process' import { randomBytes } from 'crypto' import { basename, join } from 'path' @@ -622,8 +622,13 @@ export async function recordAntigravityStatusLinePayload(input: unknown): Promis if (!event) return false const path = getAntigravityStatusLineEventsPath() - await mkdir(getCacheDir(), { recursive: true }) - await appendFile(path, `${JSON.stringify(event)}\n`, { encoding: 'utf-8', mode: 0o600 }) + await mkdir(getCacheDir(), { recursive: true, mode: 0o700 }) + const fd = await open(path, 'a', 0o600) + try { + await fd.appendFile(`${JSON.stringify(event)}\n`, { encoding: 'utf-8' }) + } finally { + await fd.close() + } return true }