mixdog 1.2.2 → 1.2.4

package/src/runtime/agent/orchestrator/providers/antigravity-request.mjs

@@ -15,6 +15,12 @@ export function isAntigravityClaude(model) {
   return /^claude-/i.test(String(model || ''));
 }
 
+// Claude declarations are validated for the gateway's Anthropic bridge in
+// toGeminiTools (claudeBridge: true); a tool with an invalid schema is left
+// out of the request and reported through `onToolExcluded` / console.warn.
+const claudeRequestTools = (tools, nativeTools, onToolExcluded) =>
+  toGeminiRequestTools(tools, nativeTools, { claudeBridge: true, onToolExcluded });
+
 // Default function-calling mode when the caller sets no toolChoice. VALIDATED
 // (schema-constrained function calls) is what the Antigravity client sends;
 // MIXDOG_ANTIGRAVITY_FC_MODE=AUTO|ANY|VALIDATED overrides it for A/B runs.
@@ -35,14 +41,72 @@ function replayFamilyMatches(message, model) {
   return isAntigravityClaude(recorded) === isAntigravityClaude(model);
 }
 
+// The stream records empty text parts around Claude tool calls; Cloud Code's
+// Claude bridge rejects them on replay ("text.text: Field required").
+const isEmptyPlainTextPart = (part) =>
+  part?.text === '' && part.thought !== true && typeof part.thoughtSignature !== 'string';
+
+// Streamed Claude thinking arrives as many thought deltas with the signature
+// on the closing (often empty) delta; the bridge turns every part into its own
+// thinking block and rejects the unsigned ones ("thinking.signature: Field
+// required"). Rejoin each run into one block that ends at its signature.
+function coalesceClaudeThoughts(parts) {
+  const out = [];
+  let open = null;
+  for (const part of parts) {
+    if (part?.thought !== true) {
+      if (open) out.push(open);
+      open = null;
+      out.push(part);
+      continue;
+    }
+    const partText = typeof part.text === 'string' ? part.text : '';
+    open = open ? { ...open, ...part, text: `${open.text}${partText}` } : { ...part, text: partText };
+    if (typeof part.thoughtSignature === 'string') {
+      out.push(open);
+      open = null;
+    }
+  }
+  if (open) out.push(open);
+  return out;
+}
+
+// The validator-bypass sentinel is a Gemini wire token: the Claude bridge
+// sends a thinking block carrying it (or no signature at all) without a
+// signature and Anthropic rejects the request ("thinking.signature: Field
+// required"). Only a signature this route recorded is replayed to Claude;
+// any other signature is removed, and a thinking block left without one (an
+// interrupted run, another model's or provider's turn, legacy metadata) is
+// dropped rather than turned into visible text.
+const hasOwnClaudeSignature = (part) =>
+  typeof part.thoughtSignature === 'string' &&
+  part.thoughtSignature !== '' &&
+  part.thoughtSignature !== THOUGHT_SIGNATURE_SENTINEL;
+
+function claudeReplayParts(parts, own) {
+  const signed = parts.map((part) => {
+    if (!part || typeof part !== 'object' || !('thoughtSignature' in part)) return part;
+    if (own && hasOwnClaudeSignature(part)) return part;
+    const { thoughtSignature: _foreign, ...unsigned } = part;
+    return unsigned;
+  });
+  return coalesceClaudeThoughts(signed.filter((part) => !isEmptyPlainTextPart(part))).filter(
+    (part) => part?.thought !== true || hasOwnClaudeSignature(part)
+  );
+}
+
 function signatureSafeMessages(messages, model) {
-  return messages.map((message) => {
-    if (message?.role !== 'assistant') return message;
+  const claude = isAntigravityClaude(model);
+  return messages.flatMap((message) => {
+    if (message?.role !== 'assistant') return [message];
     const ownParts = replayFamilyMatches(message, model) ? providerReplayItems(message, 'antigravity') : undefined;
     // Repair tool signatures only after the complete history establishes
     // the active turn, not while normalizing an isolated older message.
-    const parts = ownParts || toGeminiContents([message], model, { repairToolSignatures: false })[0]?.parts;
-    if (!parts?.length) return message;
+    const sourceParts = ownParts || toGeminiContents([message], model, { repairToolSignatures: false })[0]?.parts;
+    const parts = claude && sourceParts ? claudeReplayParts(sourceParts, !!ownParts) : sourceParts;
+    // A turn that held only empty or dropped parts is omitted, not sent as parts:[].
+    if (sourceParts?.length && !parts.length) return [];
+    if (!parts?.length) return [message];
     for (const part of parts) {
       if (!part || typeof part !== 'object') continue;
       if (part.thought === true || typeof part.thoughtSignature === 'string') {
@@ -52,10 +116,60 @@ function signatureSafeMessages(messages, model) {
         if (!ownParts || !part.thoughtSignature) part.thoughtSignature = THOUGHT_SIGNATURE_SENTINEL;
       }
     }
-    return { ...message, providerReplay: createProviderReplay('antigravity', parts) };
+    return [{ ...message, providerReplay: createProviderReplay('antigravity', parts) }];
   });
 }
 
+const isToolResultTurn = (content) =>
+  content?.role === 'user' && !!content.parts?.some((part) => part?.functionResponse);
+const isPromptTurn = (content) => content?.role === 'user' && !isToolResultTurn(content);
+const hasSignedThinking = (content) =>
+  !!content?.parts?.some((part) => part?.thought === true && hasOwnClaudeSignature(part));
+
+// With thinking on, Claude rejects a tool loop whose turn did not open with a
+// signed thinking block, and a block that was dropped (unsigned, foreign or
+// compacted away) cannot be restored. The outgoing request closes that loop
+// instead: the closed part loses its thinking, a short model reply
+// acknowledges the tool results, and a user nudge opens a fresh turn that
+// Claude starts with new thinking. The closure sits after the results of the
+// turn's last unthought model message, so later steps keep the signed
+// thinking that the fresh turn produced. Wire-only: the caller's messages and
+// recorded replays are never touched.
+function closeUnthoughtToolLoop(contents) {
+  if (!isToolResultTurn(contents.at(-1))) return contents;
+  const promptIndex = contents.findLastIndex(isPromptTurn);
+  const modelIndexes = [];
+  for (let index = promptIndex + 1; index < contents.length; index += 1) {
+    if (contents[index]?.role === 'model') modelIndexes.push(index);
+  }
+  if (!modelIndexes.length || hasSignedThinking(contents[modelIndexes[0]])) return contents;
+  const unthought = modelIndexes.findLast((index) => !hasSignedThinking(contents[index]));
+  let cut = unthought + 1;
+  let results = 0;
+  while (isToolResultTurn(contents[cut])) {
+    results += contents[cut].parts.filter((part) => part?.functionResponse).length;
+    cut += 1;
+  }
+  if (!results) return contents;
+  const closed = contents.slice(0, cut).flatMap((content) => {
+    if (content?.role !== 'model') return [content];
+    const parts = content.parts.filter((part) => part?.thought !== true);
+    return parts.length ? [{ ...content, parts }] : [];
+  });
+  const ack = results === 1 ? '[Tool result received.]' : `[${results} tool results received.]`;
+  return [
+    ...closed,
+    { role: 'model', parts: [{ text: ack }] },
+    { role: 'user', parts: [{ text: '[Continue.]' }] },
+    ...contents.slice(cut),
+  ];
+}
+
+// The request turns Claude thinking on through an effort level or token
+// budget; a `-thinking` wire id thinks without either.
+const claudeThinkingEnabled = (model, thinkingConfig) =>
+  !!(thinkingConfig.thinkingLevel || thinkingConfig.thinkingBudget) || /-thinking$/i.test(String(model));
+
 function sessionId(contents) {
   const first = contents.find((content) => content?.role === 'user');
   const text = first?.parts?.map((part) => (typeof part?.text === 'string' ? part.text : '')).join('') || '';
@@ -111,19 +225,27 @@ export function buildAntigravityRequest(messages, model, tools, opts = {}, proje
     .map((message) => message.content)
     .join('\n\n');
   const chatMessages = messages.filter((message) => message.role !== 'system');
-  const contents = toGeminiContents(signatureSafeMessages(chatMessages, model), model);
-  if (!contents.length) throw new Error('No messages to send');
+  const history = toGeminiContents(signatureSafeMessages(chatMessages, model), model);
+  if (!history.length) throw new Error('No messages to send');
   const claude = isAntigravityClaude(model);
   const thinkingConfig = claude ? claudeThinkingConfig(opts) : antigravityGeminiThinkingConfig(model, opts);
+  const contents = claude && claudeThinkingEnabled(model, thinkingConfig) ? closeUnthoughtToolLoop(history) : history;
   const generationConfig = {
     ...(claude ? { maxOutputTokens: CLAUDE_MAX_OUTPUT_TOKENS } : {}),
     ...(thinkingConfig ? { thinkingConfig } : {}),
   };
   const request = { contents };
   if (systemText) request.systemInstruction = { role: 'user', parts: [{ text: systemText }] };
-  const requestTools = toGeminiRequestTools(tools, opts.nativeTools);
+  const requestTools = claude
+    ? claudeRequestTools(tools, opts.nativeTools, opts.onToolExcluded)
+    : toGeminiRequestTools(tools, opts.nativeTools, {
+        onToolExcluded: opts.onToolExcluded,
+        // Other families (GPT-OSS) keep the plain conversion.
+        validate: isGeminiFamilyModel(model),
+      });
   if (requestTools.length) request.tools = requestTools;
-  if (tools?.length) {
+  // Tool config only accompanies function declarations that were kept.
+  if (requestTools.some((entry) => entry.functionDeclarations?.length)) {
     request.toolConfig = toGeminiToolConfig(opts.toolChoice) || {
       functionCallingConfig: { mode: antigravityFunctionCallingMode() },
     };

package/src/runtime/agent/orchestrator/providers/antigravity-response.mjs

@@ -8,13 +8,20 @@
  */
 import { traceAgentUsage } from '../agent-trace.mjs';
 import { createProviderReplay } from './lib/provider-replay.mjs';
-import { parseToolCalls, collectGeminiGroundingSources, parseGeminiTextPartMetadata } from './gemini-schema.mjs';
-import { geminiPolicyFinishCode } from './gemini-error-contract.mjs';
+import {
+  parseToolCalls,
+  emitGeminiToolCalls,
+  collectGeminiGroundingSources,
+  parseGeminiTextPartMetadata,
+} from './gemini-schema.mjs';
+import { geminiFinishErrorFields, geminiPolicyFinishCode } from './gemini-error-contract.mjs';
 
 // Pure usageMetadata normalization shared by the success path and a failed
 // stream's partialUsage.
 export function normalizeAntigravityUsage(um) {
-  const inputTokens = um.promptTokenCount || um.prompt_token_count || 0;
+  const inputTokens =
+    (um.promptTokenCount || um.prompt_token_count || 0) +
+    (um.toolUsePromptTokenCount || um.tool_use_prompt_token_count || 0);
   const cachedTokens = um.cachedContentTokenCount || um.cached_content_token_count || 0;
   const outputTokens =
     (um.candidatesTokenCount || um.candidates_token_count || 0) +
@@ -28,7 +35,7 @@ export function normalizeAntigravityUsage(um) {
  * @param {ReturnType<import('./antigravity-stream.mjs').createAntigravityStreamCollector>} deps.collector
  * @param {string} deps.useModel  wire model id
  * @param {object} deps.opts  send options (session/iteration/providerState)
- * @param {Function|null} deps.onToolCall  set when calls were streamed live
+ * @param {Function|null} deps.onToolCall  receives the turn's calls after a clean finish
  */
 export function finalizeAntigravityTurn({ response, collector, useModel, opts, onToolCall }) {
   const textLeakGuard = collector.leakGuard;
@@ -46,10 +53,8 @@ export function finalizeAntigravityTurn({ response, collector, useModel, opts, o
   // Thought signatures are only valid for the model family that minted
   // them; the request builder consults this when the route changes.
   if (providerReplay) providerReplay.requestContext = { model: useModel };
-  let nativeToolCalls;
-  if (!onToolCall) nativeToolCalls = parseToolCalls(responseParts);
-  else if (collector.streamedNativeToolCalls.length) nativeToolCalls = collector.streamedNativeToolCalls;
-  if (!onToolCall && textLeakGuard?.enabled) nativeToolCalls = textLeakGuard.filterNativeToolCalls(nativeToolCalls);
+  let nativeToolCalls = parseToolCalls(responseParts);
+  if (textLeakGuard?.enabled) nativeToolCalls = textLeakGuard.filterNativeToolCalls(nativeToolCalls);
   let toolCalls = nativeToolCalls;
   if (leakedToolCalls.length) {
     toolCalls = toolCalls?.length ? [...toolCalls, ...leakedToolCalls] : leakedToolCalls;
@@ -80,7 +85,17 @@ export function finalizeAntigravityTurn({ response, collector, useModel, opts, o
     });
   }
 
-  if (finishReason && normalizedFinish !== 'STOP') {
+  // A MAX_TOKENS cut that still delivered whole native functionCall parts is
+  // a tool-use turn: the calls arrived intact, only trailing output was cut.
+  // Every other non-STOP finish fails the turn, and so does a cut whose text
+  // also leaked a call (that call may be as truncated as the text around it).
+  const completeNativeCalls =
+    normalizedFinish === 'MAX_TOKENS' &&
+    String(candidate?.finishReason || '').replace(/^FINISH_REASON_/, '') === 'MAX_TOKENS' &&
+    !leakedToolCalls.length &&
+    !!nativeToolCalls?.length &&
+    nativeToolCalls.every((call) => typeof call?.name === 'string' && call.name !== '');
+  if (finishReason && normalizedFinish !== 'STOP' && !completeNativeCalls) {
     throw Object.assign(new Error(`Antigravity response incomplete: finishReason=${finishReason}`), {
       name: 'ProviderIncompleteError',
       code: geminiPolicyFinishCode(finishReason) || 'PROVIDER_INCOMPLETE',
@@ -93,10 +108,14 @@ export function finalizeAntigravityTurn({ response, collector, useModel, opts, o
       model: useModel,
       rawUsage: response.usageMetadata || null,
       ...(usage ? { partialUsage: usage, partialModel: useModel } : {}),
-      ...(collector.emittedToolCount ? { emittedToolCall: true, unsafeToRetry: true } : {}),
+      ...geminiFinishErrorFields(response),
     });
   }
 
+  // Terminal status first: a failed turn never runs its tool calls.
+  textLeakGuard?.releaseDeferredToolCalls();
+  emitGeminiToolCalls(nativeToolCalls, onToolCall);
+
   return {
     content,
     model: useModel,

package/src/runtime/agent/orchestrator/providers/antigravity-retry-policy.mjs

@@ -4,7 +4,8 @@
  *
  *   capacity (429/503/529, MODEL_CAPACITY_EXHAUSTED, "server busy"): retried on
  *     the SAME account and host with 1s → 2s → 4s (8s cap) ±10% jitter, three
- *     retries, whatever wait the server names, then the failure surfaces.
+ *     retries, waiting at least the server's RetryInfo delay; advice that
+ *     would push total waiting past 60s surfaces the error (with resetAt).
  *   Endpoint: the reference falls back to its next endpoint after a failure;
  *     generation here deliberately stays on the single daily host and never
  *     walks an endpoint list (a spent ladder surfaces to the account layer).
@@ -13,34 +14,135 @@
  *     rotate; advice beyond 60s surfaces at once.
  *   QUOTA_EXHAUSTED (and any day/plan window): never retried here.
  *
+ *   First byte: the gateway queues or thinks server-side and sends the
+ *     response headers together with the first chunk, so the window is chosen
+ *     per model class from the catalog `family`: Pro-class Gemini (Pro tiers,
+ *     Argon) waits MIXDOG_ANTIGRAVITY_FIRST_BYTE_TIMEOUT_MS (300s, clamped to
+ *     the provider ceiling PROVIDER_MAX_BEFORE_WARN_MS, 285s by default),
+ *     Flash-class Gemini the shared Gemini window (60s); both apply to the
+ *     headers and again to the first chunk. A header or first-chunk timeout is
+ *     retried at most MIXDOG_ANTIGRAVITY_FIRST_BYTE_RETRIES times (1), each
+ *     retry doubling the window up to that ceiling; connect, 5xx and drop
+ *     failures keep the full ladder (MIXDOG_ANTIGRAVITY_RETRY_MAX_ATTEMPTS, 4
+ *     attempts, at MIXDOG_ANTIGRAVITY_RETRY_BACKOFF_MS 1s / 2s / 4s). Models
+ *     without a Gemini family (Claude, GPT-OSS, unlisted ids) keep the shared
+ *     Gemini window on that ladder.
+ *
  * Replay stays gated by the stream-outcome contract: an attempt that exposed
  * text or dispatched a tool is never re-issued.
  */
 import { readStreamOutcome } from './lib/stream-outcome.mjs';
 import { DETERMINISTIC_KINDS, providerErrorKind } from './provider-error-kind.mjs';
+import { PROVIDER_MAX_BEFORE_WARN_MS, firstBytePolicyBudget, resolveTimeoutMs } from '../stall-policy.mjs';
+import { GEMINI_FIRST_BYTE_TIMEOUT_MS } from './gemini-stream.mjs';
+import { antigravityModelFamily } from './antigravity-oauth-catalog.mjs';
 
 export const ANTIGRAVITY_CAPACITY_RETRIES = 3;
 export const ANTIGRAVITY_CAPACITY_BASE_MS = 1000;
 export const ANTIGRAVITY_CAPACITY_MAX_MS = 8000;
 export const ANTIGRAVITY_CAPACITY_JITTER_RATIO = 0.1;
+// Total time capacity retries may spend waiting on server-advised delays.
+export const ANTIGRAVITY_CAPACITY_BUDGET_MS = 60_000;
 export const ANTIGRAVITY_QUICK_RETRY_MS = 1000;
 export const ANTIGRAVITY_QUICK_RETRY_MAX_ADVICE_MS = 60_000;
 
+const DEFAULT_RETRY_BACKOFF_MS = Object.freeze([1000, 1000, 2000, 4000]);
+
+function envInt(env, name, fallback, min, max) {
+  const raw = env?.[name];
+  if (raw == null || String(raw).trim() === '') return fallback;
+  const n = Number(raw);
+  if (!Number.isInteger(n)) return fallback;
+  return Math.min(Math.max(n, min), max);
+}
+
+function envBackoffMs(env, name) {
+  const raw = String(env?.[name] ?? '').trim();
+  if (!raw) return DEFAULT_RETRY_BACKOFF_MS;
+  const steps = raw.split(',').map((step) => Number(step.trim()));
+  if (!steps.every((step) => Number.isInteger(step) && step >= 0)) return DEFAULT_RETRY_BACKOFF_MS;
+  return Object.freeze(steps);
+}
+
+/**
+ * The same-host ladder for everything that is not a 429, read from the
+ * environment: attempts (1–6), the backoff list (index n is the wait before
+ * attempt n) and how many first-byte timeouts a thinking model may retry.
+ */
+export function resolveAntigravityRetryConfig(env = process.env) {
+  return {
+    maxAttempts: envInt(env, 'MIXDOG_ANTIGRAVITY_RETRY_MAX_ATTEMPTS', ANTIGRAVITY_CAPACITY_RETRIES + 1, 1, 6),
+    backoffMs: envBackoffMs(env, 'MIXDOG_ANTIGRAVITY_RETRY_BACKOFF_MS'),
+    firstByteRetries: envInt(env, 'MIXDOG_ANTIGRAVITY_FIRST_BYTE_RETRIES', 1, 0, 5),
+  };
+}
+
+const RETRY_CONFIG = resolveAntigravityRetryConfig();
+
 // withRetry options for everything that is not a 429: transient transport and
 // 5xx failures use the capacity schedule (the reference has no other same-host
 // retry; the leading duplicate is the sleep-before-attempt index).
 export const ANTIGRAVITY_RETRY = Object.freeze({
-  maxAttempts: ANTIGRAVITY_CAPACITY_RETRIES + 1,
-  backoffMs: Object.freeze([1000, 1000, 2000, 4000]),
+  maxAttempts: RETRY_CONFIG.maxAttempts,
+  backoffMs: RETRY_CONFIG.backoffMs,
   retryJitterRatio: ANTIGRAVITY_CAPACITY_JITTER_RATIO,
   // 429s are handled by antigravity429WaitMs, not the generic loop.
   retry429: false,
 });
 
+export const ANTIGRAVITY_FIRST_BYTE_RETRIES = RETRY_CONFIG.firstByteRetries;
+
+// Pro-class window: 300s requested, clamped to the provider ceiling.
+export const ANTIGRAVITY_FIRST_BYTE_TIMEOUT_MS = resolveTimeoutMs('MIXDOG_ANTIGRAVITY_FIRST_BYTE_TIMEOUT_MS', 300_000, {
+  minMs: 30_000,
+  maxMs: PROVIDER_MAX_BEFORE_WARN_MS,
+});
+
+// First-attempt window per catalog family.
+const FAMILY_FIRST_BYTE_MS = Object.freeze({
+  'gemini-pro': ANTIGRAVITY_FIRST_BYTE_TIMEOUT_MS,
+  'gemini-flash': GEMINI_FIRST_BYTE_TIMEOUT_MS,
+});
+
+// Models without a Gemini family: the shared Gemini window on the ordinary ladder.
+export const ANTIGRAVITY_UNCLASSED_FIRST_BYTE = Object.freeze({
+  timeoutMs: GEMINI_FIRST_BYTE_TIMEOUT_MS,
+  maxMs: GEMINI_FIRST_BYTE_TIMEOUT_MS,
+  retries: null,
+});
+
+/**
+ * First-byte policy (see firstByteWindowMs) for one model, picker or wire id.
+ * The class comes from the catalog record's `family`.
+ */
+export function antigravityFirstBytePolicy(model, models) {
+  const timeoutMs = FAMILY_FIRST_BYTE_MS[antigravityModelFamily(model, models)];
+  if (!timeoutMs) return ANTIGRAVITY_UNCLASSED_FIRST_BYTE;
+  return {
+    timeoutMs,
+    maxMs: PROVIDER_MAX_BEFORE_WARN_MS,
+    retries: Math.min(ANTIGRAVITY_FIRST_BYTE_RETRIES, ANTIGRAVITY_RETRY.maxAttempts - 1),
+  };
+}
+
+/**
+ * Longest a send may wait before its first output across the whole transport
+ * ladder ({ transportMs, semanticMs }, see firstBytePolicyBudget): the agent
+ * watchdog must not fire inside it. Drops and 5xx share the ladder with
+ * first-byte timeouts, so the budget is every backoff step of the ladder plus
+ * the windows of the first-byte timeouts the policy permits (all attempts for
+ * a policy without its own retry cap).
+ */
+export function antigravityFirstResponseBudget(policy, backoffMs = ANTIGRAVITY_RETRY.backoffMs) {
+  return firstBytePolicyBudget(policy, { attempts: ANTIGRAVITY_RETRY.maxAttempts, backoffMs });
+}
+
 /** Capacity wait for retry number `retryIndex` (0-based): 1s, 2s, 4s, 8s cap, ±10%. */
-export function antigravityCapacityDelayMs(retryIndex) {
+function antigravityCapacityDelayMs(retryIndex) {
   const base = Math.min(ANTIGRAVITY_CAPACITY_BASE_MS * 2 ** retryIndex, ANTIGRAVITY_CAPACITY_MAX_MS);
-  return Math.round(base * (1 - ANTIGRAVITY_CAPACITY_JITTER_RATIO + Math.random() * 2 * ANTIGRAVITY_CAPACITY_JITTER_RATIO));
+  return Math.round(
+    base * (1 - ANTIGRAVITY_CAPACITY_JITTER_RATIO + Math.random() * 2 * ANTIGRAVITY_CAPACITY_JITTER_RATIO)
+  );
 }
 
 export function isAntigravityCapacityError(err) {
@@ -57,13 +159,28 @@ export function antigravity429WaitMs(err, state) {
   if (status !== 429 && !(capacity && (status === 503 || status === 529))) return null;
   if (readStreamOutcome(err).replaySafe !== true) return null;
   const { kind, retryAfterMs } = providerErrorKind(err);
-  if (DETERMINISTIC_KINDS.has(kind) || err.code === 'quota_exhausted' || err.quotaScope === 'day' || err.quotaScope === 'plan') {
+  if (
+    DETERMINISTIC_KINDS.has(kind) ||
+    err.code === 'quota_exhausted' ||
+    err.quotaScope === 'day' ||
+    err.quotaScope === 'plan'
+  ) {
     return null;
   }
-  // Capacity follows the local schedule whatever wait the server names.
+  // Capacity waits at least the advised RetryInfo delay (the local schedule
+  // is the floor). Total advised waiting stays within the retry budget; an
+  // advice that does not fit surfaces with its reset time rather than
+  // retrying early.
   if (capacity) {
     if (state.capacityRetries >= ANTIGRAVITY_CAPACITY_RETRIES) return null;
-    return antigravityCapacityDelayMs(state.capacityRetries++);
+    const spent = state.capacityWaitedMs ?? 0;
+    if (retryAfterMs != null && spent + retryAfterMs > ANTIGRAVITY_CAPACITY_BUDGET_MS) {
+      err.resetAt ??= Date.now() + retryAfterMs;
+      return null;
+    }
+    const wait = Math.max(antigravityCapacityDelayMs(state.capacityRetries++), retryAfterMs ?? 0);
+    state.capacityWaitedMs = spent + wait;
+    return wait;
   }
   // An RPM hit gets one quick retry; advice up to a minute is honored in
   // full, a longer wait surfaces so the account pool can rotate.

package/src/runtime/agent/orchestrator/providers/antigravity-stream.mjs

@@ -1,18 +1,15 @@
 /**
  * antigravity-stream.mjs — what one Antigravity SSE turn accumulates while it
  * streams: the live text mirror a retirement notice is recognised from, the
- * ordered parts of every chunk, the native tool calls parsed per chunk with
- * their id dedup, the leak guard that recovers calls the model wrote as text,
- * the first terminal finishReason, and the exposure marks a failed stream must
- * carry (emittedToolCall / unsafeToRetry / partial calls + partial replay).
+ * leak guard that recovers calls the model wrote as text (held back, like the
+ * native calls, until the turn finishes cleanly), and the first terminal
+ * finishReason. Nothing is dispatched while streaming.
  *
  * The transport (antigravity-transport.mjs) owns the socket and the retry
  * policy; the projection (antigravity-response.mjs) reads this record after
  * the turn settles.
  */
-import { createProviderReplay } from './lib/provider-replay.mjs';
 import { createGeminiTextLeakGuard, consumeGeminiRestStreamResponse } from './gemini-stream.mjs';
-import { parseToolCalls, emitGeminiToolCalls } from './gemini-schema.mjs';
 import { normalizeAntigravityUsage } from './antigravity-response.mjs';
 
 // A retired wire id answers with one plain-text notice and no finishReason.
@@ -46,10 +43,6 @@ export function createAntigravityStreamCollector({ tools, useModel, onToolCall,
   // Streamed text is kept so a retirement notice can be told apart from
   // a truncated stream when the gateway omits the finishReason.
   let streamedText = '';
-  let streamedParts = [];
-  let streamedNativeToolCalls = [];
-  const seenNativeToolIds = new Set();
-  const emittedToolIds = new Set();
 
   const relayText = onTextDelta
     ? (text) => {
@@ -57,15 +50,9 @@ export function createAntigravityStreamCollector({ tools, useModel, onToolCall,
         onTextDelta(text);
       }
     : null;
-  const dispatchToolCall = onToolCall
-    ? (call) => {
-        // A failure already observed in this chunk must win over both
-        // native calls and calls recovered from its text.
-        if (terminalFailure || emittedToolIds.has(call.id)) return;
-        emittedToolIds.add(call.id);
-        onToolCall(call);
-      }
-    : null;
+  // Tool calls (native and text-recovered) are never dispatched while the
+  // stream runs: the turn's terminal verdict decides, and
+  // finalizeAntigravityTurn releases them only after a clean finish.
   const onChunk = (chunk) => {
     const candidate = chunk?.candidates?.[0];
     const finishReason =
@@ -74,46 +61,30 @@ export function createAntigravityStreamCollector({ tools, useModel, onToolCall,
     if (finishReason && String(finishReason).replace(/^FINISH_REASON_/, '') !== 'STOP') {
       terminalFailure ||= finishReason;
     }
-    if (!onToolCall) return;
-    const parts = candidate?.content?.parts ?? [];
-    streamedParts.push(...parts);
-    if (terminalFailure || !parts.some((part) => part?.functionCall)) return;
-    // Parse against the turn's parts so anonymous call IDs keep the
-    // same ordinal as final parsing, even across separate SSE chunks.
-    const fresh = (parseToolCalls(streamedParts) || []).filter((call) => {
-      if (seenNativeToolIds.has(call.id)) return false;
-      seenNativeToolIds.add(call.id);
-      return true;
-    });
-    const calls = textLeakGuard?.enabled ? textLeakGuard.filterNativeToolCalls(fresh) : fresh;
-    if (calls?.length) streamedNativeToolCalls.push(...calls);
-    emitGeminiToolCalls(calls, dispatchToolCall);
   };
 
   // Each attempt starts from an empty record: a replayed request must not
-  // inherit the dead attempt's text, parts or dispatched-call ids.
+  // inherit the dead attempt's text or failure.
   const beginAttempt = () => {
     textLeakGuard = createGeminiTextLeakGuard({
       knownToolNames: tools?.map((t) => t.name).filter(Boolean) ?? [],
       onTextDelta: relayText,
-      onToolCall: dispatchToolCall,
+      onToolCall,
       onStreamDelta,
+      deferToolCalls: true,
     });
     streamedText = '';
     terminalFailure = null;
-    streamedParts = [];
-    streamedNativeToolCalls = [];
-    seenNativeToolIds.clear();
-    emittedToolIds.clear();
   };
   const releaseGuard = () => {
     textLeakGuard = null;
   };
 
-  const consume = async (res, attemptSignal) => {
+  const consume = async (res, attemptSignal, { firstByteTimeoutMs } = {}) => {
     try {
       return await consumeGeminiRestStreamResponse(res, {
         signal: attemptSignal,
+        firstByteTimeoutMs,
         onStreamDelta,
         onTextDelta: relayText,
         onChunk,
@@ -125,22 +96,7 @@ export function createAntigravityStreamCollector({ tools, useModel, onToolCall,
         unwrapChunk: (chunk) => (chunk && typeof chunk === 'object' && chunk.response ? chunk.response : chunk),
       });
     } catch (streamErr) {
-      const error = retiredModelError(streamErr, streamedText, useModel) || streamErr;
-      // Native calls now run before EOF. Preserve
-      // their history and prohibit resampling after
-      // a tool callback.
-      if (emittedToolIds.size) {
-        error.emittedToolCall = true;
-        error.unsafeToRetry = true;
-        const leaked = textLeakGuard.getLeakedToolCalls();
-        error.partialToolCalls = [...streamedNativeToolCalls, ...leaked];
-        const replay = createProviderReplay('antigravity', leaked.length ? [] : streamedParts);
-        if (replay) {
-          replay.requestContext = { model: useModel };
-          error.partialProviderReplay = replay;
-        }
-      }
-      throw error;
+      throw retiredModelError(streamErr, streamedText, useModel) || streamErr;
     }
   };
 
@@ -155,11 +111,5 @@ export function createAntigravityStreamCollector({ tools, useModel, onToolCall,
     get terminalFailure() {
       return terminalFailure;
     },
-    get streamedNativeToolCalls() {
-      return streamedNativeToolCalls;
-    },
-    get emittedToolCount() {
-      return emittedToolIds.size;
-    },
   };
 }