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.
This commit is contained in:
Resham Joshi 2026-05-02 08:58:23 -07:00 committed by GitHub
parent 791f2b077d
commit 8c845253c2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 406 additions and 1 deletions

View file

@ -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<string, SessionSummary[]>()

View file

@ -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<string, string>
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<string, CachedCascade>
}
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<AntigravityCache> {
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<string>): Promise<void> {
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<ServerInfo | null> {
if (cachedServer !== undefined) return cachedServer
try {
const output = await new Promise<string>((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<string, unknown> = {}): Promise<unknown> {
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<ModelMap> {
if (cachedModelMap) return cachedModelMap
const map: ModelMap = {}
try {
const resp = await rpc(server, 'GetAvailableModels') as {
response?: { models?: Record<string, { model?: string }> }
}
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<SessionSource[]> {
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<string>): SessionParser {
return {
async *parse(): AsyncGenerator<ParsedProviderCall> {
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<string, string> = {
'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<SessionSource[]> {
return discoverSessions()
},
createSessionParser(source: SessionSource, seenKeys: Set<string>): SessionParser {
return createParser(source, seenKeys)
},
}
}
export async function flushAntigravityCache(liveCascadeIds?: Set<string>): Promise<void> {
await flushCache(liveCascadeIds)
}
export const antigravity = createAntigravityProvider()

View file

@ -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<Provider | null> {
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<Provider | null> {
const coreProviders: Provider[] = [claude, codex, copilot, droid, gemini, kiloCode, kiro, openclaw, pi, omp, qwen, rooCode]
export async function getAllProviders(): Promise<Provider[]> {
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<Sess
}
export async function getProvider(name: string): Promise<Provider | undefined> {
if (name === 'antigravity') {
const ag = await loadAntigravity()
return ag ?? undefined
}
if (name === 'cursor') {
const cursor = await loadCursor()
return cursor ?? undefined