From 2deeb8af98e7206201f3488b36377e402eb10809 Mon Sep 17 00:00:00 2001 From: iamtoruk Date: Sat, 2 May 2026 08:56:07 -0700 Subject: [PATCH] Add Antigravity IDE provider Fetch token usage from Antigravity's local language server via RPC. Falls back to cached results when the IDE is closed. --- src/parser.ts | 5 + src/providers/antigravity.ts | 380 +++++++++++++++++++++++++++++++++++ src/providers/index.ts | 22 +- 3 files changed, 406 insertions(+), 1 deletion(-) create mode 100644 src/providers/antigravity.ts diff --git a/src/parser.ts b/src/parser.ts index 530a99b..fa6a345 100644 --- a/src/parser.ts +++ b/src/parser.ts @@ -4,6 +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 { flushAntigravityCache } from './providers/antigravity.js' import type { ParsedProviderCall } from './providers/types.js' import type { AssistantMessageContent, @@ -437,6 +438,10 @@ async function parseProviderSources( } } finally { if (providerName === 'codex') await flushCodexCache() + if (providerName === 'antigravity') { + const liveIds = new Set(sources.map(s => basename(s.path, '.pb'))) + await flushAntigravityCache(liveIds) + } } const projectMap = new Map() diff --git a/src/providers/antigravity.ts b/src/providers/antigravity.ts new file mode 100644 index 0000000..6c93bb7 --- /dev/null +++ b/src/providers/antigravity.ts @@ -0,0 +1,380 @@ +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' +import { homedir } from 'os' +import https from 'https' + +import { calculateCost } from '../models.js' +import type { Provider, SessionSource, SessionParser, ParsedProviderCall } from './types.js' + +const CONVERSATIONS_DIR = join(homedir(), '.gemini', 'antigravity', 'conversations') +const CACHE_VERSION = 1 + +const RPC_TIMEOUT_MS = 5000 +const MAX_RESPONSE_BYTES = 16 * 1024 * 1024 + +type ServerInfo = { + port: number + csrfToken: string +} + +type ModelMap = Record + +type UsageEntry = { + model: string + inputTokens: string + outputTokens: string + thinkingOutputTokens?: string + responseOutputTokens?: string + apiProvider: string + responseId?: string +} + +type GeneratorMetadata = { + stepIndices?: number[] + chatModel?: { + model: string + usage: UsageEntry + chatStartMetadata?: { + createdAt?: string + } + } +} + +type CachedCascade = { + mtimeMs: number + sizeBytes: number + calls: ParsedProviderCall[] +} + +type AntigravityCache = { + version: number + cascades: Record +} + +let cachedServer: ServerInfo | null | undefined +let cachedModelMap: ModelMap | undefined +let memCache: AntigravityCache | null = null +let cacheDirty = false +let httpsAgent: https.Agent | undefined + +function getAgent(): https.Agent { + if (!httpsAgent) httpsAgent = new https.Agent({ rejectUnauthorized: false }) + return httpsAgent +} + +function getCacheDir(): string { + return process.env['CODEBURN_CACHE_DIR'] ?? join(homedir(), '.cache', 'codeburn') +} + +function getCachePath(): string { + return join(getCacheDir(), 'antigravity-results.json') +} + +async function loadCache(): Promise { + if (memCache) return memCache + try { + const raw = await readFile(getCachePath(), 'utf-8') + const cache = JSON.parse(raw) as AntigravityCache + if (cache.version === CACHE_VERSION && cache.cascades && typeof cache.cascades === 'object') { + memCache = cache + return cache + } + } catch { /* no cache or invalid */ } + memCache = { version: CACHE_VERSION, cascades: {} } + return memCache +} + +async function flushCache(liveCascadeIds?: Set): Promise { + if (!memCache || !cacheDirty) return + try { + if (liveCascadeIds) { + for (const id of Object.keys(memCache.cascades)) { + if (!liveCascadeIds.has(id)) delete memCache.cascades[id] + } + } + + const dir = getCacheDir() + await mkdir(dir, { recursive: true }) + const finalPath = getCachePath() + const tempPath = `${finalPath}.${randomBytes(8).toString('hex')}.tmp` + const handle = await open(tempPath, 'w', 0o600) + try { + await handle.writeFile(JSON.stringify(memCache), { encoding: 'utf-8' }) + await handle.sync() + } finally { + await handle.close() + } + try { + await rename(tempPath, finalPath) + } catch { + try { await unlink(tempPath) } catch { /* cleanup */ } + } + cacheDirty = false + } catch { /* best-effort */ } +} + +async function detectServer(): Promise { + if (cachedServer !== undefined) return cachedServer + try { + const output = await new Promise((resolve, reject) => { + execFile('ps', ['-eo', 'args'], { encoding: 'utf-8', timeout: 3000 }, (err, stdout) => { + if (err) reject(err) + else resolve(stdout) + }) + }) + for (const line of output.split('\n')) { + if (!line.includes('language_server') || !line.includes('antigravity')) continue + if (!line.includes('--https_server_port')) continue + + const csrfMatch = line.match(/--csrf_token\s+([0-9a-f-]{32,})/) + const portMatch = line.match(/--https_server_port\s+(\d+)/) + if (csrfMatch && portMatch) { + cachedServer = { csrfToken: csrfMatch[1]!, port: parseInt(portMatch[1]!, 10) } + return cachedServer + } + } + } catch { /* ps failed or timed out */ } + cachedServer = null + return null +} + +async function rpc(server: ServerInfo, method: string, body: Record = {}): Promise { + return new Promise((resolve, reject) => { + const data = JSON.stringify(body) + const req = https.request({ + hostname: '127.0.0.1', + port: server.port, + path: `/exa.language_server_pb.LanguageServerService/${method}`, + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'Connect-Protocol-Version': '1', + 'X-Codeium-Csrf-Token': server.csrfToken, + 'Content-Length': Buffer.byteLength(data), + }, + agent: getAgent(), + timeout: RPC_TIMEOUT_MS, + }, (res) => { + const chunks: Buffer[] = [] + let totalBytes = 0 + res.on('data', (chunk: Buffer) => { + totalBytes += chunk.length + if (totalBytes > MAX_RESPONSE_BYTES) { + res.destroy() + reject(new Error(`RPC ${method}: response too large`)) + return + } + chunks.push(chunk) + }) + res.on('end', () => { + if (res.statusCode !== 200) { + reject(new Error(`RPC ${method}: HTTP ${res.statusCode}`)) + return + } + try { + resolve(JSON.parse(Buffer.concat(chunks).toString('utf-8'))) + } catch { + reject(new Error(`RPC ${method}: invalid JSON`)) + } + }) + res.on('error', reject) + }) + req.on('error', reject) + req.on('timeout', () => { req.destroy(); reject(new Error(`RPC ${method}: timeout`)) }) + req.write(data) + req.end() + }) +} + +async function getModelMap(server: ServerInfo): Promise { + if (cachedModelMap) return cachedModelMap + const map: ModelMap = {} + try { + const resp = await rpc(server, 'GetAvailableModels') as { + response?: { models?: Record } + } + const models = resp?.response?.models + if (models) { + for (const [key, info] of Object.entries(models)) { + if (info.model) map[info.model] = key + } + } + } catch { /* best-effort */ } + cachedModelMap = map + return map +} + +// Strip Antigravity-specific suffixes so the pricing DB can match +function normalizePricingModel(model: string): string { + return model.replace(/-(high|low|agent)$/, '') +} + +async function discoverSessions(): Promise { + const sources: SessionSource[] = [] + let files: string[] + try { + files = await readdir(CONVERSATIONS_DIR) + } catch { + return sources + } + + for (const file of files) { + if (!file.endsWith('.pb')) continue + sources.push({ + path: join(CONVERSATIONS_DIR, file), + project: 'antigravity', + provider: 'antigravity', + }) + } + return sources +} + +function createParser(source: SessionSource, seenKeys: Set): SessionParser { + return { + async *parse(): AsyncGenerator { + const cascadeId = basename(source.path, '.pb') + const cache = await loadCache() + + const s = await stat(source.path).catch(() => null) + if (!s) return + + const cached = cache.cascades[cascadeId] + if (cached && cached.mtimeMs === s.mtimeMs && cached.sizeBytes === s.size) { + for (const call of cached.calls) { + if (seenKeys.has(call.deduplicationKey)) continue + seenKeys.add(call.deduplicationKey) + yield call + } + return + } + + const server = await detectServer() + if (!server) { + if (cached) { + for (const call of cached.calls) { + if (seenKeys.has(call.deduplicationKey)) continue + seenKeys.add(call.deduplicationKey) + yield call + } + } + return + } + + const modelMap = await getModelMap(server) + + let metadata: GeneratorMetadata[] + try { + const resp = await rpc(server, 'GetCascadeTrajectoryGeneratorMetadata', { cascadeId }) as { + generatorMetadata?: GeneratorMetadata[] + } + metadata = resp?.generatorMetadata ?? [] + } catch { + if (cached) { + for (const call of cached.calls) { + if (seenKeys.has(call.deduplicationKey)) continue + seenKeys.add(call.deduplicationKey) + yield call + } + } + return + } + + const results: ParsedProviderCall[] = [] + + for (const entry of metadata) { + 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 ?? '' + 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, + }) + } + + cache.cascades[cascadeId] = { + mtimeMs: s.mtimeMs, + sizeBytes: s.size, + calls: results, + } + cacheDirty = true + + for (const call of results) { + if (seenKeys.has(call.deduplicationKey)) continue + seenKeys.add(call.deduplicationKey) + yield call + } + }, + } +} + +const modelDisplayNames: Record = { + 'gemini-3.1-pro-high': 'Gemini 3.1 Pro', + 'gemini-3.1-pro-low': 'Gemini 3.1 Pro (Low)', + 'gemini-3-flash': 'Gemini 3 Flash', + 'gemini-3-flash-agent': 'Gemini 3 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', + 'claude-sonnet-4-6': 'Sonnet 4.6', +} + +export function createAntigravityProvider(): Provider { + return { + name: 'antigravity', + displayName: 'Antigravity', + + modelDisplayName(model: string): string { + return modelDisplayNames[model] ?? model + }, + + toolDisplayName(rawTool: string): string { + return rawTool + }, + + async discoverSessions(): Promise { + return discoverSessions() + }, + + createSessionParser(source: SessionSource, seenKeys: Set): SessionParser { + return createParser(source, seenKeys) + }, + } +} + +export async function flushAntigravityCache(liveCascadeIds?: Set): Promise { + await flushCache(liveCascadeIds) +} + +export const antigravity = createAntigravityProvider() diff --git a/src/providers/index.ts b/src/providers/index.ts index a754ada..182c219 100644 --- a/src/providers/index.ts +++ b/src/providers/index.ts @@ -11,6 +11,21 @@ import { qwen } from './qwen.js' import { rooCode } from './roo-code.js' import type { Provider, SessionSource } from './types.js' +let antigravityProvider: Provider | null = null +let antigravityLoadAttempted = false + +async function loadAntigravity(): Promise { + if (antigravityLoadAttempted) return antigravityProvider + antigravityLoadAttempted = true + try { + const { antigravity } = await import('./antigravity.js') + antigravityProvider = antigravity + return antigravity + } catch { + return null + } +} + let cursorProvider: Provider | null = null let cursorLoadAttempted = false @@ -59,8 +74,9 @@ async function loadCursorAgent(): Promise { const coreProviders: Provider[] = [claude, codex, copilot, droid, gemini, kiloCode, kiro, openclaw, pi, omp, qwen, rooCode] export async function getAllProviders(): Promise { - const [cursor, opencode, cursorAgent] = await Promise.all([loadCursor(), loadOpenCode(), loadCursorAgent()]) + const [ag, cursor, opencode, cursorAgent] = await Promise.all([loadAntigravity(), loadCursor(), loadOpenCode(), loadCursorAgent()]) const all = [...coreProviders] + if (ag) all.push(ag) if (cursor) all.push(cursor) if (opencode) all.push(opencode) if (cursorAgent) all.push(cursorAgent) @@ -83,6 +99,10 @@ export async function discoverAllSessions(providerFilter?: string): Promise { + if (name === 'antigravity') { + const ag = await loadAntigravity() + return ag ?? undefined + } if (name === 'cursor') { const cursor = await loadCursor() return cursor ?? undefined