diff --git a/src/providers/antigravity.ts b/src/providers/antigravity.ts index 946553c1..710b4bf7 100644 --- a/src/providers/antigravity.ts +++ b/src/providers/antigravity.ts @@ -7,7 +7,7 @@ import { fileURLToPath } from 'url' import https from 'https' import { calculateCost } from '../models.js' -import { isSqliteAvailable, openDatabase } from '../sqlite.js' +import { isSqliteAvailable, isSqliteBusyError, openDatabase } from '../sqlite.js' import type { Provider, SessionSource, SessionParser, ParsedProviderCall } from './types.js' type AntigravityConversationRoot = { @@ -25,7 +25,7 @@ const CONVERSATION_ROOTS: readonly AntigravityConversationRoot[] = [ { dir: join(homedir(), '.gemini', 'antigravity-cli', 'conversations'), project: 'antigravity-cli', - extensions: ['.pb'], + extensions: ['.pb', '.db'], }, { dir: join(homedir(), '.gemini', 'antigravity-cli', 'implicit'), @@ -137,11 +137,29 @@ type AntigravityCache = { cascades: Record } +type ProtoField = { + number: number + wireType: number + value?: bigint + bytes?: Uint8Array +} + +type ProtoVarint = { + value: bigint + offset: number +} + +type AntigravityGenMetadataRow = { + idx: number + data: Uint8Array | string +} + const cachedServers = new Map() const cachedModelMaps = new Map() let memCache: AntigravityCache | null = null let cacheDirty = false let httpsAgent: https.Agent | undefined +const protoTextDecoder = new TextDecoder('utf-8', { fatal: false }) const SERVER_PORT_FLAGS = ['https_server_port', 'extension_server_port', 'https-server-port', 'extension-server-port'] const CSRF_TOKEN_FLAGS = ['csrf_token', 'extension_server_csrf_token', 'csrf-token', 'extension-server-csrf-token'] @@ -270,15 +288,15 @@ export function extractAntigravityModelMap(resp: unknown): ModelMap { if (!resp || typeof resp !== 'object') return {} const data = resp as ModelMapResponse const models = data.response?.models ?? data.models - const map: ModelMap = {} - if (!models) return map + const map = new Map() + if (!models) return {} for (const [key, info] of Object.entries(models)) { if (info && typeof info === 'object' && typeof info.model === 'string') { const canonicalKey = getCanonicalModelId(key, info.displayName) - map[info.model] = canonicalKey + map.set(info.model, canonicalKey) } } - return map + return Object.fromEntries(map) } export function extractAntigravityGeneratorMetadata(resp: unknown): GeneratorMetadata[] { @@ -544,6 +562,215 @@ function normalizePricingModel(model: string): string { return PRICING_ALIASES[stripped] ?? stripped } +function readProtoVarint(data: Uint8Array, startOffset: number): ProtoVarint | null { + let value = 0n + let shift = 0n + let offset = startOffset + + while (offset < data.length) { + const byte = BigInt(data[offset]!) + offset += 1 + value |= (byte & 0x7fn) << shift + if ((byte & 0x80n) === 0n) return { value, offset } + shift += 7n + if (shift > 70n) return null + } + + return null +} + +function parseProtoFields(data: Uint8Array): ProtoField[] { + const fields: ProtoField[] = [] + let offset = 0 + + while (offset < data.length) { + const key = readProtoVarint(data, offset) + if (!key) break + offset = key.offset + + const fieldNumber = Number(key.value >> 3n) + const wireType = Number(key.value & 0x7n) + if (!Number.isSafeInteger(fieldNumber) || fieldNumber <= 0) break + + if (wireType === 0) { + const value = readProtoVarint(data, offset) + if (!value) break + fields.push({ number: fieldNumber, wireType, value: value.value }) + offset = value.offset + continue + } + + if (wireType === 1) { + if (offset + 8 > data.length) break + fields.push({ number: fieldNumber, wireType, bytes: data.subarray(offset, offset + 8) }) + offset += 8 + continue + } + + if (wireType === 2) { + const length = readProtoVarint(data, offset) + if (!length) break + offset = length.offset + const byteLength = Number(length.value) + if (!Number.isSafeInteger(byteLength) || byteLength < 0 || offset + byteLength > data.length) break + fields.push({ number: fieldNumber, wireType, bytes: data.subarray(offset, offset + byteLength) }) + offset += byteLength + continue + } + + if (wireType === 5) { + if (offset + 4 > data.length) break + fields.push({ number: fieldNumber, wireType, bytes: data.subarray(offset, offset + 4) }) + offset += 4 + continue + } + + break + } + + return fields +} + +function firstProtoField(fields: readonly ProtoField[], fieldNumber: number): ProtoField | undefined { + return fields.find(field => field.number === fieldNumber) +} + +function protoFieldText(field: ProtoField | undefined): string | undefined { + if (!field?.bytes || field.bytes.length === 0) return undefined + const text = protoTextDecoder.decode(field.bytes) + if (!text || /[\u0000-\u0008\u000E-\u001F\u007F\uFFFD]/.test(text)) return undefined + return text +} + +function protoFieldPositiveInteger(field: ProtoField | undefined): number { + if (field?.value === undefined) return 0 + const value = Number(field.value) + return Number.isSafeInteger(value) && value > 0 ? value : 0 +} + +function protoFieldBytes(field: ProtoField | undefined): Uint8Array | undefined { + return field?.bytes +} + +function isAntigravityResponseId(value: string): boolean { + return /^[^\s]+$/.test(value) +} + +function antigravitySqliteResponseId(usageFields: readonly ProtoField[], fallback: string): string { + const responseId = protoFieldText(firstProtoField(usageFields, 11)) + return responseId && isAntigravityResponseId(responseId) ? responseId : fallback +} + +function genMetadataDataBytes(value: Uint8Array | string): Uint8Array { + return typeof value === 'string' + ? new TextEncoder().encode(value) + : value +} + +function antigravitySqliteMetadataAttributes(chatFields: readonly ProtoField[]): Map { + const attributes = new Map() + for (const field of chatFields) { + if (field.number !== 20) continue + const pairFields = parseProtoFields(protoFieldBytes(field) ?? new Uint8Array()) + const key = protoFieldText(firstProtoField(pairFields, 1)) + const value = protoFieldText(firstProtoField(pairFields, 2)) + if (key && value) attributes.set(key, value) + } + return attributes +} + +function antigravitySqliteModel(chatFields: readonly ProtoField[]): string { + const attributes = antigravitySqliteMetadataAttributes(chatFields) + const displayName = protoFieldText(firstProtoField(chatFields, 21)) + const rawModel = protoFieldText(firstProtoField(chatFields, 19)) + ?? attributes.get('model_enum') + ?? displayName + ?? 'unknown' + + return getCanonicalModelId(rawModel, displayName) +} + +function buildCallFromSqliteGenMetadataRow(cascadeId: string, row: AntigravityGenMetadataRow): ParsedProviderCall | null { + const rootFields = parseProtoFields(genMetadataDataBytes(row.data)) + const chatFields = parseProtoFields(protoFieldBytes(firstProtoField(rootFields, 1)) ?? new Uint8Array()) + const usageFields = parseProtoFields(protoFieldBytes(firstProtoField(chatFields, 4)) ?? new Uint8Array()) + if (usageFields.length === 0) return null + + const inputTokens = protoFieldPositiveInteger(firstProtoField(usageFields, 2)) + || protoFieldPositiveInteger(firstProtoField(usageFields, 1)) + const totalOutputTokens = protoFieldPositiveInteger(firstProtoField(usageFields, 3)) + let responseTokens = protoFieldPositiveInteger(firstProtoField(usageFields, 9)) + let thinkingTokens = protoFieldPositiveInteger(firstProtoField(usageFields, 10)) + + if (responseTokens === 0 && thinkingTokens === 0) { + responseTokens = totalOutputTokens + } else if (totalOutputTokens > 0 && responseTokens + thinkingTokens !== totalOutputTokens) { + const adjustedResponseTokens = totalOutputTokens - thinkingTokens + if (adjustedResponseTokens >= 0) responseTokens = adjustedResponseTokens + } + + if (inputTokens === 0 && totalOutputTokens === 0) return null + + const responseId = antigravitySqliteResponseId(usageFields, String(row.idx)) + const model = antigravitySqliteModel(chatFields) + const pricingModel = normalizePricingModel(model) + const costUSD = calculateCost(pricingModel, inputTokens, responseTokens + thinkingTokens, 0, 0, 0) + + return { + provider: 'antigravity', + model, + inputTokens, + outputTokens: responseTokens, + cacheCreationInputTokens: 0, + cacheReadInputTokens: 0, + cachedInputTokens: 0, + reasoningTokens: thinkingTokens, + webSearchRequests: 0, + costUSD, + tools: [], + bashCommands: [], + timestamp: '', + speed: 'standard', + deduplicationKey: `antigravity:${cascadeId}:${responseId}`, + userMessage: '', + sessionId: cascadeId, + } +} + +function buildCallsFromSqliteGenMetadata(cascadeId: string, rows: AntigravityGenMetadataRow[]): ParsedProviderCall[] { + const calls: ParsedProviderCall[] = [] + const seenResponseIds = new Set() + + for (const row of rows) { + const call = buildCallFromSqliteGenMetadataRow(cascadeId, row) + if (!call) continue + if (seenResponseIds.has(call.deduplicationKey)) continue + seenResponseIds.add(call.deduplicationKey) + calls.push(call) + } + + return calls +} + +async function parseSqliteGenMetadataCalls(filePath: string, cascadeId: string): Promise { + if (!filePath.toLowerCase().endsWith('.db')) return [] + if (!isSqliteAvailable()) return [] + + let db: ReturnType | null = null + try { + db = openDatabase(filePath) + const rows = db.query('SELECT idx, data FROM gen_metadata ORDER BY idx') + return buildCallsFromSqliteGenMetadata(cascadeId, rows) + } catch (err) { + // Let a transient lock propagate so the run retries this file on the next + // refresh instead of treating it as empty (see parser.ts busy handling). + if (isSqliteBusyError(err)) throw err + return [] + } finally { + db?.close() + } +} + function parseFiniteToken(value: unknown): number { return typeof value === 'number' && Number.isFinite(value) && value > 0 ? Math.floor(value) @@ -954,6 +1181,22 @@ function sanitizeProject(path: string): string { return basename(path.replace(/\\/g, '/')) } +function applyAntigravityProject(call: ParsedProviderCall, source: SessionSource, projectPath: string | undefined): void { + if (source.project === 'antigravity-cli') { + call.project = source.project + delete call.projectPath + return + } + + if (projectPath) { + call.projectPath = projectPath + call.project = sanitizeProject(projectPath) + return + } + + call.project = source.project +} + function createParser(source: SessionSource, seenKeys: Set): SessionParser { return { async *parse(): AsyncGenerator { @@ -974,12 +1217,30 @@ function createParser(source: SessionSource, seenKeys: Set): SessionPars const projectPath = await extractWorkspacePath(source.path) const cached = cache.cascades[cascadeId] - if (cached && cached.mtimeMs === s.mtimeMs && cached.sizeBytes === s.size) { + if (cached && cached.mtimeMs === s.mtimeMs && cached.sizeBytes === s.size && cached.calls.length > 0) { for (const call of cached.calls) { - if (projectPath) { - call.projectPath = projectPath - call.project = sanitizeProject(projectPath) - } + applyAntigravityProject(call, source, projectPath) + if (seenKeys.has(call.deduplicationKey)) continue + seenKeys.add(call.deduplicationKey) + yield call + } + return + } + + const sqliteResults = await parseSqliteGenMetadataCalls(source.path, cascadeId) + if (sqliteResults.length > 0) { + for (const call of sqliteResults) { + applyAntigravityProject(call, source, projectPath) + } + + cache.cascades[cascadeId] = { + mtimeMs: s.mtimeMs, + sizeBytes: s.size, + calls: sqliteResults, + } + cacheDirty = true + + for (const call of sqliteResults) { if (seenKeys.has(call.deduplicationKey)) continue seenKeys.add(call.deduplicationKey) yield call @@ -991,10 +1252,7 @@ function createParser(source: SessionSource, seenKeys: Set): SessionPars if (!server) { if (cached) { for (const call of cached.calls) { - if (projectPath) { - call.projectPath = projectPath - call.project = sanitizeProject(projectPath) - } + applyAntigravityProject(call, source, projectPath) if (seenKeys.has(call.deduplicationKey)) continue seenKeys.add(call.deduplicationKey) yield call @@ -1013,10 +1271,7 @@ function createParser(source: SessionSource, seenKeys: Set): SessionPars } catch { if (cached) { for (const call of cached.calls) { - if (projectPath) { - call.projectPath = projectPath - call.project = sanitizeProject(projectPath) - } + applyAntigravityProject(call, source, projectPath) if (seenKeys.has(call.deduplicationKey)) continue seenKeys.add(call.deduplicationKey) yield call @@ -1026,12 +1281,8 @@ function createParser(source: SessionSource, seenKeys: Set): SessionPars } const results = buildCallsFromGeneratorMetadata(cascadeId, metadata, modelMap) - if (projectPath) { - const projectName = sanitizeProject(projectPath) - for (const call of results) { - call.projectPath = projectPath - call.project = projectName - } + for (const call of results) { + applyAntigravityProject(call, source, projectPath) } cache.cascades[cascadeId] = { diff --git a/tests/fixtures/antigravity-cli-current/brain/fixture-current-cli/.system_generated/logs/transcript.jsonl b/tests/fixtures/antigravity-cli-current/brain/fixture-current-cli/.system_generated/logs/transcript.jsonl new file mode 100644 index 00000000..b6428239 --- /dev/null +++ b/tests/fixtures/antigravity-cli-current/brain/fixture-current-cli/.system_generated/logs/transcript.jsonl @@ -0,0 +1,2 @@ +{"step_index":0,"source":"USER_EXPLICIT","type":"USER_INPUT","status":"DONE","created_at":"2026-06-21T13:32:13Z","content":"sanitized fixture prompt"} +{"step_index":1,"source":"MODEL","type":"PLANNER_RESPONSE","status":"DONE","created_at":"2026-06-21T13:32:14Z","content":"sanitized fixture response","thinking":"sanitized fixture reasoning"} diff --git a/tests/fixtures/antigravity-cli-current/gen-metadata.json b/tests/fixtures/antigravity-cli-current/gen-metadata.json new file mode 100644 index 00000000..d9742d88 --- /dev/null +++ b/tests/fixtures/antigravity-cli-current/gen-metadata.json @@ -0,0 +1,10 @@ +{ + "conversationId": "fixture-current-cli", + "layout": "~/.gemini/antigravity-cli/conversations/.db with sibling brain//.system_generated/logs/transcript.jsonl", + "rows": [ + { + "idx": 0, + "hex": "120202032213666978747572652d63757272656e742d636c690ace01222508f80710b9ec0118da05301848930550475a12666978747572652d726573706f6e73652d319a011267656d696e692d70726f2d64656661756c74a201230a0d7472616a6563746f72795f69641212666978747572652d7472616a6563746f7279a201140a0b757365645f636c61756465120566616c7365a201140a0f6c6173745f737465705f696e646578120132a201230a0a6d6f64656c5f656e756d12154d4f44454c5f504c414345484f4c4445525f4d3136aa011547656d696e6920332e312050726f20284869676829" + } + ] +} diff --git a/tests/providers/antigravity.test.ts b/tests/providers/antigravity.test.ts index 76cba7be..e98fd591 100644 --- a/tests/providers/antigravity.test.ts +++ b/tests/providers/antigravity.test.ts @@ -1,8 +1,10 @@ import { mkdtemp, mkdir, readFile, rm, writeFile } from 'fs/promises' import { tmpdir } from 'os' import { join } from 'path' +import { createRequire } from 'node:module' import { describe, expect, it } from 'vitest' +import { isSqliteAvailable } from '../../src/sqlite.js' import { antigravityAppDataDirFromSourcePath, antigravityCascadeIdFromPath, @@ -17,6 +19,46 @@ import { recordAntigravityStatusLinePayload, shouldReparseAntigravitySource, } from '../../src/providers/antigravity.js' +import type { ParsedProviderCall } from '../../src/providers/types.js' + +const requireForTest = createRequire(import.meta.url) + +type CurrentCliFixture = { + conversationId: string + rows: Array<{ idx: number; hex: string }> +} + +type TestDb = { + exec(sql: string): void + prepare(sql: string): { run(...params: unknown[]): void } + close(): void +} + +function createCurrentAntigravityCliDb(dbPath: string, fixture: CurrentCliFixture): void { + const { DatabaseSync: Database } = requireForTest('node:sqlite') + const db = new Database(dbPath) as TestDb + try { + db.exec('CREATE TABLE gen_metadata (idx integer, data blob, size integer NOT NULL DEFAULT 0, PRIMARY KEY (idx))') + db.exec('CREATE TABLE trajectory_metadata_blob (id text DEFAULT "main", data blob, PRIMARY KEY (id))') + db.prepare('INSERT INTO trajectory_metadata_blob (id, data) VALUES (?, ?)').run( + 'main', + Buffer.from('file:///Users/example/private-project'), + ) + for (const row of fixture.rows) { + const data = Buffer.from(row.hex, 'hex') + db.prepare('INSERT INTO gen_metadata (idx, data, size) VALUES (?, ?, ?)').run(row.idx, data, data.length) + } + } finally { + db.close() + } +} + +async function collectAntigravityCalls(source: { path: string; project: string; provider: string }): Promise { + const parser = createAntigravityProvider().createSessionParser(source, new Set()) + const calls: ParsedProviderCall[] = [] + for await (const call of parser.parse()) calls.push(call) + return calls +} describe('antigravity provider helpers', () => { it('parses legacy https server flags from POSIX process args', () => { @@ -459,4 +501,108 @@ describe('antigravity provider helpers', () => { expect(shouldReparseAntigravitySource('/tmp/antigravity/conversation.pb', 0)).toBe(true) expect(shouldReparseAntigravitySource('/tmp/antigravity/conversation.pb', 1)).toBe(false) }) + + it('parses current Antigravity CLI SQLite conversations with non-zero token usage', async () => { + if (!isSqliteAvailable()) return + + const tempHome = await mkdtemp(join(tmpdir(), 'codeburn-antigravity-current-cli-')) + const cacheDir = join(tempHome, 'cache') + const previousCacheDir = process.env['CODEBURN_CACHE_DIR'] + process.env['CODEBURN_CACHE_DIR'] = cacheDir + + try { + const fixture = JSON.parse(await readFile( + new URL('../fixtures/antigravity-cli-current/gen-metadata.json', import.meta.url), + 'utf-8', + )) as CurrentCliFixture + const conversationsDir = join(tempHome, '.gemini', 'antigravity-cli', 'conversations') + const logsDir = join( + tempHome, + '.gemini', + 'antigravity-cli', + 'brain', + fixture.conversationId, + '.system_generated', + 'logs', + ) + + await mkdir(conversationsDir, { recursive: true }) + await mkdir(logsDir, { recursive: true }) + await writeFile( + join(logsDir, 'transcript.jsonl'), + await readFile( + new URL( + '../fixtures/antigravity-cli-current/brain/fixture-current-cli/.system_generated/logs/transcript.jsonl', + import.meta.url, + ), + 'utf-8', + ), + ) + + const dbPath = join(conversationsDir, `${fixture.conversationId}.db`) + createCurrentAntigravityCliDb(dbPath, fixture) + + const sources = await discoverAntigravitySessionSources([{ + dir: conversationsDir, + project: 'antigravity-cli', + extensions: ['.pb', '.db'], + }]) + expect(sources).toEqual([{ path: dbPath, project: 'antigravity-cli', provider: 'antigravity' }]) + + const calls = await collectAntigravityCalls(sources[0]!) + + expect(calls.length).toBeGreaterThanOrEqual(1) + expect(calls[0]).toMatchObject({ + provider: 'antigravity', + model: 'gemini-3.1-pro-high', + inputTokens: 30265, + outputTokens: 659, + reasoningTokens: 71, + sessionId: fixture.conversationId, + project: 'antigravity-cli', + }) + expect(calls[0]!.projectPath).toBeUndefined() + expect(calls[0]!.costUSD).toBeGreaterThan(0) + } finally { + if (previousCacheDir === undefined) delete process.env['CODEBURN_CACHE_DIR'] + else process.env['CODEBURN_CACHE_DIR'] = previousCacheDir + await rm(tempHome, { recursive: true, force: true }) + } + }) + + it('deduplicates current SQLite rows against RPC response ids with hyphens', async () => { + if (!isSqliteAvailable()) return + + const tempHome = await mkdtemp(join(tmpdir(), 'codeburn-antigravity-current-cli-dedup-')) + const cacheDir = join(tempHome, 'cache') + const previousCacheDir = process.env['CODEBURN_CACHE_DIR'] + process.env['CODEBURN_CACHE_DIR'] = cacheDir + + try { + const fixture = JSON.parse(await readFile( + new URL('../fixtures/antigravity-cli-current/gen-metadata.json', import.meta.url), + 'utf-8', + )) as CurrentCliFixture + const conversationsDir = join(tempHome, '.gemini', 'antigravity-cli', 'conversations') + + await mkdir(conversationsDir, { recursive: true }) + + const dbPath = join(conversationsDir, `${fixture.conversationId}.db`) + createCurrentAntigravityCliDb(dbPath, fixture) + + const parser = createAntigravityProvider().createSessionParser({ + path: dbPath, + project: 'antigravity-cli', + provider: 'antigravity', + }, new Set([`antigravity:${fixture.conversationId}:fixture-response-1`])) + const calls = [] + for await (const call of parser.parse()) calls.push(call) + + expect(calls).toEqual([]) + } finally { + if (previousCacheDir === undefined) delete process.env['CODEBURN_CACHE_DIR'] + else process.env['CODEBURN_CACHE_DIR'] = previousCacheDir + await rm(tempHome, { recursive: true, force: true }) + } + }) })