Files
OmniRoute/open-sse/executors/antigravity/executeAttempt.ts
Bob.Hou 60580ffeb7 fix(antigravity): streaming passthrough for non-streaming clients (#7408)
* chore(ci): add .mergify.yml to main — Mergify only reads config from the default branch (#7168)

* fix(ci): add the auto-enqueue pull_request_rule to the Mergify config (queue_conditions alone are eligibility-only) (#7179)

* fix(ci): migrate Mergify auto-enqueue to merge_protections_settings.auto_merge_conditions (rules-based path is EOL 2026-07-16) (#7216)

* fix(ci): drop Mergify batch settings (batching is a paid-tier feature; free plan queue is serial) (#7220)

* fix(ci): merge queue tolerates the advisory dast-smoke failure (its GH-hosted build hang dequeued every attempt) (#7225)

* test(ci): make the #6634 selfref guard hermetic — main's copy hard-fails every PR (#7341)

main's copy of this test still does git I/O inside a unit test:

    const baseSrc = git(['show', 'origin/main:' + FILE]);

Runners check out a shallow single ref, so origin/main does not resolve and the
test dies with 'fatal: invalid object name origin/main'. Every PR into main
fails Unit Tests (7/8) on it — today that is #7313, #7315, #7316, #7334, #7336
and #7337, six PRs red on a defect none of them introduced. #7313 has no other
red at all.

release/v3.8.49 already carries a fix (2e42b8efc, #7174: try/catch, fetch
origin/main on demand, t.skip() when unreachable), but it only reaches main at
release time — so main stays broken for the whole cycle. Cherry-picking it would
also import a new problem: PR Test Policy classifies t.skip() as a silenced
assertion, which we watched it correctly catch on #7300 today.

This is the hermetic version instead (ported from #7327, which does the same for
the release branch): read the file straight off disk, compare against an empty
base so baseTaut/baseExtTaut are 0 — the strictest possible comparison point —
and call evaluateMasking() directly. No git ref, no fetch, no skip, nothing the
runner's checkout depth can break.

The #6634 regression stays covered: the guard's logic lives in
SELF_TEST_FIXTURE_RE (check-test-masking.mjs:337), not in the test. Proven both
ways on main before committing — neutralise SELF_TEST_FIXTURE_RE to /$^/ and
the test FAILS; restore it and it passes 2/2, with check-test-masking.mjs left
byte-identical.

Co-authored-by: growab <nekron@icloud.com>

* chore(quality): tighten main's coverage baseline to the CI's real numbers (#7347)

main's ratchet had been failing --require-tighten on every PR: 11 metrics
improved but the baseline was never tightened. Same class as the #6634
selfref guard — an infra fix that lands only on the release branch leaves
main red for the whole cycle, and every PR into main pays for it.

Values are the merged-coverage numbers from a run on main itself (a local
run measures ~68% vs CI's ~80%; the baseline's own note warns about that
gap). Only the 11 coverage values change — gitleaks and semgrepFindings
keep main's own state.

No changelog fragment: #7326 carries it on release/v3.8.49, and a second
one here would double the entry at release time.

* fix(antigravity): remove hardcoded 120s SSE collect timeout

The SSE collection in collectStreamToResponse had a hardcoded 120 s
timeout.  Reasoning-heavy models like gemini-3.1-pro-high on large
prompts (>30 KB) regularly exceed 120 s of generation time, causing
the executor to return a synthetic 504 before the model finishes.

Replace the hardcoded value with FETCH_TIMEOUT_MS (default 600 s,
overridable via FETCH_TIMEOUT_MS env var), which is the standard
upstream-request budget across all OmniRoute providers.

Signed-off-by: Minxi Hou <houminxi@gmail.com>

* fix(antigravity): streaming passthrough for non-streaming clients

When a client sends stream: false to the Antigravity executor
(Gemini models), OmniRoute buffered the entire SSE stream before
responding. Long-thinking models exceeded the 120s timeout.

Remove hardcoded SSE_COLLECT_TIMEOUT_MS. Extract shared
createCreditsExtractionTransform with 16KB buffer cap and abort
handling for client disconnect. Add parseSSEToGeminiResponse for
the non-streaming drain path. Fix hasGeminiTerminalFinishReason
to check top-level candidates (no response wrapper). Add signal
null guards for credits retry path. Return 499 on early abort
instead of piping cancelled body.

Also remove duplicate SKILLS_SANDBOX_RUNTIME from .env.example
and clarify .artifacts/ vs _artifacts/ in .gitignore.

Signed-off-by: Minxi Hou <houminxi@gmail.com>

* refactor(antigravity): extract streaming passthrough to module (file-size cap)

#7408 added the non-streaming SSE pass-through (createCreditsExtractionTransform
plus its two call sites: the credits-retry path and the main non-streaming
path) inline in antigravity.ts, growing it to 1806 lines. Combined with two
other authorized PRs touching the same file (#6979 +11, #7290 +30), the
projected total exceeds the frozen file-size gate (1813).

Extract the new streaming-passthrough logic verbatim into
open-sse/executors/antigravity/streamingPassthrough.ts
(createCreditsExtractionTransform + a new buildSsePassthroughResult that
deduplicates the two near-identical call sites), following the existing
sseCollect.ts submodule pattern -- pure, no host state, no fetch/auth.
antigravity.ts keeps a thin wrapper for createCreditsExtractionTransform
(same public signature the existing unit tests import) that injects
updateAntigravityRemainingCredits so the two modules don't import each
other.

No behavior change: same abort handling, same 499-on-early-disconnect,
same 16KB credits sliding-window cap. antigravity.ts: 1806 -> 1693 lines
(under the 1755 pre-PR baseline, with margin). New module: 176 lines
(cap 800).

Co-authored-by: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com>

* refactor(antigravity): split incremental parser + move new tests to own file (file-size caps)

Two remaining frozen file-size violations from #7408, resolved by
extraction/move with zero behavior or assert changes:

- open-sse/handlers/sseParser.ts (979 > frozen 830): the PR appended
  parseSSEToGeminiResponse (+153, the Gemini buffered-SSE ->
  chat.completion parser). Moved verbatim to
  open-sse/handlers/sseParser/geminiResponse.ts, following the handlers
  submodule pattern (chatCore/, responseSanitizer/). sseParser.ts is now
  byte-identical to its pre-PR content (825 lines; PR delta 0). Importers
  (chatCore/nonStreamingSse.ts, tests) point at the new module.

- tests/unit/executor-antigravity.test.ts (1058 > testFrozen 942): the
  PR's new streaming-passthrough tests moved verbatim (same tests, same
  asserts) to tests/unit/antigravity-streaming-passthrough.test.ts:
  the 3 createCreditsExtractionTransform tests plus the non-streaming
  passthrough drain test ("auto-retries short 429 ... collects SSE for
  non-stream clients"), which the PR rewired onto the new raw-SSE path.
  The frozen file drops to 888 lines (below its pre-PR 941).

New files: geminiResponse.ts 156 lines, passthrough test 202 lines (caps
800). Also fixes the stale sseParser.ts path in collectStreamToResponse's
deprecation note.

Co-authored-by: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com>

* refactor(antigravity): decompose execute + gemini parser below complexity gate

executeOnce() (complexity 127, 436 lines) and parseSSEToGeminiResponse()
(complexity 39, 117 lines) were both over the check-complexity.mjs gate
(complexity>15, max-lines-per-function>80). Decomposed each into small
named helpers, no behavior change:

- geminiResponse.ts: split into pure per-concern functions (markdown
  shortcut, candidate-parts walk, finishReason, usageMetadata, final
  response assembly).
- antigravity.ts: extracted the per-url-index attempt pipeline
  (runAntigravityAttempt, handleAntigravityRateLimit,
  tryResolveRetryFromErrorBody, shouldAutoRetryTransient) and moved the
  request/result-building helpers (send, credits-retry, embed-retry,
  non-streaming/streaming result builders) into a new
  antigravity/executeAttempt.ts submodule, mirroring the existing
  streamingPassthrough.ts/sseCollect.ts pattern. Also fixes the
  antigravity.ts file-size cap (was pushed to 2084 lines > 1813 frozen
  ceiling by the decomposition itself; now 1428).

check-complexity.mjs: 2054 violations (baseline 2058) — net improvement.
execute/executeOnce/parseSSEToGeminiResponse no longer appear with
ruleId complexity or max-lines-per-function.

Co-authored-by: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com>

* Merge branch 'release/v3.8.49' into fix/antigravity-streaming-passthrough

Resolves conflict in open-sse/executors/antigravity.ts between this
branch's streaming-passthrough decomposition and #7290's fallback-chain
decomposition (already merged into release/v3.8.49) — both sides added
imports from the same new antigravity/ submodule files, kept both.

Co-authored-by: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com>

* fix(antigravity): keep buffered JSON contract for non-streaming callers

#3786's Pro-family fallback-chain retry loop (execute()) calls executeOnce()
per candidate and inspects result.response directly, expecting a
synthesized chat.completion JSON body on success. The streaming-passthrough
migration made ALL non-streaming (stream: false) responses a raw SSE
pass-through instead, so a successful retry candidate's response.json()
threw ("data: {...}" is not valid JSON) — breaking the fallback chain
(tests/unit/agy-pro-fallback-chain-3786.test.ts, 3 of 13 red).

Route non-streaming (stream: false) responses back through
collectStreamToResponse (buffered collect-to-JSON), which already uses
FETCH_TIMEOUT_MS with no hardcoded 120s ceiling, so long-thinking models
are not penalized. Passthrough is reserved for actual streaming clients
(stream: true), which was the PR's real target scenario.

Extracted the branch into buildAntigravityAttemptResult() to keep
runAntigravityAttempt under the 80-line ratchet cap.

Co-authored-by: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com>

---------

Signed-off-by: Minxi Hou <houminxi@gmail.com>
Co-authored-by: Diego Rodrigues de Sa e Souza <8016841+diegosouzapw@users.noreply.github.com>
Co-authored-by: growab <nekron@icloud.com>
Co-authored-by: HouMinXi <1000+HouMinXi@users.noreply.github.com>
Co-authored-by: HouMinXi <19586012+HouMinXi@users.noreply.github.com>
2026-07-18 21:19:20 -03:00

688 lines
24 KiB
TypeScript

// Pure-ish per-attempt request/result helpers for the Antigravity executor (#7408
// complexity-gate decomposition): building + sending one upstream request, and
// building the final non-streaming/streaming result. No host state of their own —
// callers inject `provider` and `onCreditsUpdate` so this module doesn't need to
// import the executor's credit-balance cache. Extracted from antigravity.ts
// (file-size cap), mirroring the existing antigravity/streamingPassthrough.ts and
// antigravity/sseCollect.ts submodule pattern.
import { mergeAbortSignals, type ExecutorLog } from "../base.ts";
import { applyFingerprint, isCliCompatEnabled } from "../../config/cliFingerprints.ts";
import { buildAntigravityUpstreamError } from "../antigravityUpstreamError.ts";
import {
HTTP_STATUS,
STREAM_READINESS_TIMEOUT_MS,
ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE,
} from "../../config/constants.ts";
import { injectCreditsField, handleCreditsFailure } from "../../services/antigravityCredits.ts";
import { cloakAntigravityToolPayload } from "../../config/toolCloaking.ts";
import {
applyAntigravityClientProfileHeaders,
removeHeaderCaseInsensitive,
} from "../../services/antigravityClientProfile.ts";
import * as prl from "../../utils/providerRequestLogging.ts";
import {
createCreditsExtractionTransform as createCreditsExtractionTransformImpl,
buildSsePassthroughResult,
type SsePassthroughResult,
} from "./streamingPassthrough.ts";
import type { AntigravityCredentials } from "../antigravity.ts";
const LONG_RETRY_THRESHOLD_MS = 60_000;
const CREDITS_EXHAUSTED_TTL_MS = 5 * 60 * 60 * 1000; // 5 hours
/** Invoked with a fresh GOOGLE_ONE_AI credit balance to persist in the caller's cache. */
export type OnAntigravityCreditsUpdate = (accountId: string, balance: number) => void;
/**
* Per-account GOOGLE_ONE_AI credits-exhausted tracker.
* Key: accountId (OAuth subject / email). Value: expiry timestamp.
* When credits hit 0 we skip the credit retry for CREDITS_EXHAUSTED_TTL_MS.
* Lives here (not antigravity.ts) so both this module's tryCreditsRetry and
* antigravity.ts's tryResolveRetryFromErrorBody can share it via a single import
* direction (antigravity.ts -> executeAttempt.ts), avoiding a circular import.
*/
const MAX_CREDITS_EXHAUSTED_ENTRIES = 50;
const creditsExhaustedUntil = new Map<string, number>();
const _creditsExhaustedSweep = setInterval(() => {
const now = Date.now();
for (const [key, until] of creditsExhaustedUntil) {
if (now >= until) creditsExhaustedUntil.delete(key);
}
}, 60_000);
if (typeof _creditsExhaustedSweep === "object" && "unref" in _creditsExhaustedSweep) {
(_creditsExhaustedSweep as { unref?: () => void }).unref?.();
}
/** True while `accountId`'s Google One AI credits are marked exhausted. @internal */
export function isCreditsExhausted(accountId: string): boolean {
const until = creditsExhaustedUntil.get(accountId);
if (!until) return false;
if (Date.now() >= until) {
creditsExhaustedUntil.delete(accountId);
return false;
}
return true;
}
/** Mark an account's Google One AI credits as exhausted for CREDITS_EXHAUSTED_TTL_MS. */
export function markCreditsExhausted(accountId: string): void {
if (
creditsExhaustedUntil.size >= MAX_CREDITS_EXHAUSTED_ENTRIES &&
!creditsExhaustedUntil.has(accountId)
) {
const now = Date.now();
for (const [key, until] of creditsExhaustedUntil) {
if (now >= until) {
creditsExhaustedUntil.delete(key);
}
}
if (creditsExhaustedUntil.size >= MAX_CREDITS_EXHAUSTED_ENTRIES) {
const oldestKey = creditsExhaustedUntil.keys().next().value;
if (oldestKey !== undefined) creditsExhaustedUntil.delete(oldestKey);
}
}
creditsExhaustedUntil.set(accountId, Date.now() + CREDITS_EXHAUSTED_TTL_MS);
}
class AntigravityPreResponseTimeoutError extends Error {
code = ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE;
status = HTTP_STATUS.GATEWAY_TIMEOUT;
constructor(timeoutMs: number, url: string) {
super(`Antigravity upstream did not return response headers within ${timeoutMs}ms: ${url}`);
this.name = "TimeoutError";
}
}
function getAbortErrorCode(error: unknown): string | null {
if (!error || typeof error !== "object") return null;
const value = (error as { code?: unknown }).code;
return typeof value === "string" ? value : null;
}
function isAntigravityPreResponseTimeout(error: unknown): boolean {
return getAbortErrorCode(error) === ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE;
}
/**
* `fetch()` wrapper that aborts if the upstream never returns response headers
* within `timeoutMs` (default STREAM_READINESS_TIMEOUT_MS) — distinct from the
* overall FETCH_TIMEOUT_MS, which bounds the whole request including body streaming.
* Shared by every fetch attempt in executeOnce() (initial, 403-retry, credits-retry).
*/
export async function fetchAntigravityWithReadinessTimeout(
url: string,
init: RequestInit,
timeoutMs = STREAM_READINESS_TIMEOUT_MS
): Promise<Response> {
const boundedTimeoutMs = Math.max(0, Math.floor(timeoutMs));
if (boundedTimeoutMs <= 0) {
return fetch(url, init);
}
const timeoutController = new AbortController();
let timeoutId: ReturnType<typeof setTimeout> | null = setTimeout(() => {
timeoutController.abort(new AntigravityPreResponseTimeoutError(boundedTimeoutMs, url));
}, boundedTimeoutMs);
const existingSignal = init.signal instanceof AbortSignal ? init.signal : null;
const combinedSignal = existingSignal
? mergeAbortSignals(existingSignal, timeoutController.signal)
: timeoutController.signal;
try {
return await fetch(url, { ...init, signal: combinedSignal });
} catch (error) {
if (
timeoutController.signal.aborted &&
isAntigravityPreResponseTimeout(timeoutController.signal.reason)
) {
throw timeoutController.signal.reason;
}
throw error;
} finally {
if (timeoutId) {
clearTimeout(timeoutId);
timeoutId = null;
}
}
}
/** ExecutorLog with every method always callable — see toSafeAntigravityLog(). */
export type SafeAntigravityLog = Required<ExecutorLog>;
function noopLogFn(): void {}
/**
* Normalize a possibly-null/undefined ExecutorLog into an object with all four
* methods always callable, so the request/retry helpers below can call
* `l.debug(...)` directly instead of repeating `log?.debug?.(...)` at every call
* site. This isn't just style: the complexity linter (eslint `complexity` rule)
* weighs each `?.` link in a chain as its own branch — a doubly-chained
* `log?.debug?.(...)` costs +2 — so a logging-heavy helper can rack up a large
* complexity score with zero real decision points. Resolving once here keeps
* the actual branch count legible in the functions that matter.
*/
export function toSafeAntigravityLog(log: ExecutorLog | null | undefined): SafeAntigravityLog {
return {
debug: log?.debug ? log.debug.bind(log) : noopLogFn,
info: log?.info ? log.info.bind(log) : noopLogFn,
warn: log?.warn ? log.warn.bind(log) : noopLogFn,
error: log?.error ? log.error.bind(log) : noopLogFn,
};
}
/** Flatten a 429/503 error JSON body (message + `error.details[].reason`) into one string. */
export function buildAntigravity429ErrorMessage(errorJson: unknown): string {
const obj = errorJson as
| { error?: { message?: unknown; details?: unknown }; message?: unknown }
| null
| undefined;
let errorMessage = String(obj?.error?.message || obj?.message || "");
const details = obj?.error?.details;
if (Array.isArray(details)) {
for (const detail of details) {
const reason = (detail as { reason?: unknown } | null)?.reason;
if (reason) errorMessage += ` ${reason}`;
}
}
return errorMessage;
}
function getChunkedOrFixedBody(bodyStr: string, stream: boolean): BodyInit {
if (stream) {
return new ReadableStream(
{
async start(controller) {
controller.enqueue(new TextEncoder().encode(bodyStr));
controller.close();
},
},
{ highWaterMark: 16384 }
);
}
return bodyStr;
}
function cloneAntigravityRequestBody(body: unknown): unknown {
if (!body || typeof body !== "object") {
return body;
}
try {
return structuredClone(body);
} catch {
return JSON.parse(JSON.stringify(body));
}
}
function serializeAntigravityRequest(
provider: string,
headers: Record<string, string>,
body: unknown
): { headers: Record<string, string>; bodyString: string } {
const serializedBody = cloneAntigravityRequestBody(body);
if (!isCliCompatEnabled(provider)) {
return { headers, bodyString: JSON.stringify(serializedBody) };
}
return applyFingerprint(provider, { ...headers }, serializedBody);
}
function getRequestTargetModel(body: Record<string, unknown>): string {
const target = body.model;
return typeof target === "string" && target.length > 0 ? target : "unknown";
}
function attachToolNameMap<T>(payload: T, toolNameMap: Map<string, string> | null): T {
if (!toolNameMap?.size || !payload || typeof payload !== "object") {
return payload;
}
const copy = Array.isArray(payload) ? ([...payload] as T) : ({ ...(payload as object) } as T);
Object.defineProperty(copy, "_toolNameMap", {
value: toolNameMap,
enumerable: false,
configurable: true,
writable: true,
});
return copy;
}
/** Cloak the tool-name payload, then apply credits-first injection, for one attempt. */
export function finalizeAntigravityRequestBody(
transformed: Record<string, unknown>,
useCreditsFirst: boolean,
log: SafeAntigravityLog
): {
transformedBody: Record<string, unknown>;
requestToolNameMap: Map<string, string> | null;
} {
let transformedBody: Record<string, unknown> = transformed;
let requestToolNameMap: Map<string, string> | null = null;
if (transformedBody && typeof transformedBody === "object") {
const cloaked = cloakAntigravityToolPayload(transformedBody);
transformedBody = cloaked.body;
requestToolNameMap = cloaked.toolNameMap;
}
// Credits-first: inject GOOGLE_ONE_AI upfront so we never try the normal
// quota path. If credits are exhausted / disabled shouldUseCreditsFirst()
// returns false and we fall back to the legacy retry-on-429 flow.
if (useCreditsFirst) {
transformedBody = injectCreditsField(transformedBody);
log.debug("AG_CREDITS", "Credits-first enabled (ANTIGRAVITY_CREDITS=always)");
}
return { transformedBody, requestToolNameMap };
}
/** Debug-only dump of outgoing headers (mask Authorization) and envelope shape. */
function dumpAntigravityRequestDebug(
finalHeaders: Record<string, string>,
transformedBody: Record<string, unknown>,
clientProfile: unknown,
log: SafeAntigravityLog
): void {
const safeHeaders = { ...finalHeaders };
if (safeHeaders["Authorization"]) safeHeaders["Authorization"] = "Bearer ***";
log.debug("AG_REQUEST_HEADERS", JSON.stringify(safeHeaders));
const envelope = transformedBody as Record<string, unknown>;
const requestInner = envelope.request as Record<string, unknown> | undefined;
log.debug(
"AG_REQUEST_ENVELOPE",
JSON.stringify({
fieldOrder: Object.keys(envelope),
project: envelope.project,
requestId: envelope.requestId,
model: envelope.model,
userAgent: envelope.userAgent,
requestType: envelope.requestType,
enabledCreditTypes: envelope.enabledCreditTypes,
clientProfile,
sessionId: requestInner?.sessionId,
generationConfig: requestInner?.generationConfig,
})
);
}
/**
* Send one Antigravity request attempt: serialize + apply the client-profile
* fingerprint, debug-dump the outgoing envelope, fetch with a readiness timeout,
* and transparently retry once without `x-goog-user-project` on a 403 (some
* projects reject that header). Returns the (possibly 403-retried) response
* plus the headers actually used for it.
*/
export async function sendAntigravityRequest(
provider: string,
url: string,
model: string,
headers: Record<string, string>,
transformedBody: Record<string, unknown>,
credentials: AntigravityCredentials,
stream: boolean,
signal: AbortSignal | null | undefined,
log: SafeAntigravityLog,
retryAttempt: number
): Promise<{ response: Response; finalHeaders: Record<string, string> }> {
const serializedRequest = serializeAntigravityRequest(provider, headers, transformedBody);
let finalHeaders = serializedRequest.headers;
const clientProfile = applyAntigravityClientProfileHeaders(
finalHeaders,
credentials,
transformedBody
);
log.debug(
"TELEMETRY",
`[Antigravity] Execute - URL: ${url}, Model: ${model}, Target: ${getRequestTargetModel(transformedBody)}, RetryAttempt: ${retryAttempt}`
);
// Dump outgoing headers (mask Authorization) and envelope shape for debugging.
// Gated behind an explicit typeof check (not just calling log.debug() unconditionally)
// so the JSON.stringify work below is skipped entirely when debug logging is off.
if (typeof log.debug === "function") {
dumpAntigravityRequestDebug(finalHeaders, transformedBody, clientProfile, log);
}
await prl.captureCurrentProviderBody(url, finalHeaders, serializedRequest.bodyString, log);
let response = await fetchAntigravityWithReadinessTimeout(url, {
method: "POST",
headers: finalHeaders,
body: getChunkedOrFixedBody(serializedRequest.bodyString, stream),
...(stream ? { duplex: "half" } : {}),
signal,
});
if (response.status === HTTP_STATUS.FORBIDDEN && finalHeaders["x-goog-user-project"]) {
const retryHeaders = { ...finalHeaders };
removeHeaderCaseInsensitive(retryHeaders, "x-goog-user-project");
log.debug("RETRY", "403 with x-goog-user-project, retrying once without it");
await prl.captureCurrentProviderBody(url, retryHeaders, serializedRequest.bodyString, log);
response = await fetchAntigravityWithReadinessTimeout(url, {
method: "POST",
headers: retryHeaders,
body: getChunkedOrFixedBody(serializedRequest.bodyString, stream),
...(stream ? { duplex: "half" } : {}),
signal,
});
finalHeaders = retryHeaders;
}
if (!response.ok) {
log.warn(
"TELEMETRY",
`[Antigravity] Error Response - URL: ${url}, Status: ${response.status}, Model: ${model}`
);
}
return { response, finalHeaders };
}
/**
* Retry the SAME url with `enabledCreditTypes: ["GOOGLE_ONE_AI"]` injected, for a
* quota_exhausted 429 that hasn't already tried credits. Returns the result to hand
* back to the caller of execute() on success (or a non-429 status), or null if the
* credits retry also failed/429'd (caller falls through to the normal retry logic).
*/
export async function tryCreditsRetry(
provider: string,
url: string,
headers: Record<string, string>,
transformedBody: Record<string, unknown>,
requestToolNameMap: Map<string, string> | null,
credentials: AntigravityCredentials,
stream: boolean,
signal: AbortSignal | null | undefined,
log: SafeAntigravityLog,
accountId: string,
onCreditsUpdate: OnAntigravityCreditsUpdate
): Promise<SsePassthroughResult | null> {
log.info("AG_CREDITS", "Retrying with Google One AI credits");
const creditsBody = injectCreditsField(transformedBody);
const serializedCreditsRequest = serializeAntigravityRequest(provider, headers, creditsBody);
const finalCreditsHeaders = serializedCreditsRequest.headers;
try {
await prl.captureCurrentProviderBody(
url,
finalCreditsHeaders,
serializedCreditsRequest.bodyString,
log
);
const creditsResp = await fetchAntigravityWithReadinessTimeout(url, {
method: "POST",
headers: finalCreditsHeaders,
body: getChunkedOrFixedBody(serializedCreditsRequest.bodyString, stream),
...(stream ? { duplex: "half" } : {}),
signal,
});
if (creditsResp.ok || creditsResp.status !== HTTP_STATUS.RATE_LIMITED) {
log.info("AG_CREDITS", `Credits retry succeeded: ${creditsResp.status}`);
if (!stream && creditsResp.body) {
// Raw SSE pass-through + credits extraction (see
// streamingPassthrough.ts); 499s early if the client
// already disconnected instead of piping a cancelled body.
return buildSsePassthroughResult(
creditsResp.body,
creditsResp,
accountId,
onCreditsUpdate,
url,
finalCreditsHeaders,
attachToolNameMap(creditsBody, requestToolNameMap),
signal
);
}
return {
response: creditsResp,
url,
headers: finalCreditsHeaders,
transformedBody: attachToolNameMap(creditsBody, requestToolNameMap),
};
}
// Credit retry also 429'd
handleCreditsFailure(credentials?.accessToken || "");
log.warn("AG_CREDITS", "Credits retry also 429'd");
// Also mark in our legacy exhaustion map to avoid retrying other routes
markCreditsExhausted(accountId);
return null;
} catch (creditsErr) {
handleCreditsFailure(credentials?.accessToken || "");
log.warn("AG_CREDITS", `Credits retry failed: ${creditsErr}`);
return null;
}
}
/**
* If we have a 429 with a long retry time (> LONG_RETRY_THRESHOLD_MS), embed
* `retryAfterMs` in the response body so the caller (combo/account-fallback
* layer) can read it back out. Returns null (fall back to the original
* response handling) when the status/retryMs don't qualify, or on error.
*/
export async function tryEmbedLongRetryAfter(
response: Response,
retryMs: number | null,
url: string,
finalHeaders: Record<string, string>,
transformedBody: Record<string, unknown>,
requestToolNameMap: Map<string, string> | null,
log: ExecutorLog | null | undefined
): Promise<SsePassthroughResult | null> {
if (
response.status !== HTTP_STATUS.RATE_LIMITED ||
!retryMs ||
retryMs <= LONG_RETRY_THRESHOLD_MS
) {
return null;
}
try {
const respBody = await response.clone().text();
let obj;
try {
obj = JSON.parse(respBody);
} catch {
obj = {};
}
obj.retryAfterMs = retryMs;
const modifiedBody = JSON.stringify(obj);
const modifiedResponse = new Response(modifiedBody, {
status: response.status,
headers: response.headers,
});
return {
response: modifiedResponse,
url,
headers: finalHeaders,
transformedBody: attachToolNameMap(transformedBody, requestToolNameMap),
};
} catch (err) {
log?.warn?.("RETRY", `Failed to embed retryAfterMs: ${err}`);
return null;
}
}
/** Build the sanitized JSON error result shared by the non-streaming and streaming paths. */
async function buildUpstreamErrorResult(
response: Response,
url: string,
finalHeaders: Record<string, string>,
transformedBody: Record<string, unknown>,
requestToolNameMap: Map<string, string> | null
): Promise<SsePassthroughResult> {
const rawBody = await response
.clone()
.text()
.catch(() => "");
const errorBody = buildAntigravityUpstreamError(response.status, response.statusText, rawBody);
return {
response: new Response(JSON.stringify(errorBody), {
status: response.status,
headers: { "Content-Type": "application/json" },
}),
url,
headers: finalHeaders,
transformedBody: attachToolNameMap(transformedBody, requestToolNameMap),
};
}
/**
* For non-streaming clients, return the raw SSE stream with a
* credits-extraction TransformStream. chatCore's non-streaming path
* (readNonStreamingResponseBody + parseNonStreamingSSEPayload with
* Gemini format support) handles draining and conversion to JSON.
* This replaces the previous collectStreamToResponse() approach which
* had an artificial timeout (now the standard FETCH_BODY_TIMEOUT_MS
* of 10 min applies).
*/
async function buildNonStreamingExecuteOnceResult(
response: Response,
url: string,
finalHeaders: Record<string, string>,
transformedBody: Record<string, unknown>,
requestToolNameMap: Map<string, string> | null,
accountId: string,
signal: AbortSignal | null | undefined,
onCreditsUpdate: OnAntigravityCreditsUpdate
): Promise<SsePassthroughResult> {
// #3229: surface a real upstream error instead of masking a 4xx/5xx as an
// empty `chat.completion` envelope.
if (!response.ok) {
return buildUpstreamErrorResult(response, url, finalHeaders, transformedBody, requestToolNameMap);
}
if (response.body) {
// Raw SSE pass-through + credits extraction (see
// streamingPassthrough.ts); 499s early if the client already
// disconnected instead of piping a cancelled body.
return buildSsePassthroughResult(
response.body,
response,
accountId,
onCreditsUpdate,
url,
finalHeaders,
attachToolNameMap(transformedBody, requestToolNameMap),
signal
);
}
// No body -- return as-is
return {
response,
url,
headers: finalHeaders,
transformedBody: attachToolNameMap(transformedBody, requestToolNameMap),
};
}
/**
* Streaming path: wrap the response body in a pass-through TransformStream
* that extracts remainingCredits from the final SSE chunk(s) without
* consuming the stream. The client receives the unmodified SSE data.
*
* #2461: a non-ok upstream response (e.g. 403) must never be piped through the
* streaming pass-through below as if it were an SSE body. Google occasionally
* returns non-UTF8/binary error bodies (observed: gzip-magic-byte payloads) for
* 403s on this endpoint; reading/forwarding those raw bytes corrupts the
* client-visible error message. Mirror the non-streaming branch above and build
* a sanitized JSON error via buildAntigravityUpstreamError (hard rule #12)
* instead of streaming unknown bytes straight through.
*/
async function buildStreamingExecuteOnceResult(
response: Response,
url: string,
finalHeaders: Record<string, string>,
transformedBody: Record<string, unknown>,
requestToolNameMap: Map<string, string> | null,
accountId: string,
signal: AbortSignal | null | undefined,
onCreditsUpdate: OnAntigravityCreditsUpdate
): Promise<SsePassthroughResult> {
if (!response.ok) {
return buildUpstreamErrorResult(response, url, finalHeaders, transformedBody, requestToolNameMap);
}
if (response.body) {
// If the downstream client aborts, cancel the upstream fetch body immediately
// to release the socket back to the Undici agent pool and prevent memory leaks.
if (signal) {
const abortHandler = () => {
try {
response.body?.cancel().catch(() => {});
} catch (_) {}
};
if (signal.aborted) {
abortHandler();
} else {
signal.addEventListener("abort", abortHandler, { once: true });
}
}
const passThrough = createCreditsExtractionTransformImpl(
accountId,
onCreditsUpdate,
16 * 1024 // 16KB sliding-window cap to prevent OOM
);
const tappedBody = response.body.pipeThrough(passThrough);
const tappedResponse = new Response(tappedBody, {
status: response.status,
statusText: response.statusText,
headers: response.headers,
});
return {
response: tappedResponse,
url,
headers: finalHeaders,
transformedBody: attachToolNameMap(transformedBody, requestToolNameMap),
};
}
return {
response,
url,
headers: finalHeaders,
transformedBody: attachToolNameMap(transformedBody, requestToolNameMap),
};
}
/** Dispatch to the non-streaming or streaming final-result builder. */
export async function buildFinalAntigravityResult(
stream: boolean,
response: Response,
url: string,
finalHeaders: Record<string, string>,
transformedBody: Record<string, unknown>,
requestToolNameMap: Map<string, string> | null,
accountId: string,
signal: AbortSignal | null | undefined,
onCreditsUpdate: OnAntigravityCreditsUpdate
): Promise<SsePassthroughResult> {
if (!stream) {
return buildNonStreamingExecuteOnceResult(
response,
url,
finalHeaders,
transformedBody,
requestToolNameMap,
accountId,
signal,
onCreditsUpdate
);
}
return buildStreamingExecuteOnceResult(
response,
url,
finalHeaders,
transformedBody,
requestToolNameMap,
accountId,
signal,
onCreditsUpdate
);
}