pjandro/OmniRoute
1
0
Fork 0
mirror of https://github.com/diegosouzapw/OmniRoute.git synced 2026-09-17 20:32:25 +03:00
Code Issues Packages Projects Releases Wiki Activity

feat(sse): reserve the Antigravity account for the request's stream lifecycle

Browse Source 
Re-land of the account-lease half of #10011 on the current release branch.
Its exact-model-scoping half had already shipped in #8050 and its quota half
lost to the tip's aggregate-family design (selectAntigravityQuotaWindowNames /
antigravityQuotaFamily.ts); none of that is reintroduced here. The lease is a
concurrency reservation only and never reads or writes quota state.

The Antigravity account selected for a request is reserved for the whole
streaming lifecycle of that request, so a concurrent retry — or the credential
handoff inside getProviderCredentialsWithQuotaPreflight — cannot re-pick an
account already committed to an in-flight upstream stream. The reservation is
scoped to (connection, callable upstream model) rather than the whole account,
so one account can still serve two different models at once; catalog ids that
resolve to the same upstream id (the gemini-3.7-flash tiers, all
gemini-3.7-flash-tiered) share one lease. When every eligible account is leased
for that model the request returns a structured 503 antigravity_pool_busy with
a bounded Retry-After instead of piling onto a busy account.

Opt-in behind ANTIGRAVITY_ACCOUNT_LEASE_ENABLED (runtime, default false). With
the flag off no reservation is taken, credentials carry no routing descriptor,
every release/hold is a no-op on an undefined lease id, and account selection
and dispatch behave exactly as before.

#10011's original test suite asserted family semantics for a lease that was
exact-model scoped and failed deterministically on its own head; the model ids
it used (gemini-3.5-flash / gemini-3-flash-agent) no longer exist in the
catalog. The contradiction is resolved in favour of one coherent semantic —
exact callable upstream model — and the tests assert it against the alias
tables as they are on this branch.

Co-authored-by: Ardem2025 <openclaw-auto@example.invalid>
This commit is contained in:
diegosouzapw
2026-09-16 20:55:55 -03:00
parent a2c6f7188d
commit 0321bc8e2e
16 changed files with 925 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
changelog.d/fixes/10011-antigravity-account-lease.md Normal file
View File

@@ -0,0 +1,5 @@
- **fix(sse):** the Antigravity account picked for a request can now be reserved for that
request's streaming lifecycle, so a concurrent retry or the credential handoff cannot re-pick
an account already committed to an in-flight stream; a fully leased pool answers with a
structured 503 `antigravity_pool_busy` carrying a bounded `Retry-After`. Opt-in behind the
new `ANTIGRAVITY_ACCOUNT_LEASE_ENABLED` flag (default off) (#10011) — thanks @Ardem2025

5
config/quality/file-size-baseline.json
View File

@@ -1,4 +1,5 @@
{
"_rebaseline_2026_09_16_10011_antigravity_account_lease": "Re-land of #10011 (@Ardem2025): the Antigravity account lease. src/sse/handlers/chat.ts->2533 (+35), src/sse/services/auth.ts->3577 (+20). The lease registry, its lifecycle glue and its selection glue were extracted into three NEW modules (src/sse/services/antigravityRoutingState.ts, antigravityLeaseLifecycle.ts, antigravityLeaseSelection.ts) precisely to keep this growth to the call sites; what remains in chat.ts/auth.ts is the wiring itself, which cannot be moved out of the selection loop and the dispatch path. Every added hunk is inert unless ANTIGRAVITY_ACCOUNT_LEASE_ENABLED (default false) is on. Covered by tests/unit/antigravity-routing-state.test.ts, antigravity-lease-lifecycle.test.ts and antigravity-account-lease-flag.test.ts.",
"_rebaseline_2026_09_16_jxnlexn_wave_growth": "Combined growth of the 2026-09-15 maxmad64bis uplift batch (each PR rebaselined its own growth; the merged sum is larger): src/sse/handlers/chat.ts->2498; open-sse/handlers/chatCore.ts->6181. Every hunk is flag-gated or a verified fix covered by that PR's tests; see the batch report.",
"_rebaseline_2026_09_16_wave22_growth": "Combined growth of the 2026-09-15 maxmad64bis uplift batch (each PR rebaselined its own growth; the merged sum is larger): src/sse/handlers/chatHelpers.ts->1231; open-sse/executors/cursor.ts->1808. Every hunk is flag-gated or a verified fix covered by that PR's tests; see the batch report.",
"_rebaseline_2026_09_15_13572_combined_growth": "Combined growth of the 2026-09-15 maxmad64bis uplift batch (each PR rebaselined its own growth; the merged sum is larger): open-sse/executors/base.ts->1754. Every hunk is flag-gated or a verified fix covered by that PR's tests; see the batch report.",
@@ -501,8 +502,8 @@
"src/shared/components/RequestLoggerV2.tsx": 1718,
"src/shared/constants/providers/apikey/gateways.ts": 1502,
"src/shared/services/cliRuntime.ts": 1296,
"src/sse/handlers/chat.ts": 2498,
"src/sse/services/auth.ts": 3557,
"src/sse/handlers/chat.ts": 2533,
"src/sse/services/auth.ts": 3577,
"tests/unit/account-fallback-service.test.ts": 2453,
"tests/unit/provider-validation-specialty.test.ts": 4656,
"open-sse/services/autoCombo/virtualFactory.ts": 1230,

73
docs/reference/FEATURE_FLAGS.md
View File

@@ -46,7 +46,7 @@ A boolean flag is considered **enabled** when its effective value is `"true"`,
## Flag Catalog
72 flags across 6 categories. **Default** is the definition default — the value
73 flags across 6 categories. **Default** is the definition default — the value
used when neither a DB override nor an environment variable is present.
### Security (10)
@@ -94,42 +94,43 @@ used when neither a DB override nor an environment variable is present.
| `CAPABILITY_FILTER_ENABLED` | boolean | `false` | Reject requests before dispatch when the target model lacks required capabilities (vision, tools, structured output, context window). Protects direct single-provider requests that bypass the combo-layer compatibility filter. |
| `RADAR_ENABLED` | boolean | `false` | Enable the OmniRoute Radar module (catalog feed screens and sync). Off by default; enabling only unlocks the UI — data sync remains a separate opt-in. |
### Runtime (32)
### Runtime (33)
| Key | Type | Default | Restart | Description |
| ------------------------------------------- | ------- | ------- | ------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `UNIVERSAL_CONTEXT_HANDOFF_ENABLED` | boolean | `true` | | Generate and inject conversation summaries when combo routing switches models. Disable to treat model switches independently and prevent background handoff requests for all existing and future combos. |
| `RESPONSES_PASSTHROUGH_DROP_COMMENTARY` | boolean | `true` | | Drop internal commentary-phase output items from Responses API passthrough streams before forwarding to clients. Disable to receive raw upstream commentary. |
| `OMNIROUTE_MCP_ENFORCE_SCOPES` | boolean | `true` | | Enforce scope restrictions on MCP tool access. |
| `OMNIROUTE_MCP_COMPRESS_DESCRIPTIONS` | boolean | `false` | | Compress MCP tool descriptions to reduce token usage. |
| `OMNIROUTE_ENABLE_RUNTIME_BACKGROUND_TASKS` | boolean | `false` | | Enable background task processing at runtime. |
| `OMNIROUTE_DISABLE_BACKGROUND_SERVICES` | boolean | `false` | ✓ | Disable all background services (quota refresh, sync, etc). |
| `OMNIROUTE_RTK_TRUST_PROJECT_FILTERS` | boolean | `false` | | Trust project-level RTK filters without validation. |
| `OMNIROUTE_ENABLE_LIVE_WS` | boolean | `true` | ✓ | Start the real-time dashboard WebSocket server on import (port 20132 by default). |
| `OMNIROUTE_CODEX_WS_ENABLED` | boolean | `true` | | Allow Codex to use the Responses-over-WebSocket transport. When off, Codex falls back to HTTP Responses. |
| `OMNIROUTE_CODEX_APP_SERVER_ENABLED` | boolean | `true` | | Allow Codex to use the local app-server WebSocket JSON-RPC transport (codexTransport=app-server). When off, connections opted into app-server fall back to Codex's other transports. |
| `OMNIROUTE_EMERGENCY_FALLBACK` | boolean | `true` | | Route budget-exhausted requests to the emergency free fallback provider/model. (See [Emergency Budget Fallback](#emergency-budget-fallback) below.) |
| `STREAM_RECOVERY_ENABLED` | boolean | `false` | | Enable transparent early retry for truncated upstream SSE streams before any response bytes reach the client. |
| `STREAM_RECOVERY_MIDSTREAM_ENABLED` | boolean | `false` | | Allow stream recovery to re-request and stitch a response after bytes have already reached the client. |
| `STREAM_RECOVERY_TOOLCALL_ORDER_FIX` | boolean | `false` | | Make mid-stream continuation tool-call safe: never resume a cut stream once a tool call was emitted (in flight or already finished with finish_reason tool_calls), and close after one empty continuation instead of spending the whole budget. Off: release behavior. |
| `STREAM_EARLY_EOF_SIBLING_FAILOVER_ENABLED` | boolean | `false` | | Fail over once to a sibling connection when an SSE stream closes before emitting any useful frame and the bounded same-connection retry is spent; with no usable sibling the original `STREAM_EARLY_EOF` 502 is returned. Off by default: early-EOF stays terminal after the same-connection retry. |
| `MODEL_CATALOG_INCLUDE_NAMES` | boolean | `true` | | Include display-friendly name fields in `/v1/models` responses. Disable for clients that expect model IDs only. |
| `MODELS_CATALOG_PREFIX_MODE` | enum | `dual` | | Controls how model IDs are prefixed in /v1/models. 'dual' (default) emits both alias and canonical provider-id prefixes for backward compatibility. 'alias' emits only the short alias prefix (e.g. ds-web/model, not deepseek-web/model). 'canonical' emits only the full provider-id prefix. Values: `dual`, `alias`, `canonical`. |
| `ARENA_ELO_SYNC_ENABLED` | boolean | `true` | | Enable periodic Arena AI leaderboard ELO sync for model intelligence rankings. |
| `EXPOSE_CC_DISCOVERY_ALIASES` | boolean | `false` | | Advertise `claude/<provider>/<model>` mirror ids on `/v1/models` so Claude Code gateway model discovery lists non-Claude models. Global level of the three-level gate (env wins over the dashboard override). See [Claude Code configuration](../guides/CLAUDE-CODE-CONFIGURATION.md#discovery-aliases--surface-non-claude-models-in-the-model-picker). |
| `NO_THINKING_ALIAS_ENABLED` | boolean | `true` | | Master switch for the no-think/<provider>/<model> gateway aliases. On (default): /v1/models advertises a no-thinking variant for every eligible thinking-capable Claude model, and a no-think/ id sent on a request resolves back to the real model with reasoning suppressed. Off: no variants are advertised and a no-think/ id is treated like any other unknown model id. The per-model ModelSpec.noThinkingAlias opt-in/opt-out still applies while this is on. |
| `OMNIROUTE_DISABLE_THINKING_LEVEL_VARIANTS` | boolean | `false` | | Disable the generation of thinking level variants (e.g. -low, -medium, -high) in the /v1/models catalog. |
| `OMNIROUTE_CHAT_VIRTUAL_LANES` | boolean | `false` | ✓ | Enable per-tenant adaptive virtual admission lanes for provider dispatch (#9654): one tenant's burst no longer 503s another. The OMNIROUTE_CHAT_VIRTUAL_LANES env var wins over this dashboard override; changes take effect at server restart. |
| `EXPOSE_FUNCTIONAL_GATEWAY_MIRRORS` | boolean | `false` | | Advertise <gateway-alias>/<model> mirror ids on /v1/models for models whose canonical owner has no active credential but a passthrough gateway with an active credential routes them. Warning: adds catalog entries for all clients when enabled globally. |
| `NEWAPI_AGGREGATOR_BALANCE` | boolean | `false` | | Enable balance detection for New-API / One-API / Sub2API aggregator compatible nodes. When enabled, compatible nodes with the aggregator flag set will report their balance in the dashboard and quota-preflight routing. |
| `SERVER_OWNED_TOOL_LOOP_ENABLED` | boolean | `false` | | Continue non-streaming server-owned tool calls until the model returns a client-usable response. |
| `SEARCH_STATS_HIDE_DELETED_CONNECTIONS` | boolean | `false` | | Search stats and recent searches only count providers that still have a live connection (keyless providers such as duckduckgo-free always count). Off keeps every retained search row with a provider id. |
| `FREE_BADGE_REQUIRES_PROVIDER_FREE_TIER` | boolean | `false` | | Dashboard provider pages: show the Free badge only on signals the provider honors — drops the display-name heuristic, non-boolean free fields and :free suffixes on registered providers without a documented free tier. Off keeps the historical badge rule. |
| `RETRY_AFTER_PROVENANCE_ENABLED` | boolean | `false` | | On aggregated 429/503 unavailable responses, omit `Retry-After` when no concrete future retry time is known (instead of a synthetic 1s), add `error.retry_after_provenance` (`signal` \| `none`), and let combo drain paths read prose retry hints from JSON and plain-text upstream bodies. The field only appears on responses built by `unavailableResponse()`; other 429/503 bodies are unchanged. |
| `PROTECTED_PRIORITY_INFRA_502_ENABLED` | boolean | `false` | | When a `priority` combo target marked fallback-only-on-quota-exhaustion stops the combo for a cause that is provably not quota (provider circuit breaker open, predictive latency skip), answer 502 instead of the quota-looking 503. Lockout, cooldown, unavailable, exhaustion and concurrency-cap stops keep 503. |
| `MISTRAL_AMBIGUOUS_401_SOFT_LOCKOUT` | boolean | `false` | | A bare Mistral 401 (`{"detail":"Unauthorized"}`, no explicit auth signal) is identical for a revoked key and for exhausted quota. When on, it cools the connection down instead of parking it as `expired`, at most 3 times per hour per connection; the next one parks it, so a revoked key still converges. Off by default: every bare Mistral 401 parks the connection as before. |
| `XAI_OAUTH_LIVE_MODEL_DISCOVERY` | boolean | `false` | | Fetch the live xAI model catalog for `xai-oauth` connections from `https://api.x.ai/v1/models` using the OAuth bearer token, instead of the frozen static seed. Off by default: `xai-oauth` keeps serving the static seed unchanged. On any resolution error, discovery falls back to the seed (unverified whether x.ai accepts an OAuth bearer at this endpoint). |
| `BATCH_AND_FILE_AUTO_CLEANUP_ENABLED` | boolean | `false` | | Let the automatic cleanup sweep delete terminal (completed/failed/cancelled/expired) Batch API jobs older than `OMNIROUTE_BATCH_RETENTION_DAYS`, along with their per-line checkpoints, and clear the BLOB content of uploaded files past their own `expires_at`. Off by default: every existing install keeps this data exactly as before until an operator opts in. The operator-triggered `DELETE /api/v1/batches/delete-completed` route is unaffected either way — it is a separate, unconditional public API contract. |
| Key | Type | Default | Restart | Description |
| ------------------------------------------- | ------- | ------- | ------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `UNIVERSAL_CONTEXT_HANDOFF_ENABLED` | boolean | `true` | | Generate and inject conversation summaries when combo routing switches models. Disable to treat model switches independently and prevent background handoff requests for all existing and future combos. |
| `RESPONSES_PASSTHROUGH_DROP_COMMENTARY` | boolean | `true` | | Drop internal commentary-phase output items from Responses API passthrough streams before forwarding to clients. Disable to receive raw upstream commentary. |
| `OMNIROUTE_MCP_ENFORCE_SCOPES` | boolean | `true` | | Enforce scope restrictions on MCP tool access. |
| `OMNIROUTE_MCP_COMPRESS_DESCRIPTIONS` | boolean | `false` | | Compress MCP tool descriptions to reduce token usage. |
| `OMNIROUTE_ENABLE_RUNTIME_BACKGROUND_TASKS` | boolean | `false` | | Enable background task processing at runtime. |
| `OMNIROUTE_DISABLE_BACKGROUND_SERVICES` | boolean | `false` | ✓ | Disable all background services (quota refresh, sync, etc). |
| `OMNIROUTE_RTK_TRUST_PROJECT_FILTERS` | boolean | `false` | | Trust project-level RTK filters without validation. |
| `OMNIROUTE_ENABLE_LIVE_WS` | boolean | `true` | ✓ | Start the real-time dashboard WebSocket server on import (port 20132 by default). |
| `OMNIROUTE_CODEX_WS_ENABLED` | boolean | `true` | | Allow Codex to use the Responses-over-WebSocket transport. When off, Codex falls back to HTTP Responses. |
| `OMNIROUTE_CODEX_APP_SERVER_ENABLED` | boolean | `true` | | Allow Codex to use the local app-server WebSocket JSON-RPC transport (codexTransport=app-server). When off, connections opted into app-server fall back to Codex's other transports. |
| `OMNIROUTE_EMERGENCY_FALLBACK` | boolean | `true` | | Route budget-exhausted requests to the emergency free fallback provider/model. (See [Emergency Budget Fallback](#emergency-budget-fallback) below.) |
| `STREAM_RECOVERY_ENABLED` | boolean | `false` | | Enable transparent early retry for truncated upstream SSE streams before any response bytes reach the client. |
| `STREAM_RECOVERY_MIDSTREAM_ENABLED` | boolean | `false` | | Allow stream recovery to re-request and stitch a response after bytes have already reached the client. |
| `STREAM_RECOVERY_TOOLCALL_ORDER_FIX` | boolean | `false` | | Make mid-stream continuation tool-call safe: never resume a cut stream once a tool call was emitted (in flight or already finished with finish_reason tool_calls), and close after one empty continuation instead of spending the whole budget. Off: release behavior. |
| `STREAM_EARLY_EOF_SIBLING_FAILOVER_ENABLED` | boolean | `false` | | Fail over once to a sibling connection when an SSE stream closes before emitting any useful frame and the bounded same-connection retry is spent; with no usable sibling the original `STREAM_EARLY_EOF` 502 is returned. Off by default: early-EOF stays terminal after the same-connection retry. |
| `MODEL_CATALOG_INCLUDE_NAMES` | boolean | `true` | | Include display-friendly name fields in `/v1/models` responses. Disable for clients that expect model IDs only. |
| `MODELS_CATALOG_PREFIX_MODE` | enum | `dual` | | Controls how model IDs are prefixed in /v1/models. 'dual' (default) emits both alias and canonical provider-id prefixes for backward compatibility. 'alias' emits only the short alias prefix (e.g. ds-web/model, not deepseek-web/model). 'canonical' emits only the full provider-id prefix. Values: `dual`, `alias`, `canonical`. |
| `ARENA_ELO_SYNC_ENABLED` | boolean | `true` | | Enable periodic Arena AI leaderboard ELO sync for model intelligence rankings. |
| `EXPOSE_CC_DISCOVERY_ALIASES` | boolean | `false` | | Advertise `claude/<provider>/<model>` mirror ids on `/v1/models` so Claude Code gateway model discovery lists non-Claude models. Global level of the three-level gate (env wins over the dashboard override). See [Claude Code configuration](../guides/CLAUDE-CODE-CONFIGURATION.md#discovery-aliases--surface-non-claude-models-in-the-model-picker). |
| `NO_THINKING_ALIAS_ENABLED` | boolean | `true` | | Master switch for the no-think/<provider>/<model> gateway aliases. On (default): /v1/models advertises a no-thinking variant for every eligible thinking-capable Claude model, and a no-think/ id sent on a request resolves back to the real model with reasoning suppressed. Off: no variants are advertised and a no-think/ id is treated like any other unknown model id. The per-model ModelSpec.noThinkingAlias opt-in/opt-out still applies while this is on. |
| `OMNIROUTE_DISABLE_THINKING_LEVEL_VARIANTS` | boolean | `false` | | Disable the generation of thinking level variants (e.g. -low, -medium, -high) in the /v1/models catalog. |
| `OMNIROUTE_CHAT_VIRTUAL_LANES` | boolean | `false` | ✓ | Enable per-tenant adaptive virtual admission lanes for provider dispatch (#9654): one tenant's burst no longer 503s another. The OMNIROUTE_CHAT_VIRTUAL_LANES env var wins over this dashboard override; changes take effect at server restart. |
| `EXPOSE_FUNCTIONAL_GATEWAY_MIRRORS` | boolean | `false` | | Advertise <gateway-alias>/<model> mirror ids on /v1/models for models whose canonical owner has no active credential but a passthrough gateway with an active credential routes them. Warning: adds catalog entries for all clients when enabled globally. |
| `NEWAPI_AGGREGATOR_BALANCE` | boolean | `false` | | Enable balance detection for New-API / One-API / Sub2API aggregator compatible nodes. When enabled, compatible nodes with the aggregator flag set will report their balance in the dashboard and quota-preflight routing. |
| `SERVER_OWNED_TOOL_LOOP_ENABLED` | boolean | `false` | | Continue non-streaming server-owned tool calls until the model returns a client-usable response. |
| `SEARCH_STATS_HIDE_DELETED_CONNECTIONS` | boolean | `false` | | Search stats and recent searches only count providers that still have a live connection (keyless providers such as duckduckgo-free always count). Off keeps every retained search row with a provider id. |
| `FREE_BADGE_REQUIRES_PROVIDER_FREE_TIER` | boolean | `false` | | Dashboard provider pages: show the Free badge only on signals the provider honors — drops the display-name heuristic, non-boolean free fields and :free suffixes on registered providers without a documented free tier. Off keeps the historical badge rule. |
| `RETRY_AFTER_PROVENANCE_ENABLED` | boolean | `false` | | On aggregated 429/503 unavailable responses, omit `Retry-After` when no concrete future retry time is known (instead of a synthetic 1s), add `error.retry_after_provenance` (`signal` \| `none`), and let combo drain paths read prose retry hints from JSON and plain-text upstream bodies. The field only appears on responses built by `unavailableResponse()`; other 429/503 bodies are unchanged. |
| `PROTECTED_PRIORITY_INFRA_502_ENABLED` | boolean | `false` | | When a `priority` combo target marked fallback-only-on-quota-exhaustion stops the combo for a cause that is provably not quota (provider circuit breaker open, predictive latency skip), answer 502 instead of the quota-looking 503. Lockout, cooldown, unavailable, exhaustion and concurrency-cap stops keep 503. |
| `MISTRAL_AMBIGUOUS_401_SOFT_LOCKOUT` | boolean | `false` | | A bare Mistral 401 (`{"detail":"Unauthorized"}`, no explicit auth signal) is identical for a revoked key and for exhausted quota. When on, it cools the connection down instead of parking it as `expired`, at most 3 times per hour per connection; the next one parks it, so a revoked key still converges. Off by default: every bare Mistral 401 parks the connection as before. |
| `XAI_OAUTH_LIVE_MODEL_DISCOVERY` | boolean | `false` | | Fetch the live xAI model catalog for `xai-oauth` connections from `https://api.x.ai/v1/models` using the OAuth bearer token, instead of the frozen static seed. Off by default: `xai-oauth` keeps serving the static seed unchanged. On any resolution error, discovery falls back to the seed (unverified whether x.ai accepts an OAuth bearer at this endpoint). |
| `BATCH_AND_FILE_AUTO_CLEANUP_ENABLED` | boolean | `false` | | Let the automatic cleanup sweep delete terminal (completed/failed/cancelled/expired) Batch API jobs older than `OMNIROUTE_BATCH_RETENTION_DAYS`, along with their per-line checkpoints, and clear the BLOB content of uploaded files past their own `expires_at`. Off by default: every existing install keeps this data exactly as before until an operator opts in. The operator-triggered `DELETE /api/v1/batches/delete-completed` route is unaffected either way — it is a separate, unconditional public API contract. |
| `ANTIGRAVITY_ACCOUNT_LEASE_ENABLED` | boolean | `false` | | Reserve the selected Antigravity account for the streaming lifecycle of the request that picked it, so a concurrent retry or the credential handoff cannot re-pick an account already committed to an in-flight stream. The reservation is scoped to (connection, callable upstream model), so one account can still serve two different models at once. When every eligible account is already leased for that model, the request returns a structured 503 `antigravity_pool_busy` with a bounded `Retry-After` instead of piling onto a busy account. Off by default: account selection stays exactly as before, and no reservation is taken. |
### CLI (5)

1
open-sse/utils/error.ts
View File

@@ -51,6 +51,7 @@ const SAFE_PUBLIC_ERROR_IDENTIFIERS = new Set([
"admission_unavailable",
"all_accounts_inactive",
"all_targets_skipped",
"antigravity_pool_busy",
"antigravity_pre_response_timeout",
"api_error",
"auth_error",

12
src/shared/constants/featureFlagDefinitions.ts
View File

@@ -738,6 +738,18 @@ export const FEATURE_FLAG_DEFINITIONS: FeatureFlagDefinition[] = [
requiresRestart: false,
warningLevel: "danger",
},
{
key: "ANTIGRAVITY_ACCOUNT_LEASE_ENABLED",
label: "Antigravity Account Lease",
description:
"Reserve the selected Antigravity account for the streaming lifecycle of the request that picked it, so a concurrent retry or the credential handoff cannot re-pick an account already committed to an in-flight stream. The reservation is scoped to (connection, callable upstream model), so one account can still serve two different models at once. When every eligible account is already leased for that model, the request returns a structured 503 POOL_BUSY with a bounded Retry-After instead of piling onto a busy account. Off by default: account selection stays exactly as before, and no reservation is taken.",
descriptionI18nKey: "featureFlagAntigravityAccountLeaseEnabledDescription",
category: "runtime",
defaultValue: "false",
type: "boolean",
requiresRestart: false,
warningLevel: "caution",
},
// ──────────────── CLI (5) ────────────────
{

20
src/shared/utils/featureFlags.ts
View File

@@ -306,6 +306,26 @@ export function isOpencodeRateLimited429EarlyStopEnabled(): boolean {
}
}
/**
* Antigravity account lease (re-land of #10011). Opt-in: when off, Antigravity
* account selection and the dispatch path behave exactly as before — no
* reservation is taken and no POOL_BUSY response can be produced.
* Fail closed: an unreadable flag store keeps the pre-flag behavior (disabled).
*/
export function isAntigravityAccountLeaseEnabled(
reader: (key: string) => boolean = isFeatureFlagEnabled
): boolean {
try {
return reader("ANTIGRAVITY_ACCOUNT_LEASE_ENABLED");
} catch (error) {
console.error(
"[featureFlags] Failed to resolve ANTIGRAVITY_ACCOUNT_LEASE_ENABLED, defaulting to disabled:",
error instanceof Error ? error.message : error
);
return false;
}
}
export function isServerOwnedToolLoopEnabled(
reader: (key: string) => boolean = isFeatureFlagEnabled
): boolean {

45
src/sse/handlers/chat.ts
View File

@@ -138,6 +138,7 @@ import { classify429FromError, type FailureKind } from "@/shared/utils/classify4
import { isSubscriptionQuotaText } from "@omniroute/open-sse/services/quotaTextCooldowns.ts";
import { resolveUseUpstream429BreakerHints } from "@/shared/utils/providerHints";
import { isFeatureFlagEnabled } from "@/shared/utils/featureFlags";
import * as agyLease from "../services/antigravityLeaseLifecycle";
import { shouldIsolateProbeFailures } from "@/shared/utils/probeOrigin";
import { getCircuitBreaker, isLocalStreamLifecycleError } from "../../shared/utils/circuitBreaker";
import { markAccountExhaustedFrom429 } from "../../domain/quotaCache";
@@ -1661,9 +1662,12 @@ async function handleSingleModelChat(
const occupancySessionKey =
runtimeOptions.sessionAffinityKey ?? runtimeOptions.sessionId ?? `request:${randomUUID()}`;
let initialPreselectedCredentials = runtimeOptions.preselectedCredentials;
// ANTIGRAVITY_ACCOUNT_LEASE_ENABLED (#10011 re-land): off ⇒ every `agy.*` branch is inert
// and selection/dispatch behave exactly as before. `attempted` survives a loop restart.
const agy = agyLease.startAntigravityLeaseRequest(provider, runtimeOptions.correlationId);
requestAttemptLoop: while (true) {
const excludedConnectionIds = new Set<string>();
const excludedConnectionIds = new Set<string>(agy.on ? agy.attempted : []);
let lastError = requestRetryLastError;
let lastStatus = requestRetryLastStatus;
let lastCooldownMs = requestRetryLastCooldownMs;
@@ -1672,7 +1676,7 @@ async function handleSingleModelChat(
while (true) {
const credentials =
preselectedCredentials && excludedConnectionIds.size === 0
preselectedCredentials && excludedConnectionIds.size === 0 && !agy.on
? preselectedCredentials
: await getProviderCredentialsWithQuotaPreflight(
provider,
@@ -1683,6 +1687,9 @@ async function handleSingleModelChat(
sessionKey: occupancySessionKey,
reserveOAuthSession: true,
excludeConnectionIds: Array.from(excludedConnectionIds),
...(agy.on
? { reserveAntigravityLease: true, routingRequestId: agy.requestId }
: {}),
...(runtimeOptions.allowRateLimitedConnection
? { allowRateLimitedConnections: true }
: {}),
@@ -1715,6 +1722,12 @@ async function handleSingleModelChat(
);
preselectedCredentials = null;
if (credentials && "leaseUnavailable" in credentials && credentials.leaseUnavailable) {
excludedConnectionIds.add(agyLease.trackAntigravityLeaseBusy(agy, credentials));
if (!hasForcedConnection) continue;
return agyLease.buildAntigravityPoolBusyResponse(agy.earliestRetryHintAtMs ?? Date.now());
}
if (runtimeOptions.managedLease && credentials) {
const leaseError = buildManagedLeaseSelectionErrorResponse(credentials);
if (leaseError) return leaseError;
@@ -1730,6 +1743,8 @@ async function handleSingleModelChat(
!credentials.connectionId
) {
if (earlyEofOriginal) return earlyEofOriginal;
if (!credentials?.allRateLimited && agy.earliestRetryHintAtMs !== null)
return agyLease.buildAntigravityPoolBusyResponse(agy.earliestRetryHintAtMs);
if (credentials?.allRateLimited) {
const retryDecision = getCooldownAwareRetryDecision({
retryAfter: credentials.retryAfter,
@@ -1820,6 +1835,9 @@ async function handleSingleModelChat(
const accountId = credentials.connectionId.slice(0, 8);
const releaseOAuthSession = credentials.releaseOAuthSession ?? (() => {});
// Undefined whenever the lease flag is off, which makes every release/hold below a no-op.
const leaseId: string | undefined = credentials.routing?.leaseId;
if (agy.on) agy.attempted.add(credentials.connectionId);
// #10348: redact the account prefix by default. Gated on the narrow
// AUTH_LOG_INCLUDE_ACCOUNT_ID flag (default off) rather than the broad
// `debugMode` setting — `debugMode` is a general dashboard-visibility
@@ -1863,9 +1881,10 @@ async function handleSingleModelChat(
reasoningIntent: runtimeOptions.reasoningIntent,
reasoningDecision: runtimeOptions.reasoningDecision,
requestRoutingTags: runtimeOptions.reasoningRequestTags,
});
}).catch(agyLease.releasingRethrow(leaseId));
if (connectionRouting.response) {
releaseOAuthSession();
agyLease.release(leaseId);
return connectionRouting.response;
}
requestBody = connectionRouting.body;
@@ -1894,7 +1913,9 @@ async function handleSingleModelChat(
}
let refreshedCredentials;
try {
refreshedCredentials = await checkAndRefreshToken(provider, credentials);
refreshedCredentials = await checkAndRefreshToken(provider, credentials).catch(
agyLease.releasingRethrow(leaseId)
);
} catch (error) {
releaseOAuthSession();
throw error;
@@ -1930,7 +1951,11 @@ async function handleSingleModelChat(
}
let proxyInfo;
try {
proxyInfo = await safeResolveProxy(credentials.connectionId, apiKeyInfo?.id, provider);
proxyInfo = await safeResolveProxy(
credentials.connectionId,
apiKeyInfo?.id,
provider
).catch(agyLease.releasingRethrow(leaseId));
} catch (error) {
releaseOAuthSession();
throw error;
@@ -1988,14 +2013,24 @@ async function handleSingleModelChat(
);
} catch (error) {
releaseOAuthSession();
agyLease.release(leaseId);
throw error;
}
if (telemetry) telemetry.endPhase();
if ("localResourcePressureResult" in execution) {
agyLease.release(leaseId);
return execution.localResourcePressureResult.response;
}
const { result, tlsFingerprintUsed } = execution;
if (!result.success) releaseOAuthSession();
// Hand the lease to the SSE body's terminal lifecycle; anything else frees it now.
if (result.success && agyLease.isStreamingAntigravityResponse(result.response))
result.response = agyLease.holdAntigravityLeaseThroughResponse(
result.response,
leaseId,
clientRawRequest?.signal
);
else agyLease.release(leaseId);
const proxyLatency = Date.now() - proxyStartTime;
const providerAlias = PROVIDER_ID_TO_ALIAS[provider] || provider;

160
src/sse/services/antigravityLeaseLifecycle.ts Normal file
View File

@@ -0,0 +1,160 @@
import { randomUUID } from "crypto";
import { buildErrorBody } from "@omniroute/open-sse/utils/error.ts";
import { isAntigravityAccountLeaseEnabled } from "@/shared/utils/featureFlags";
import { releaseAntigravityLease } from "./antigravityRoutingState";
import type { AntigravityLeaseUnavailable } from "./antigravityLeaseSelection";
/**
* Lifecycle glue for the Antigravity account lease (re-land of @Ardem2025's
* #10011). The lease is acquired during credential selection; these helpers own
* the two ways it ends: a pre-dispatch failure, or the terminal event of the
* SSE body it was reserved for.
*/
/**
* `.catch()` handler that frees a selector-owned lease when a pre-dispatch step throws
* and rethrows untouched, so the awaited call keeps its original shape.
*/
export function releasingRethrow(leaseId: string | null | undefined) {
return (error: unknown): never => {
releaseAntigravityLease(leaseId);
throw error;
};
}
export { releaseAntigravityLease as release };
/**
* Transfer the lease to the terminal lifecycle of an SSE body: it is released
* exactly once, on EOF, cancel, reader error, or client abort. Without a lease
* id (feature flag off) or a body, the response is returned untouched.
*/
export function holdAntigravityLeaseThroughResponse(
response: Response,
leaseId: string | null | undefined,
signal: AbortSignal | null | undefined
): Response {
if (!leaseId || !response.body) {
releaseAntigravityLease(leaseId);
return response;
}
const reader = response.body.getReader();
let released = false;
let abortListener: (() => void) | null = null;
const releaseOnce = () => {
if (released) return;
released = true;
if (abortListener && signal) signal.removeEventListener("abort", abortListener);
releaseAntigravityLease(leaseId);
};
const stream = new ReadableStream<Uint8Array>({
async pull(controller) {
try {
const { done, value } = await reader.read();
if (done) {
releaseOnce();
controller.close();
} else {
controller.enqueue(value);
}
} catch (error) {
releaseOnce();
controller.error(error);
}
},
async cancel(reason) {
try {
await reader.cancel(reason);
} finally {
releaseOnce();
}
},
});
abortListener = () => {
void reader
.cancel(signal?.reason)
.catch(() => {})
.finally(releaseOnce);
};
if (signal?.aborted) abortListener();
else signal?.addEventListener("abort", abortListener, { once: true });
return new Response(stream, {
status: response.status,
statusText: response.statusText,
headers: response.headers,
});
}
/** Only a streaming body has a lifecycle worth holding the lease through. */
export function isStreamingAntigravityResponse(response: unknown): response is Response {
return (
response instanceof Response &&
/(?:^|[,;\s])text\/event-stream(?:$|[;,\s])/i.test(response.headers.get("content-type") || "")
);
}
/**
* Every eligible Antigravity account is already leased for this model. Local lease
* saturation is neither provider quota/cooldown nor breaker failure, so it gets its own
* structured 503 instead of being reported as an upstream error.
*/
export function buildAntigravityPoolBusyResponse(retryHintAtMs: number): Response {
const retryAfter = Math.max(1, Math.ceil((retryHintAtMs - Date.now()) / 1000));
return new Response(
JSON.stringify(
buildErrorBody(
503,
"All eligible Antigravity accounts are currently handling a request for this model",
undefined,
{ type: "server_error", code: "antigravity_pool_busy" }
)
),
{
status: 503,
headers: { "Content-Type": "application/json", "Retry-After": String(retryAfter) },
}
);
}
/** Per-request lease state: whether the flag is on, and which accounts came back busy. */
export type AntigravityLeaseRequestState = {
on: boolean;
requestId: string;
attempted: Set<string>;
earliestRetryHintAtMs: number | null;
};
/**
* Resolve the opt-in flag once per request. Off ⇒ every helper below is inert, and the
* provider check short-circuits before the flag store is read at all.
*/
export function startAntigravityLeaseRequest(
provider: string,
correlationId: string | null | undefined
): AntigravityLeaseRequestState {
const on = provider === "antigravity" && isAntigravityAccountLeaseEnabled();
return {
on,
requestId: on ? (correlationId ?? randomUUID()) : "",
attempted: new Set<string>(),
earliestRetryHintAtMs: null,
};
}
/** Record a busy verdict and return the account id the caller must now exclude. */
export function trackAntigravityLeaseBusy(
state: AntigravityLeaseRequestState,
busy: AntigravityLeaseUnavailable
): string {
const hint = busy.retryHintAtMs;
if (
Number.isFinite(hint) &&
(state.earliestRetryHintAtMs === null || hint < state.earliestRetryHintAtMs)
) {
state.earliestRetryHintAtMs = hint;
}
state.attempted.add(busy.selectedConnectionId);
return busy.selectedConnectionId;
}

71
src/sse/services/antigravityLeaseSelection.ts Normal file
View File

@@ -0,0 +1,71 @@
import {
canonicalizeAntigravityExactModel,
releaseAntigravityLease,
tryAcquireAntigravityLease,
type AntigravityLease,
} from "./antigravityRoutingState";
export type { AntigravityLease };
/** Returned instead of credentials when the selected account is already leased. */
export type AntigravityLeaseUnavailable = {
leaseUnavailable: true;
selectedConnectionId: string;
retryHintAtMs: number;
};
type CredentialsWithRouting = { routing?: { leaseId?: string } };
/**
* Selection-side glue for the Antigravity account lease. Kept out of auth.ts so the
* selector only calls two functions and the lease stays a self-contained concern.
*/
export function reserveAntigravityLeaseForSelection(
provider: string,
connection: { id: string } | null | undefined,
requestedModel: string | null | undefined,
options: { reserveAntigravityLease?: boolean; routingRequestId?: string | null }
): { busy: AntigravityLeaseUnavailable } | { busy?: undefined; lease?: AntigravityLease } {
if (provider !== "antigravity" || !connection || options.reserveAntigravityLease !== true)
return {};
const acquired = tryAcquireAntigravityLease({
connectionId: connection.id,
requestedModel,
requestId: options.routingRequestId,
});
if (acquired.kind === "busy") {
return {
busy: {
leaseUnavailable: true,
selectedConnectionId: connection.id,
retryHintAtMs: acquired.retryHintAtMs,
},
};
}
return { lease: acquired.lease };
}
/** The `routing` descriptor carried on materialized credentials, or nothing. */
export function buildAntigravityRoutingFields(
lease: AntigravityLease | undefined,
connectionId: string,
requestedModel: string | null | undefined
) {
if (!lease) return {};
return {
routing: {
provider: "antigravity" as const,
connectionId,
exactModel: canonicalizeAntigravityExactModel(requestedModel),
leaseId: lease.id,
},
};
}
/**
* Hand a routing lease back. Every selector path that abandons the connection it just
* picked must call this — otherwise the account stays fenced for the rest of the process.
*/
export function releaseRoutingLeaseFromCredentials(credentials: unknown): void {
releaseAntigravityLease((credentials as CredentialsWithRouting)?.routing?.leaseId);
}

130
src/sse/services/antigravityRoutingState.ts Normal file
View File

@@ -0,0 +1,130 @@
import { randomUUID } from "crypto";
import { resolveAntigravityModelId } from "@omniroute/open-sse/config/antigravityModelAliases.ts";
import { resolveModelAlias } from "@/shared/constants/modelSpecs";
/**
* Process-local Antigravity account lease (re-land of @Ardem2025's #10011).
*
* The account selected for a request is reserved for the whole streaming
* lifecycle of that request, so a concurrent retry — or the credential handoff
* inside `getProviderCredentialsWithQuotaPreflight` — cannot re-pick an account
* that is already committed to an in-flight upstream stream.
*
* Scope is (connection, canonical upstream model), NOT the whole account: an
* Antigravity account can serve two different upstream models at once, and
* fencing the whole account would needlessly serialize unrelated traffic. Two
* catalog ids that resolve to the SAME callable upstream id (e.g. the
* `gemini-3.7-flash-{high,medium,low}` tiers, all `gemini-3.7-flash-tiered`)
* therefore share one lease.
*
* This is a concurrency reservation only. It is deliberately orthogonal to
* quota accounting (`selectAntigravityQuotaWindowNames` /
* `antigravityQuotaFamily.ts`) and never reads or writes quota state.
*/
export type AntigravityLease = {
id: string;
connectionId: string;
exactModel: string;
acquiredAtMs: number;
/**
* Bounded `Retry-After` hint for a POOL_BUSY response — NOT an expiry. A lease
* is released by its response lifecycle, never by elapsed time, so nothing in
* this module ever compares this against the clock to free a lease.
*/
retryHintAtMs: number;
};
export type LeaseAcquireResult =
{ kind: "acquired"; lease: AntigravityLease } | { kind: "busy"; retryHintAtMs: number };
export type LeaseAvailability = { available: true } | { available: false; retryHintAtMs: number };
/** Bounded POOL_BUSY `Retry-After` hint (ms). */
export const ANTIGRAVITY_LEASE_RETRY_HINT_MS = 30_000;
const leasesByKey = new Map<string, AntigravityLease>();
const keysByLeaseId = new Map<string, string>();
/**
* Collapse a requested model id onto the callable upstream id, so every catalog
* alias of one upstream model contends for the same lease.
*/
export function canonicalizeAntigravityExactModel(model: string | null | undefined): string {
const requested = typeof model === "string" ? model.trim().toLowerCase() : "";
if (!requested) return "";
const unprefixed = requested.replace(/^(?:antigravity|agy)\//, "");
// The Antigravity alias table is authoritative when it knows the requested id
// verbatim: `gemini-3.1-pro-high` is callable only as `gemini-pro-agent`, and the
// generic catalog alias would otherwise collapse it onto plain `gemini-3.1-pro`.
const upstream = resolveAntigravityModelId(unprefixed);
if (upstream !== unprefixed) return upstream;
const canonical = resolveModelAlias(unprefixed);
return resolveAntigravityModelId(canonical) || canonical;
}
function leaseKey(connectionId: string, exactModel: string): string {
return `antigravity::${connectionId}::${exactModel}`;
}
export function getAntigravityLeaseAvailability({
connectionId,
requestedModel,
}: {
connectionId: string;
requestedModel: string | null | undefined;
}): LeaseAvailability {
const held = leasesByKey.get(
leaseKey(connectionId, canonicalizeAntigravityExactModel(requestedModel))
);
return held ? { available: false, retryHintAtMs: held.retryHintAtMs } : { available: true };
}
export function tryAcquireAntigravityLease({
connectionId,
requestedModel,
requestId,
now = Date.now(),
}: {
connectionId: string;
requestedModel: string | null | undefined;
requestId?: string | null;
now?: number;
}): LeaseAcquireResult {
const exactModel = canonicalizeAntigravityExactModel(requestedModel);
const key = leaseKey(connectionId, exactModel);
const held = leasesByKey.get(key);
if (held) return { kind: "busy", retryHintAtMs: held.retryHintAtMs };
const lease: AntigravityLease = {
id: `${requestId || "request"}:${randomUUID()}`,
connectionId,
exactModel,
acquiredAtMs: now,
retryHintAtMs: now + ANTIGRAVITY_LEASE_RETRY_HINT_MS,
};
leasesByKey.set(key, lease);
keysByLeaseId.set(lease.id, key);
return { kind: "acquired", lease };
}
/**
* Release by lease id. The id fences a late `finally`: a stale id cannot free a
* newer lease that has since claimed the same (connection, model) slot.
*/
export function releaseAntigravityLease(leaseId: string | null | undefined): boolean {
if (!leaseId) return false;
const key = keysByLeaseId.get(leaseId);
if (key === undefined) return false;
keysByLeaseId.delete(leaseId);
if (leasesByKey.get(key)?.id === leaseId) {
leasesByKey.delete(key);
return true;
}
return false;
}
export function __resetAntigravityRoutingStateForTests(): void {
leasesByKey.clear();
keysByLeaseId.clear();
}

31
src/sse/services/auth.ts
View File

@@ -168,6 +168,12 @@ import { loadOptionalNoAuthApiKeyCredentials } from "./noAuthOptionalApiKey";
import { getResource404Bypass } from "./requestResourceHealth";
import { isVertexConnectionWidePermissionDenied } from "./vertexErrorClassifier";
import { maybeAutoDisableBannedAccount } from "./autoDisableBannedAccount";
import {
buildAntigravityRoutingFields,
releaseRoutingLeaseFromCredentials,
reserveAntigravityLeaseForSelection,
type AntigravityLease,
} from "./antigravityLeaseSelection";
import * as log from "../utils/logger";
import {
fisherYatesShuffle,
@@ -213,6 +219,9 @@ export interface CredentialSelectionOptions {
_leaseRetryWithLockHeld?: boolean;
/** Internal: freeze the original policy-valid candidate set across lease race/preflight retry. */
_leaseCandidateIds?: string[];
/** Antigravity account lease (#10011): only the final chat dispatch opts in. */
reserveAntigravityLease?: boolean;
routingRequestId?: string | null;
}
export type ExclusiveLeaseSelectionResult = {
exclusiveLease: ExclusiveConnectionLease;
@@ -1046,6 +1055,8 @@ async function materializeConnection(
extra: DeferredLeaseSelection & {
exclusiveLease?: ExclusiveConnectionLease;
reactivatedFromInactive?: boolean;
routingLease?: AntigravityLease;
requestedModel?: string | null;
} = {}
) {
const providerSpecificData = await hydrateConnectionProviderSpecificData(connection);
@@ -1081,6 +1092,7 @@ async function materializeConnection(
maxConcurrent: connection.maxConcurrent,
quotaWindowThresholds: connection.quotaWindowThresholds ?? null,
...(releaseOAuthSession ? { releaseOAuthSession } : {}),
...buildAntigravityRoutingFields(extra.routingLease, connection.id, extra.requestedModel),
...extra,
};
}
@@ -2187,6 +2199,14 @@ export async function getProviderCredentials(
}
}
const reserved = reserveAntigravityLeaseForSelection(
provider,
connection,
requestedModel,
options
);
if (reserved.busy) return reserved.busy;
if (provider === "antigravity" && connection) {
log.info(
"AUTH",
@@ -2197,6 +2217,8 @@ export async function getProviderCredentials(
return materializeConnection(connection, options, {
exclusiveLease,
...probeStamp,
routingLease: reserved.lease,
requestedModel,
});
} finally {
selectionLock?.release();
@@ -2306,15 +2328,20 @@ export async function getProviderCredentialsWithQuotaPreflight(
);
if (claim.kind === "LOST") {
selectedCredentials.releaseOAuthSession?.();
releaseRoutingLeaseFromCredentials(credentials);
excludedConnectionIds.add(connectionId);
pendingCredentialSelection =
await selectedCredentials.selectNextLeaseCandidate?.(connectionId);
return null;
}
if (claim.kind === "STALE") return { leaseFenceStale: true };
if (claim.kind === "STALE") {
releaseRoutingLeaseFromCredentials(credentials);
return { leaseFenceStale: true };
}
await selectedCredentials.commitSelectionSideEffects?.();
if (options.materializeCredentials === false) {
selectedCredentials.releaseOAuthSession?.();
releaseRoutingLeaseFromCredentials(credentials);
return { exclusiveLease: claim.lease, connectionId, provider };
}
return { ...credentials, exclusiveLease: claim.lease };
@@ -2413,6 +2440,7 @@ export async function getProviderCredentialsWithQuotaPreflight(
);
} catch (error) {
selectedCredentials.releaseOAuthSession?.();
releaseRoutingLeaseFromCredentials(credentials);
throw error;
}
if (preflight.proceed) {
@@ -2422,6 +2450,7 @@ export async function getProviderCredentialsWithQuotaPreflight(
}
selectedCredentials.releaseOAuthSession?.();
releaseRoutingLeaseFromCredentials(credentials);
const unavailableUntil = await markQuotaPreflightAccountUnavailable(
provider,

82
tests/unit/antigravity-account-lease-flag.test.ts Normal file
View File

@@ -0,0 +1,82 @@
import test from "node:test";
import assert from "node:assert/strict";
import fs from "node:fs";
import os from "node:os";
import path from "node:path";
process.env.DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-agy-lease-flag-"));
const { FEATURE_FLAG_DEFINITIONS } =
await import("../../src/shared/constants/featureFlagDefinitions.ts");
const { isAntigravityAccountLeaseEnabled } = await import("../../src/shared/utils/featureFlags.ts");
const chat = await import("../../src/sse/services/antigravityLeaseLifecycle.ts");
const { buildErrorBody } = await import("../../open-sse/utils/error.ts");
test("ANTIGRAVITY_ACCOUNT_LEASE_ENABLED is a runtime boolean flag that defaults to OFF", () => {
const def = FEATURE_FLAG_DEFINITIONS.find((d) => d.key === "ANTIGRAVITY_ACCOUNT_LEASE_ENABLED");
assert.ok(def, "ANTIGRAVITY_ACCOUNT_LEASE_ENABLED should exist");
assert.equal(def.category, "runtime");
assert.equal(def.type, "boolean");
assert.equal(def.defaultValue, "false");
assert.equal(def.requiresRestart, false);
assert.equal(def.descriptionI18nKey, "featureFlagAntigravityAccountLeaseEnabledDescription");
});
test("the flag reader is off by default and fails closed", () => {
assert.equal(
isAntigravityAccountLeaseEnabled(() => false),
false
);
assert.equal(
isAntigravityAccountLeaseEnabled(() => true),
true
);
const logs: unknown[] = [];
const original = console.error;
console.error = (...args: unknown[]) => logs.push(args);
try {
const result = isAntigravityAccountLeaseEnabled(() => {
throw new Error("flag store unavailable");
});
assert.equal(result, false);
assert.ok(
logs.some((args) =>
String(args).includes("Failed to resolve ANTIGRAVITY_ACCOUNT_LEASE_ENABLED")
),
"the failure should be logged with the flag key"
);
} finally {
console.error = original;
}
});
test("POOL_BUSY is a structured, sanitized 503 with a bounded positive Retry-After", async () => {
const response = chat.buildAntigravityPoolBusyResponse(Date.now() + 1_250);
assert.equal(response.status, 503);
const retryAfter = Number(response.headers.get("Retry-After"));
assert.ok(Number.isInteger(retryAfter) && retryAfter >= 1, `Retry-After was ${retryAfter}`);
const body = await response.json();
assert.equal(body.error.code, "antigravity_pool_busy");
assert.equal(body.error.type, "server_error");
assert.ok(!String(body.error.message).includes("at /"), "must not leak a stack trace");
// A hint already in the past must still produce a usable Retry-After, never 0 or negative.
assert.equal(
Number(chat.buildAntigravityPoolBusyResponse(Date.now() - 60_000).headers.get("Retry-After")),
1
);
});
test("antigravity_pool_busy survives the public error-identifier projection", () => {
// buildErrorBody projects any code outside SAFE_PUBLIC_ERROR_IDENTIFIERS onto the
// status fallback, which would silently turn POOL_BUSY into a generic
// service_unavailable and make the condition indistinguishable from an upstream 503.
const body = buildErrorBody(503, "pool busy", undefined, {
type: "server_error",
code: "antigravity_pool_busy",
});
assert.equal(body.error.code, "antigravity_pool_busy");
assert.notEqual(body.error.code, "service_unavailable");
});

170
tests/unit/antigravity-lease-lifecycle.test.ts Normal file
View File

@@ -0,0 +1,170 @@
import test from "node:test";
import assert from "node:assert/strict";
import fs from "node:fs";
import os from "node:os";
import path from "node:path";
process.env.DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-agy-lease-lifecycle-"));
const state = await import("../../src/sse/services/antigravityRoutingState.ts");
const lifecycle = await import("../../src/sse/services/antigravityLeaseLifecycle.ts");
const MODEL = "gemini-3.7-flash-high";
test.beforeEach(() => state.__resetAntigravityRoutingStateForTests());
function acquire() {
const result = state.tryAcquireAntigravityLease({
connectionId: "account-a",
requestedModel: MODEL,
});
assert.equal(result.kind, "acquired");
if (result.kind !== "acquired") throw new Error("unreachable");
return result.lease;
}
function held(): boolean {
return !state.getAntigravityLeaseAvailability({
connectionId: "account-a",
requestedModel: MODEL,
}).available;
}
function sseResponse(body: ReadableStream<Uint8Array>): Response {
return new Response(body, { headers: { "content-type": "text/event-stream" } });
}
test("only an event-stream response is treated as streaming", () => {
assert.equal(
lifecycle.isStreamingAntigravityResponse(
new Response("{}", { headers: { "content-type": "text/event-stream; charset=utf-8" } })
),
true
);
assert.equal(
lifecycle.isStreamingAntigravityResponse(
new Response("{}", { headers: { "content-type": "application/json" } })
),
false
);
assert.equal(lifecycle.isStreamingAntigravityResponse({ headers: {} }), false);
});
test("the lease is held for the whole stream and released on EOF", async () => {
const lease = acquire();
let controller!: ReadableStreamDefaultController<Uint8Array>;
const response = lifecycle.holdAntigravityLeaseThroughResponse(
sseResponse(
new ReadableStream({
start(c) {
controller = c;
},
})
),
lease.id,
null
);
const reader = response.body!.getReader();
controller.enqueue(new TextEncoder().encode("data: chunk\n\n"));
const first = await reader.read();
assert.equal(first.done, false);
// Mid-stream the account is still reserved.
assert.equal(held(), true);
controller.close();
await reader.read();
assert.equal(held(), false);
});
test("the lease is released when the stream is cancelled", async () => {
const lease = acquire();
const response = lifecycle.holdAntigravityLeaseThroughResponse(
sseResponse(new ReadableStream()),
lease.id,
null
);
assert.equal(held(), true);
await response.body!.cancel();
assert.equal(held(), false);
});
test("the lease is released when the client aborts", async () => {
const lease = acquire();
const abort = new AbortController();
const response = lifecycle.holdAntigravityLeaseThroughResponse(
sseResponse(new ReadableStream()),
lease.id,
abort.signal
);
assert.equal(held(), true);
abort.abort();
// The abort listener cancels the reader asynchronously.
await new Promise((resolve) => setTimeout(resolve, 10));
assert.equal(held(), false);
// Consume so the wrapper stream is not left dangling.
await response.body!.cancel().catch(() => {});
});
test("without a lease id the response is passed through untouched", () => {
const original = sseResponse(new ReadableStream());
assert.equal(lifecycle.holdAntigravityLeaseThroughResponse(original, undefined, null), original);
assert.equal(lifecycle.holdAntigravityLeaseThroughResponse(original, null, null), original);
});
test("a pre-dispatch failure frees the selector-owned lease and rethrows untouched", async () => {
const lease = acquire();
const failing = Promise.reject(new Error("reasoning rule failed")).catch(
lifecycle.releasingRethrow(lease.id)
);
await assert.rejects(failing, /reasoning rule failed/);
assert.equal(held(), false);
});
test("a successful pre-dispatch step keeps the lease and its resolved value", async () => {
const lease = acquire();
const value = await Promise.resolve(42).catch(lifecycle.releasingRethrow(lease.id));
assert.equal(value, 42);
assert.equal(held(), true);
});
test("release() is the fenced registry release re-exported for the dispatch path", () => {
const lease = acquire();
assert.equal(lifecycle.release("not-a-lease"), false);
assert.equal(held(), true);
assert.equal(lifecycle.release(lease.id), true);
assert.equal(held(), false);
});
test("startAntigravityLeaseRequest is inert for every non-Antigravity provider", () => {
const openai = lifecycle.startAntigravityLeaseRequest("openai", "req-1");
assert.equal(openai.on, false);
// Off ⇒ no correlation id is even carried, and the flag store is never read.
assert.equal(openai.requestId, "");
assert.equal(openai.earliestRetryHintAtMs, null);
assert.equal(openai.attempted.size, 0);
});
test("trackAntigravityLeaseBusy memoizes the account and the earliest retry hint", () => {
const state = lifecycle.startAntigravityLeaseRequest("antigravity", "req-2");
assert.equal(
lifecycle.trackAntigravityLeaseBusy(state, {
leaseUnavailable: true,
selectedConnectionId: "account-a",
retryHintAtMs: 5_000,
}),
"account-a"
);
lifecycle.trackAntigravityLeaseBusy(state, {
leaseUnavailable: true,
selectedConnectionId: "account-b",
retryHintAtMs: 3_000,
});
lifecycle.trackAntigravityLeaseBusy(state, {
leaseUnavailable: true,
selectedConnectionId: "account-c",
retryHintAtMs: 9_000,
});
assert.deepEqual([...state.attempted], ["account-a", "account-b", "account-c"]);
assert.equal(state.earliestRetryHintAtMs, 3_000);
});

162
tests/unit/antigravity-routing-state.test.ts Normal file
View File

@@ -0,0 +1,162 @@
import test from "node:test";
import assert from "node:assert/strict";
import fs from "node:fs";
import os from "node:os";
import path from "node:path";
// Any module reachable from src/lib/db opens a real SQLite file and runs migrations.
process.env.DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-agy-lease-state-"));
const state = await import("../../src/sse/services/antigravityRoutingState.ts");
test.beforeEach(() => state.__resetAntigravityRoutingStateForTests());
test("canonicalizes prefix, case and aliases onto the callable upstream id", () => {
// Every `gemini-3.7-flash-*` tier is callable only as the tiered endpoint.
assert.equal(
state.canonicalizeAntigravityExactModel("gemini-3.7-flash-high"),
"gemini-3.7-flash-tiered"
);
assert.equal(
state.canonicalizeAntigravityExactModel(" ANTIGRAVITY/Gemini-3.7-Flash-Medium "),
"gemini-3.7-flash-tiered"
);
assert.equal(
state.canonicalizeAntigravityExactModel("agy/gemini-3.7-flash-low"),
"gemini-3.7-flash-tiered"
);
// The Antigravity alias table wins over the generic catalog alias: the `-high`
// pro tier is callable as gemini-pro-agent, NOT as plain gemini-3.1-pro.
assert.equal(state.canonicalizeAntigravityExactModel("gemini-3.1-pro-high"), "gemini-pro-agent");
assert.equal(state.canonicalizeAntigravityExactModel("gemini-3.1-pro"), "gemini-3.1-pro");
assert.equal(state.canonicalizeAntigravityExactModel(null), "");
});
test("a lease fences the same connection and model, and only that pair", () => {
const first = state.tryAcquireAntigravityLease({
connectionId: "account-a",
requestedModel: "gemini-3.1-pro",
requestId: "request-a",
now: 1_000,
});
assert.equal(first.kind, "acquired");
if (first.kind !== "acquired") return;
assert.equal(first.lease.exactModel, "gemini-3.1-pro");
assert.equal(first.lease.connectionId, "account-a");
assert.ok(first.lease.id.startsWith("request-a:"));
assert.equal(first.lease.retryHintAtMs, 1_000 + state.ANTIGRAVITY_LEASE_RETRY_HINT_MS);
// Same account + same model → busy, carrying the held lease's retry hint.
assert.deepEqual(
state.tryAcquireAntigravityLease({
connectionId: "account-a",
requestedModel: "gemini-3.1-pro",
now: 1_001,
}),
{ kind: "busy", retryHintAtMs: first.lease.retryHintAtMs }
);
// A different account for the same model is free.
assert.equal(
state.tryAcquireAntigravityLease({
connectionId: "account-b",
requestedModel: "gemini-3.1-pro",
now: 1_001,
}).kind,
"acquired"
);
// A different upstream model on the SAME account is free: the lease scopes a
// (connection, model) pair, it does not serialize the whole account.
assert.equal(
state.tryAcquireAntigravityLease({
connectionId: "account-a",
requestedModel: "gemini-3.7-flash-high",
now: 1_001,
}).kind,
"acquired"
);
});
test("aliases of one upstream model contend for the same lease", () => {
const held = state.tryAcquireAntigravityLease({
connectionId: "account-a",
requestedModel: "gemini-3.7-flash-high",
requestId: "request-a",
now: 1_000,
});
assert.equal(held.kind, "acquired");
if (held.kind !== "acquired") return;
// `-medium` and `-high` are two catalog names for gemini-3.7-flash-tiered.
assert.deepEqual(
state.getAntigravityLeaseAvailability({
connectionId: "account-a",
requestedModel: "gemini-3.7-flash-medium",
}),
{ available: false, retryHintAtMs: held.lease.retryHintAtMs }
);
assert.deepEqual(
state.getAntigravityLeaseAvailability({
connectionId: "account-a",
requestedModel: "antigravity/gemini-3.7-flash-tiered",
}),
{ available: false, retryHintAtMs: held.lease.retryHintAtMs }
);
// A genuinely different upstream model is untouched.
assert.deepEqual(
state.getAntigravityLeaseAvailability({
connectionId: "account-a",
requestedModel: "gemini-3.1-pro",
}),
{ available: true }
);
});
test("release is fenced by lease id and restores availability", () => {
const acquired = state.tryAcquireAntigravityLease({
connectionId: "account-a",
requestedModel: "gemini-3.7-flash-high",
requestId: "request-a",
now: 1_000,
});
assert.equal(acquired.kind, "acquired");
if (acquired.kind !== "acquired") return;
// A foreign / stale id must not free somebody else's lease.
assert.equal(state.releaseAntigravityLease("wrong-lease"), false);
assert.equal(state.releaseAntigravityLease(null), false);
assert.equal(
state.getAntigravityLeaseAvailability({
connectionId: "account-a",
requestedModel: "gemini-3.7-flash-medium",
}).available,
false
);
assert.equal(state.releaseAntigravityLease(acquired.lease.id), true);
assert.deepEqual(
state.getAntigravityLeaseAvailability({
connectionId: "account-a",
requestedModel: "gemini-3.7-flash-medium",
}),
{ available: true }
);
// Double release is a no-op, and a late release cannot free the NEXT holder.
const second = state.tryAcquireAntigravityLease({
connectionId: "account-a",
requestedModel: "gemini-3.7-flash-medium",
requestId: "request-b",
now: 2_000,
});
assert.equal(second.kind, "acquired");
assert.equal(state.releaseAntigravityLease(acquired.lease.id), false);
assert.equal(
state.getAntigravityLeaseAvailability({
connectionId: "account-a",
requestedModel: "gemini-3.7-flash-high",
}).available,
false
);
});

2
tests/unit/feature-flags-settings.test.ts
View File

@@ -40,7 +40,7 @@ const {
// the dead ONEPROXY_ENABLED (readerless since the 1proxy purge, #12091)
// brought it back to 53. UNIVERSAL_CONTEXT_HANDOFF_ENABLED bumped it to 54.
// #13641 added SEARCH_STATS_HIDE_DELETED_CONNECTIONS, bumping the count to 56.
const EXPECTED_FEATURE_FLAG_COUNT = 72;
const EXPECTED_FEATURE_FLAG_COUNT = 73;
// ──────────────────────────────────────────────────────
// Test group 1 — Flag definitions registry

2
tests/unit/server-owned-tool-loop-flag.test.ts
View File

@@ -68,7 +68,7 @@ describe("isServerOwnedToolLoopEnabled wrapper", () => {
describe("feature-flags-settings count update", () => {
it("flag count matches updated expected value", () => {
assert.equal(FEATURE_FLAG_DEFINITIONS.length, 72);
assert.equal(FEATURE_FLAG_DEFINITIONS.length, 73);
});
});