From 1e83611289aec19cf55920593f6a267f6c3cc12e Mon Sep 17 00:00:00 2001 From: Dan Caldr <22105838+dcaldr@users.noreply.github.com> Date: Sun, 4 Oct 2026 20:05:07 +0200 Subject: [PATCH] feat(antigravity): align tool parsing, MCP attribution, and activity classification (#1620) * feat(antigravity): align tool parsing, MCP attribution, and activity classification Extends Antigravity tool classification and MCP normalization: - Classify manage_task as planning and search_web/read_url_content as exploration in TASK_TOOLS/SEARCH_TOOLS. - Recognize invoke_subagent in hasAgentSpawn alongside Agent in parser.ts to categorize delegation turns. - Implement and export normalizeAntigravityToolCall to normalize both eager mcp__ and lazy call_mcp_tool into canonical mcp____. - Connect toolDisplayName to normalizeAntigravityToolCall. - Bump ANTIGRAVITY_CACHE_VERSION to 7 and PROVIDER_PARSE_VERSIONS.antigravity to worktree-project-grouping-v7 so existing cached sessions re-parse and gain eager MCP attribution. - Add comprehensive unit tests in classifier and provider suites. * refactor(antigravity): map tool names in the provider, not the shared classifier; daily cache bump --------- Co-authored-by: Resham Joshi <65915470+iamtoruk@users.noreply.github.com> --- src/daily-cache.ts | 5 +- src/providers/antigravity.ts | 66 ++++++-- src/session-cache.ts | 2 +- tests/providers/antigravity.test.ts | 231 +++++++++++++++++++++++++++- 4 files changed, 290 insertions(+), 14 deletions(-) diff --git a/src/daily-cache.ts b/src/daily-cache.ts index fae626e6..7c58d2f7 100644 --- a/src/daily-cache.ts +++ b/src/daily-cache.ts @@ -243,7 +243,10 @@ import type { DateRange, ProjectSummary } from './types.js' // v47: Copilot chat-session journals read request-level promptTokens, repairing // missing input usage and input-only calls. Re-derive settled days from the // corrected session cache; calls and cost only rise. -export const DAILY_CACHE_VERSION = 47 +// v48: #1620 Antigravity manage_task/search_web/read_url_content/invoke_subagent +// reclassify turns. Calls and cost are unchanged, but settled antigravity +// category totals need re-derivation. +export const DAILY_CACHE_VERSION = 48 const MIN_SUPPORTED_VERSION = 28 /// Providers whose per-day CALL COUNT means something different at diff --git a/src/providers/antigravity.ts b/src/providers/antigravity.ts index c941876c..2122cdee 100644 --- a/src/providers/antigravity.ts +++ b/src/providers/antigravity.ts @@ -50,7 +50,7 @@ function conversationRoots(): readonly AntigravityConversationRoot[] { }, ] } -const CACHE_VERSION = 6 +const CACHE_VERSION = 7 export const ANTIGRAVITY_CACHE_VERSION = CACHE_VERSION export const ANTIGRAVITY_LEGACY_CACHE_FILE = 'antigravity-results.json' export function antigravityCacheFileName(version = CACHE_VERSION): string { @@ -829,6 +829,56 @@ function antigravitySqliteCreatedAt(chatFields: readonly ProtoField[]): string { const SKILL_MD_PATTERN = /(?:^|[\\/])([^\\/]+)[\\/]SKILL\.md$/i +const toolNameMap: Record = { + invoke_subagent: 'Agent', + manage_task: 'TodoWrite', + search_web: 'WebSearch', + read_url_content: 'WebFetch', +} + +/** + * Normalizes Antigravity tool calls to the canonical Codeburn format (`mcp____`). + * + * Antigravity emits MCP tool calls in two distinct shapes: + * 1. Eager: `mcp__` (single underscore delimiter between prefix, server, and tool). + * 2. Lazy: `call_mcp_tool` with JSON arguments `{ ServerName: "...", ToolName: "..." }`. + * + * Both are normalized to `mcp____` so they match `extractMcpTools`, + * attribute properly to MCP inventory, and deduplicate cleanly in dashboards. + */ +export function normalizeAntigravityToolCall(toolName: string, args?: Record | null): string { + if (!toolName) return '' + + if (toolName === 'call_mcp_tool') { + const rawServer = args?.['ServerName'] ?? args?.['server_name'] ?? args?.['serverName'] ?? args?.['server'] ?? args?.['Server'] + const server = typeof rawServer === 'string' ? rawServer.trim() : '' + const rawTool = args?.['ToolName'] ?? args?.['tool_name'] ?? args?.['toolName'] ?? args?.['tool'] ?? args?.['Tool'] + const tool = typeof rawTool === 'string' ? rawTool.trim() : '' + if (server && tool) { + return `mcp__${server}__${tool}` + } + return 'call_mcp_tool' + } + + if (toolName.startsWith('mcp__')) { + return toolName + } + + if (toolName.startsWith('mcp_')) { + const rest = toolName.slice(4) + // ponytail: first-underscore split misattributes servers whose names contain '_'; upgrade: longest-prefix match against mcp_config.json server keys + const sep = rest.indexOf('_') + if (sep > 0 && sep < rest.length - 1) { + const server = rest.slice(0, sep) + const tool = rest.slice(sep + 1) + return `mcp__${server}__${tool}` + } + return toolName + } + + return toolName +} + function extractAntigravityToolFromStep(metadataBytes: Uint8Array, turn: TurnTools): void { const fields = parseProtoFields(metadataBytes) for (const field of fields) { @@ -848,13 +898,7 @@ function extractAntigravityToolFromStep(metadataBytes: Uint8Array, turn: TurnToo } if (toolName === 'call_mcp_tool') { - const server = typeof args?.['ServerName'] === 'string' ? args['ServerName'].trim() : '' - const tool = typeof args?.['ToolName'] === 'string' ? args['ToolName'].trim() : '' - if (server && tool) { - turn.tools.push(`mcp__${server}__${tool}`) - } else { - turn.tools.push('call_mcp_tool') - } + turn.tools.push(normalizeAntigravityToolCall(toolName, args)) continue } @@ -881,7 +925,7 @@ function extractAntigravityToolFromStep(metadataBytes: Uint8Array, turn: TurnToo } if (toolName === 'invoke_subagent') { - turn.tools.push(toolName) + turn.tools.push(toolNameMap[toolName]) if (Array.isArray(args?.['Subagents'])) { for (const sa of args['Subagents']) { if (sa && typeof sa === 'object') { @@ -898,7 +942,7 @@ function extractAntigravityToolFromStep(metadataBytes: Uint8Array, turn: TurnToo continue } - turn.tools.push(toolName) + turn.tools.push(toolNameMap[toolName] ?? normalizeAntigravityToolCall(toolName, args)) } } @@ -1642,7 +1686,7 @@ export function createAntigravityProvider(): Provider { }, toolDisplayName(rawTool: string): string { - return rawTool + return normalizeAntigravityToolCall(rawTool) }, async probeRoots(): Promise { diff --git a/src/session-cache.ts b/src/session-cache.ts index 3d4d47f7..18843639 100644 --- a/src/session-cache.ts +++ b/src/session-cache.ts @@ -565,7 +565,7 @@ export const PROVIDER_PARSE_VERSIONS: Record = { // token floor on every read, so they must re-parse once for the real dollars // to land. warp: 'worktree-project-grouping-v1-est-cost-billing-cost-v1', - antigravity: 'worktree-project-grouping-v6', + antigravity: 'worktree-project-grouping-v7', // pr-attribution-v1: the parser now reads the `message`/`part` tables for // per-turn user prompt text and the GitHub PR URLs it references. Cached // ZCode sessions hold empty userMessage turns and no session prLinks, so diff --git a/tests/providers/antigravity.test.ts b/tests/providers/antigravity.test.ts index 31e1c835..88e72ee6 100644 --- a/tests/providers/antigravity.test.ts +++ b/tests/providers/antigravity.test.ts @@ -18,11 +18,43 @@ import { parseAntigravityServerInfoFromLine, recordAntigravityStatusLinePayload, shouldReparseAntigravitySource, + normalizeAntigravityToolCall, } from '../../src/providers/antigravity.js' import type { ParsedProviderCall } from '../../src/providers/types.js' +import { classifyTurn } from '../../src/classifier.js' +import type { ParsedApiCall, ParsedTurn } from '../../src/types.js' const requireForTest = createRequire(import.meta.url) +// Mirrors the parsed-call -> turn shape src/parser.ts builds (cachedCallToApiCall). +function turnFromCall(call: ParsedProviderCall): ParsedTurn { + const apiCall: ParsedApiCall = { + provider: call.provider, + model: call.model, + usage: { + inputTokens: call.inputTokens, + outputTokens: call.outputTokens, + cacheCreationInputTokens: call.cacheCreationInputTokens, + cacheReadInputTokens: call.cacheReadInputTokens, + cachedInputTokens: call.cachedInputTokens, + reasoningTokens: call.reasoningTokens, + webSearchRequests: call.webSearchRequests, + }, + costUSD: call.costUSD, + tools: call.tools, + mcpTools: call.tools.filter(t => t.startsWith('mcp__')), + skills: call.skills ?? [], + subagentTypes: call.subagentTypes ?? [], + hasAgentSpawn: call.tools.includes('Agent'), + hasPlanMode: call.tools.includes('EnterPlanMode'), + speed: call.speed, + timestamp: call.timestamp, + bashCommands: call.bashCommands, + deduplicationKey: call.deduplicationKey, + } + return { userMessage: '', assistantCalls: [apiCall], timestamp: call.timestamp, sessionId: call.sessionId } +} + type CurrentCliFixture = { conversationId: string rows: Array<{ idx: number; hex: string }> @@ -802,7 +834,7 @@ describe('antigravity provider helpers', () => { 'run_command', 'mcp__dart-mcp-server__analyze_files', 'view_file', - 'invoke_subagent', + 'Agent', 'find_by_name', 'write_to_file', ]) @@ -812,6 +844,135 @@ describe('antigravity provider helpers', () => { }) }) + it('normalizes eager mcp__ calls from steps table', async () => { + if (!isSqliteAvailable()) return + + await withTempAntigravityHome('codeburn-antigravity-eager-mcp-', async (tempHome) => { + 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', 'conversations') + await mkdir(conversationsDir, { recursive: true }) + const dbPath = join(conversationsDir, `${fixture.conversationId}.db`) + const varint = (n: number): number[] => { + const out: number[] = [] + let v = n + while (v > 0x7f) { out.push((v & 0x7f) | 0x80); v = Math.floor(v / 128) } + out.push(v) + return out + } + const tag = (field: number, wire: number): number[] => varint(field * 8 + wire) + const lenField = (field: number, bytes: number[]): number[] => [...tag(field, 2), ...varint(bytes.length), ...bytes] + const strField = (field: number, str: string): number[] => lenField(field, Array.from(Buffer.from(str, 'utf-8'))) + + const encodeToolStepMetadata = (toolName: string, argsJson?: string): Buffer => { + const toolCallSub = [ + ...strField(1, 'call_test_123'), + ...strField(2, toolName), + ...(argsJson !== undefined ? strField(3, argsJson) : []), + ] + return Buffer.from(lenField(4, toolCallSub)) + } + + const withStepIndices = (fixtureHex: string, indices: number[]): Buffer => { + const rest = Buffer.from(fixtureHex, 'hex').subarray(4) + const packed = Buffer.from(indices.flatMap(n => varint(n))) + const field2 = Buffer.from([...tag(2, 2), ...varint(packed.length), ...packed]) + return Buffer.concat([field2, rest]) + } + + createCurrentAntigravityCliDb(dbPath, fixture) + + const { DatabaseSync: Database } = requireForTest('node:sqlite') + const db = new Database(dbPath) as TestDb + try { + db.prepare('UPDATE gen_metadata SET data = ? WHERE idx = 0').run( + withStepIndices(fixture.rows[0]!.hex, [1, 2]), + ) + db.exec('CREATE TABLE steps (idx integer PRIMARY KEY, step_type integer, metadata blob)') + const stmt = db.prepare('INSERT INTO steps (idx, step_type, metadata) VALUES (?, ?, ?)') + stmt.run(0, 15, null) + stmt.run(1, 38, encodeToolStepMetadata('mcp_context7_resolve-library-id')) + stmt.run(2, 38, encodeToolStepMetadata('call_mcp_tool', JSON.stringify({ ServerName: 'context7', ToolName: 'resolve-library-id' }))) + } finally { + db.close() + } + + const calls = await collectAntigravityCalls({ path: dbPath, project: 'antigravity', provider: 'antigravity' }) + expect(calls.length).toBeGreaterThan(0) + expect(calls[0]!.tools).toEqual([ + 'mcp__context7__resolve-library-id', + 'mcp__context7__resolve-library-id', + ]) + }) + }) + + it('maps manage_task, search_web, read_url_content and invoke_subagent to canonical tools that classify', async () => { + if (!isSqliteAvailable()) return + + await withTempAntigravityHome('codeburn-antigravity-tool-map-', async (tempHome) => { + const fixture = JSON.parse(await readFile( + new URL('../fixtures/antigravity-cli-current/gen-metadata.json', import.meta.url), + 'utf-8', + )) as CurrentCliFixture + + const varint = (n: number): number[] => { + const out: number[] = [] + let v = n + while (v > 0x7f) { out.push((v & 0x7f) | 0x80); v = Math.floor(v / 128) } + out.push(v) + return out + } + const tag = (field: number, wire: number): number[] => varint(field * 8 + wire) + const lenField = (field: number, bytes: number[]): number[] => [...tag(field, 2), ...varint(bytes.length), ...bytes] + const strField = (field: number, str: string): number[] => lenField(field, Array.from(Buffer.from(str, 'utf-8'))) + const encodeToolStepMetadata = (toolName: string, argsJson: string): Buffer => + Buffer.from(lenField(4, [...strField(1, 'call_test_123'), ...strField(2, toolName), ...strField(3, argsJson)])) + const withStepIndices = (fixtureHex: string, indices: number[]): Buffer => { + const rest = Buffer.from(fixtureHex, 'hex').subarray(4) + const packed = Buffer.from(indices.flatMap(n => varint(n))) + return Buffer.concat([Buffer.from([...tag(2, 2), ...varint(packed.length), ...packed]), rest]) + } + + const callWithTools = async (name: string, toolNames: string[]): Promise => { + const dir = join(tempHome, name) + await mkdir(dir, { recursive: true }) + const dbPath = join(dir, `${fixture.conversationId}.db`) + createCurrentAntigravityCliDb(dbPath, fixture) + const { DatabaseSync: Database } = requireForTest('node:sqlite') + const db = new Database(dbPath) as TestDb + try { + db.prepare('UPDATE gen_metadata SET data = ? WHERE idx = 0').run( + withStepIndices(fixture.rows[0]!.hex, toolNames.map((_, i) => i + 1)), + ) + db.exec('CREATE TABLE steps (idx integer PRIMARY KEY, step_type integer, metadata blob)') + const stmt = db.prepare('INSERT INTO steps (idx, step_type, metadata) VALUES (?, ?, ?)') + stmt.run(0, 15, null) + toolNames.forEach((tool, i) => stmt.run(i + 1, 38, encodeToolStepMetadata(tool, '{}'))) + } finally { + db.close() + } + const calls = await collectAntigravityCalls({ path: dbPath, project: 'antigravity', provider: 'antigravity' }) + expect(calls.length).toBeGreaterThan(0) + return calls[0]! + } + + const planning = await callWithTools('planning', ['manage_task']) + expect(planning.tools).toEqual(['TodoWrite']) + expect(classifyTurn(turnFromCall(planning)).category).toBe('planning') + + const exploration = await callWithTools('exploration', ['search_web', 'read_url_content']) + expect(exploration.tools).toEqual(['WebSearch', 'WebFetch']) + expect(classifyTurn(turnFromCall(exploration)).category).toBe('exploration') + + const delegation = await callWithTools('delegation', ['invoke_subagent', 'search_web']) + expect(delegation.tools).toEqual(['Agent', 'WebSearch']) + expect(classifyTurn(turnFromCall(delegation)).category).toBe('delegation') + }) + }) + it('gracefully handles missing steps table in legacy DB', async () => { if (!isSqliteAvailable()) return @@ -1109,3 +1270,71 @@ describe('antigravity provider helpers', () => { }) }) +describe('normalizeAntigravityToolCall', () => { + it('formats call_mcp_tool with PascalCase ServerName and ToolName as mcp____', () => { + expect(normalizeAntigravityToolCall('call_mcp_tool', { ServerName: 'context7', ToolName: 'resolve-library-id' })) + .toBe('mcp__context7__resolve-library-id') + }) + + it('formats call_mcp_tool with snake_case keys', () => { + expect(normalizeAntigravityToolCall('call_mcp_tool', { server_name: 'context7', tool_name: 'resolve-library-id' })) + .toBe('mcp__context7__resolve-library-id') + }) + + it('falls back to call_mcp_tool when args are missing or empty', () => { + expect(normalizeAntigravityToolCall('call_mcp_tool', null)).toBe('call_mcp_tool') + expect(normalizeAntigravityToolCall('call_mcp_tool', {})).toBe('call_mcp_tool') + expect(normalizeAntigravityToolCall('call_mcp_tool', { ServerName: '' })).toBe('call_mcp_tool') + }) + + it('normalizes eager mcp__ into mcp____', () => { + expect(normalizeAntigravityToolCall('mcp_context7_resolve-library-id')) + .toBe('mcp__context7__resolve-library-id') + expect(normalizeAntigravityToolCall('mcp_dart-mcp-server_analyze_files')) + .toBe('mcp__dart-mcp-server__analyze_files') + }) + + it('preserves already canonical mcp____', () => { + expect(normalizeAntigravityToolCall('mcp__context7__resolve-library-id')) + .toBe('mcp__context7__resolve-library-id') + }) + + it('leaves non-mcp tools unchanged', () => { + expect(normalizeAntigravityToolCall('run_command')).toBe('run_command') + expect(normalizeAntigravityToolCall('view_file')).toBe('view_file') + expect(normalizeAntigravityToolCall('manage_task')).toBe('manage_task') + expect(normalizeAntigravityToolCall('search_web')).toBe('search_web') + }) + + it('splits eager mcp names at the first underscore, misattributing servers with underscores', () => { + expect(normalizeAntigravityToolCall('mcp_github_mcp_server_create_issue')) + .toBe('mcp__github__mcp_server_create_issue') + }) + + it('leaves incomplete mcp prefixes unchanged', () => { + expect(normalizeAntigravityToolCall('mcp_')).toBe('mcp_') + expect(normalizeAntigravityToolCall('mcp_noservertool')).toBe('mcp_noservertool') + }) + it('formats call_mcp_tool with camelCase keys', () => { + expect(normalizeAntigravityToolCall('call_mcp_tool', { serverName: 'context7', toolName: 'resolve-library-id' })) + .toBe('mcp__context7__resolve-library-id') + }) + + it('formats call_mcp_tool with short server/tool keys', () => { + expect(normalizeAntigravityToolCall('call_mcp_tool', { server: 'context7', tool: 'resolve-library-id' })) + .toBe('mcp__context7__resolve-library-id') + }) + + it('handles empty or falsy toolName safely', () => { + expect(normalizeAntigravityToolCall('')).toBe('') + }) +}) + +describe('antigravity provider toolDisplayName', () => { + it('normalizes eager and canonical MCP tools to canonical display names', () => { + const provider = createAntigravityProvider() + expect(provider.toolDisplayName('mcp_context7_resolve-library-id')).toBe('mcp__context7__resolve-library-id') + expect(provider.toolDisplayName('mcp__context7__resolve-library-id')).toBe('mcp__context7__resolve-library-id') + expect(provider.toolDisplayName('run_command')).toBe('run_command') + }) +})