From b51e466629ad246103a735810042301ad210cb02 Mon Sep 17 00:00:00 2001 From: BeiLuoqu Date: Fri, 23 Oct 2026 23:35:52 +0800 Subject: [PATCH] fix(antigravity): keep Conversation totals out of the context reading A Turn's result usage sums every request of the native Conversation: on a resumed Conversation its input_tokens was 31,422 while the latest request's prompt was 8,499. hostUsage fell back to input_tokens for context, so each Turn ended by reporting the Conversation's cumulative input as its context, growing with every Turn until the next step replaced it. Estimate context only from a step's usage, which describes one request, and count that request's whole prompt. agy reports the cached part separately as cache_read_tokens (a follow-up request's cache_read_tokens is about the previous request's prompt), so input_tokens alone undercounted context whenever the prompt cache was hit. Result usage still publishes its totals, but carries context only when agy names it explicitly. Co-Authored-By: Claude Opus 5.5 --- .../antigravity/src/antigravity-adapter.ts | 25 ++++-- .../test/antigravity-adapter.test.ts | 77 +++++++++++++++++++ 2 files changed, 96 insertions(+), 6 deletions(-) diff --git a/packages/adapters/antigravity/src/antigravity-adapter.ts b/packages/adapters/antigravity/src/antigravity-adapter.ts index c9fdc7495..b0a30249a 100644 --- a/packages/adapters/antigravity/src/antigravity-adapter.ts +++ b/packages/adapters/antigravity/src/antigravity-adapter.ts @@ -393,15 +393,28 @@ async function pollAntigravityContextUsage( return null; } -function hostUsage(value: AntigravityUsage | undefined, modelId?: string): HostUsage | null { +/** + * Host usage from one agy usage object. A step's usage describes one model request, whose prompt + * estimates the context in use; agy reports that prompt as `input_tokens` plus the separate + * `cache_read_tokens`. A result's usage sums every request of the native Conversation, so it + * carries context only when agy names it. + */ +function hostUsage( + value: AntigravityUsage | undefined, + modelId: string | undefined, + scope: "request" | "conversation", +): HostUsage | null { if (!value) return null; const inputTokens = safeToken(value.input_tokens); const outputTokens = safeToken(value.output_tokens); const reasoningOutputTokens = safeToken(value.thinking_tokens); const totalTokens = safeToken(value.total_tokens); - const contextUsedTokens = safeToken( - value.context_used_tokens ?? value.estimated_tokens_used ?? value.input_tokens, - ); + const requestPromptTokens = + scope === "request" && inputTokens !== undefined + ? inputTokens + (safeToken(value.cache_read_tokens) ?? 0) + : undefined; + const contextUsedTokens = + safeToken(value.context_used_tokens ?? value.estimated_tokens_used) ?? requestPromptTokens; const rawWindow = safeToken(value.context_window_tokens ?? value.max_context_tokens); const contextWindowTokens = contextUsedTokens !== undefined || rawWindow !== undefined @@ -902,7 +915,7 @@ class AntigravitySession implements HarnessSession { if (event.event === "step_update") { if (event.step_update.conversation_id !== this.#nativeRef?.nativeSessionId) return; await this.#handleStep(active, event.step_update); - const usage = hostUsage(event.step_update.usage, active.model?.id); + const usage = hostUsage(event.step_update.usage, active.model?.id, "request"); if (usage) this.#publishUsage(active, usage); this.#ensureContextUsage(active, event.step_update.conversation_id); return; @@ -930,7 +943,7 @@ class AntigravitySession implements HarnessSession { }); return; } - const usage = hostUsage(event.result.usage, active.model?.id); + const usage = hostUsage(event.result.usage, active.model?.id, "conversation"); if (usage) this.#publishUsage(active, usage); this.#ensureContextUsage(active, event.result.conversation_id); if (active.contextUsagePromise) { diff --git a/packages/adapters/antigravity/test/antigravity-adapter.test.ts b/packages/adapters/antigravity/test/antigravity-adapter.test.ts index 3ad558903..59dd2ce18 100644 --- a/packages/adapters/antigravity/test/antigravity-adapter.test.ts +++ b/packages/adapters/antigravity/test/antigravity-adapter.test.ts @@ -1751,6 +1751,83 @@ if (count === 0) { } }); + it("estimates context from the latest request's prompt, not the Conversation's result totals", async () => { + const streamLines = [ + JSON.stringify({ + event: "init", + init: { permission_mode: "dangerously-skip-permissions" }, + conversation_id: "conv-context", + }), + JSON.stringify({ + event: "step_update", + step_update: { + conversation_id: "conv-context", + step_index: 7, + state: "DONE", + step_type: "agent_response", + text_delta: "Done.", + // agy reports the cached part of this request's prompt separately. + usage: { + input_tokens: 2_751, + cache_read_tokens: 80_432, + output_tokens: 146, + thinking_tokens: 0, + total_tokens: 83_329, + }, + }, + }), + JSON.stringify({ + event: "result", + result: { + conversation_id: "conv-context", + status: "SUCCESS", + num_turns: 3, + response: "Done.", + // Totals over every request of the resumed Conversation, not one prompt. + usage: { + input_tokens: 120_418, + cache_read_tokens: 412_903, + output_tokens: 1_204, + thinking_tokens: 0, + total_tokens: 534_525, + }, + }, + }), + ]; + const { command, cwd, cleanup } = await fakeStreamingAgy(streamLines); + const adapter = new AntigravityAdapter({ command }); + try { + const opened = await adapter.open({ kind: "create", cwd }); + if (!opened.ok) throw new Error(opened.error.message); + const session = opened.value; + const iterator = session.outputs[Symbol.asyncIterator](); + await session.execute({ + type: "turn.start", + turnId: hostTurnIdSchema.parse("turn-context"), + input: [{ type: "text", text: "continue" }], + }); + const usages: Array> = []; + while (true) { + const event = await nextEvent(iterator); + if (event.type === "session.usage.changed" && event.usage) { + usages.push(event.usage as Record); + } + if (event.type === "turn.completed") break; + } + expect(usages.map(({ contextUsedTokens }) => contextUsedTokens)).not.toContain(120_418); + expect(usages.at(-1)).toMatchObject({ + contextUsedTokens: 83_183, + contextWindowTokens: 1_048_576, + inputTokens: 120_418, + outputTokens: 1_204, + }); + await session.close(); + } finally { + await adapter.close(); + await cleanup(); + } + }); + it("formats turn prompt with workspace file instructions and leaves slash commands untouched", () => { const normalPrompt = "Create a hello world python file"; const formatted = formatAntigravityTurnPrompt(normalPrompt);