From 8acd799af7d0ad6aa562f294a6c5c79acaebd61b Mon Sep 17 00:00:00 2001 From: KaspaPulse Date: Tue, 18 Aug 2026 17:25:46 +0300 Subject: [PATCH] feat(routing): add exclusive managed session connection leases (#10362) Co-authored-by: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com> --- ...leased-exclusive-managed-session-leases.md | 1 + docs/architecture/RESILIENCE_GUIDE.md | 34 + docs/openapi.yaml | 97 +++ docs/reference/API_REFERENCE.md | 59 ++ open-sse/handlers/chatCore.ts | 48 ++ .../chatCore/executorClientHeaders.ts | 6 + open-sse/utils/requestLogger.ts | 5 + .../api/internal/codex-responses-ws/route.ts | 16 + src/app/api/keys/[id]/route.ts | 8 + src/app/api/keys/route.ts | 4 +- src/app/api/providers/[id]/models/route.ts | 6 +- src/app/api/providers/[id]/test/route.ts | 12 + src/app/api/translator/send/route.ts | 10 +- src/app/api/translator/translate/route.ts | 18 +- src/app/api/v1/session-leases/route.ts | 139 ++++ src/lib/api/modelTestRunner.ts | 22 +- src/lib/credentialHealth/scheduler.ts | 4 +- src/lib/db/apiKeys.ts | 88 ++- src/lib/db/exclusiveConnectionLeases.ts | 394 +++++++++++ .../155_exclusive_connection_leases.sql | 30 + src/lib/exclusiveLeaseIsolation.ts | 8 + src/lib/localDb.ts | 3 + src/lib/services/quotaAutoPing.ts | 10 +- src/lib/usage/codexResetCredits.ts | 8 + src/lib/usage/providerLimits.ts | 18 +- src/lib/vncSession/service.ts | 7 + src/lib/warmupScheduler.ts | 2 + src/server/cors/origins.ts | 2 +- src/shared/services/modelSyncScheduler.ts | 6 + src/shared/utils/cors.ts | 2 +- src/shared/utils/shuffleDeck.ts | 27 + src/shared/validation/schemas/keys.ts | 37 +- src/sse/handlers/chat.ts | 121 +++- src/sse/handlers/chat/clientRawRequest.ts | 5 +- src/sse/handlers/chatHelpers.ts | 2 + src/sse/services/auth.ts | 413 ++++++++---- .../exclusiveConnectionLeasePolicy.ts | 125 ++++ src/sse/services/leaseContext.ts | 136 ++++ src/sse/services/sessionAffinityPin.ts | 37 + tests/unit/api-key-scope-validation.test.ts | 29 +- tests/unit/chat-managed-lease-routing.test.ts | 635 ++++++++++++++++++ .../chatcore-executor-client-headers.test.ts | 10 + tests/unit/chatcore-translation-paths.test.ts | 4 + .../codex-ws-policy-enforcement-6564.test.ts | 19 + tests/unit/cors/origins.test.ts | 6 + .../unit/exclusive-connection-leases.test.ts | 455 +++++++++++++ .../exclusive-lease-api-key-policy.test.ts | 105 +++ ...xclusive-lease-auxiliary-isolation.test.ts | 192 ++++++ ...ve-lease-connection-test-isolation.test.ts | 96 +++ .../unit/exclusive-lease-managed-set.test.ts | 99 +++ ...ard-session-lease-bypass-inventory.test.ts | 295 ++++++++ ...ard-session-lease-zero-model-gates.test.ts | 180 +++++ tests/unit/lease-context.test.ts | 86 +++ tests/unit/pick-internal-api-key-6372.test.ts | 9 + tests/unit/quota-auto-ping.test.ts | 81 ++- tests/unit/request-logger-endpoints.test.ts | 41 ++ tests/unit/session-leases-route.test.ts | 268 ++++++++ tests/unit/sse-auth-exclusive-leases.test.ts | 547 +++++++++++++++ tests/unit/warmupScheduler.test.ts | 36 + 59 files changed, 4981 insertions(+), 182 deletions(-) create mode 100644 changelog.d/features/unreleased-exclusive-managed-session-leases.md create mode 100644 src/app/api/v1/session-leases/route.ts create mode 100644 src/lib/db/exclusiveConnectionLeases.ts create mode 100644 src/lib/db/migrations/155_exclusive_connection_leases.sql create mode 100644 src/lib/exclusiveLeaseIsolation.ts create mode 100644 src/sse/services/exclusiveConnectionLeasePolicy.ts create mode 100644 src/sse/services/leaseContext.ts create mode 100644 tests/unit/chat-managed-lease-routing.test.ts create mode 100644 tests/unit/exclusive-connection-leases.test.ts create mode 100644 tests/unit/exclusive-lease-api-key-policy.test.ts create mode 100644 tests/unit/exclusive-lease-auxiliary-isolation.test.ts create mode 100644 tests/unit/exclusive-lease-connection-test-isolation.test.ts create mode 100644 tests/unit/exclusive-lease-managed-set.test.ts create mode 100644 tests/unit/hard-session-lease-bypass-inventory.test.ts create mode 100644 tests/unit/hard-session-lease-zero-model-gates.test.ts create mode 100644 tests/unit/lease-context.test.ts create mode 100644 tests/unit/session-leases-route.test.ts create mode 100644 tests/unit/sse-auth-exclusive-leases.test.ts diff --git a/changelog.d/features/unreleased-exclusive-managed-session-leases.md b/changelog.d/features/unreleased-exclusive-managed-session-leases.md new file mode 100644 index 0000000000..9db23724ef --- /dev/null +++ b/changelog.d/features/unreleased-exclusive-managed-session-leases.md @@ -0,0 +1 @@ +- **feat(routing):** add client-, provider-, and model-neutral exclusive managed session connection leases with API-key-bound generation fencing, durable SQLite ownership, explicit allowlist policy, and bounded 429 capacity retry semantics. diff --git a/docs/architecture/RESILIENCE_GUIDE.md b/docs/architecture/RESILIENCE_GUIDE.md index da4095e2fc..d92aebb512 100644 --- a/docs/architecture/RESILIENCE_GUIDE.md +++ b/docs/architecture/RESILIENCE_GUIDE.md @@ -107,6 +107,40 @@ Before #7274, `resolveSessionAffinityTtlMs()` hard-bailed to `0` for every provi The three session-affinity headers are never forwarded upstream — executors build their own upstream headers from scratch rather than passing client headers through, so this stays an internal correlation id only. +### Exclusive managed session connection leases + +**Scope:** one active managed HTTP client/session owns one eligible OmniRoute connection. + +**Purpose:** provide durable exclusive connection ownership for clients that need a hard routing +fence across requests. This differs from session affinity, which is a soft continuity preference: +an exclusive lease persists lifecycle state in SQLite, enforces global active-owner and +active-connection uniqueness, and rejects a stale generation before provider dispatch. + +The feature is opt-in per API key. A managed key must have the `lease:exclusive` scope and an +explicit non-empty `allowedConnections` list. Any HTTP client can use the lifecycle endpoint; no +client name, user-agent, provider, OAuth method, or model is required. The lease owns a connection, +not a model, so a model change retains the binding while the connection remains ordinarily +eligible. Normal model, quota, health, cooldown, and allowlist rules remain authoritative and may +transition the same generation to another free eligible connection. + +The lifecycle is `POST /api/v1/session-leases` with JSON actions `acquire`, `renew`, and `release`. +Managed inference requests present the opaque `X-OmniRoute-Lease-Owner` value and exact +`X-OmniRoute-Lease-Generation`. The owner uses `vlo_` followed by 43 base64url characters; only +its SHA-256 hash is stored. Every final dispatch fence also binds the authenticated API key ID and +active connection ID. Lease control headers are removed from logs, retained request snapshots, and +upstream executor headers. + +If ordinary routing has eligible managed candidates but every free candidate is occupied by a +foreign active lease, OmniRoute returns HTTP `429`, lease-capacity-unavailable code, a +waiting-for-capacity state, and a bounded `Retry-After` derived from the earliest relevant expiry. +Ordinary empty eligibility is not lease contention and keeps its existing routing error semantics. + +Related mechanisms remain separate: + +- OAuth session occupancy is process-local soft distribution for OAuth accounts. +- Account semaphores grant request-concurrency permits and end when a request completes. +- Exclusive managed session leases are durable lifecycle ownership with a generation fence. + --- ## 3. Model Lockout diff --git a/docs/openapi.yaml b/docs/openapi.yaml index fc5f85c60b..79cee7893b 100644 --- a/docs/openapi.yaml +++ b/docs/openapi.yaml @@ -56,6 +56,8 @@ tags: background scheduler tick. - name: API Keys description: API key management + - name: Session Leases + description: Client-neutral exclusive managed session connection leases - name: Combos description: Routing combo management - name: Settings @@ -103,6 +105,76 @@ tags: See docs/frameworks/TRAFFIC_INSPECTOR.md. paths: + /api/v1/session-leases: + post: + tags: + - Session Leases + summary: Acquire, renew, or release an exclusive managed connection lease + description: | + Requires an API key with `lease:exclusive` and an explicit non-empty + `allowedConnections` policy. The opaque owner is bound to the authenticated API key; + the lease owns an eligible connection, not a provider or model. Managed inference + requests present the owner and exact generation headers. Temporary foreign occupancy + returns 429 `WAITING_FOR_CAPACITY` with `Retry-After`. + security: + - BearerAuth: [] + parameters: + - name: X-OmniRoute-Lease-Owner + in: header + required: true + schema: + type: string + pattern: ^vlo_[A-Za-z0-9_-]{43}$ + requestBody: + required: true + content: + application/json: + schema: + oneOf: + - type: object + required: [action, model] + properties: + action: { type: string, const: acquire } + model: { type: string, minLength: 1, maxLength: 512 } + - type: object + required: [action, generation] + properties: + action: { type: string, const: renew } + generation: { type: integer, minimum: 1 } + - type: object + required: [action, generation] + properties: + action: { type: string, const: release } + generation: { type: integer, minimum: 1 } + reason: + type: string + enum: [OWNER_EXIT, CLIENT_CANCELLED] + responses: + "200": + description: Lease lifecycle state without connection or credential disclosure + content: + application/json: + schema: + $ref: "#/components/schemas/ExclusiveConnectionLeaseLifecycle" + "400": + description: Missing or invalid lease context/action + "401": + description: Missing or invalid API key + "403": + description: Managed lease scope or key configuration required + "409": + description: Stale generation, missing binding, or connection fence rejection + "415": + description: Lifecycle mutations require application/json + "429": + description: Eligible managed connections are held by foreign active leases + headers: + Retry-After: + schema: { type: integer, minimum: 1, maximum: 3600 } + content: + application/json: + schema: + $ref: "#/components/schemas/ExclusiveConnectionLeaseCapacity" # --- Playground + Search Tools (plans 17+18) --- /api/playground/improve-prompt: post: @@ -7346,6 +7418,31 @@ components: requestId: 0a1b2c3d-4e5f-6a7b-8c9d-0e1f2a3b4c5d schemas: + ExclusiveConnectionLeaseLifecycle: + type: object + required: [state, generation, acquiredAt, renewedAt, expiresAt] + properties: + state: { type: string, enum: [ACTIVE, RELEASED] } + generation: { type: integer, minimum: 1 } + acquiredAt: { type: string, format: date-time } + renewedAt: { type: string, format: date-time } + expiresAt: { type: string, format: date-time } + ExclusiveConnectionLeaseCapacity: + type: object + required: [state, error, reason, retryAfter, eligibleCount, freeCount] + properties: + state: { type: string, const: WAITING_FOR_CAPACITY } + error: + type: object + required: [type, code, message] + properties: + type: { type: string, const: lease_error } + code: { type: string, const: LEASE_CAPACITY_UNAVAILABLE } + message: { type: string } + reason: { type: string, const: NO_FREE_ELIGIBLE_CONNECTION } + retryAfter: { type: integer, minimum: 1, maximum: 3600 } + eligibleCount: { type: integer, minimum: 0 } + freeCount: { type: integer, minimum: 0 } EmbeddingMultimodalItem: oneOf: - type: object diff --git a/docs/reference/API_REFERENCE.md b/docs/reference/API_REFERENCE.md index b11b55a72f..23125894b2 100644 --- a/docs/reference/API_REFERENCE.md +++ b/docs/reference/API_REFERENCE.md @@ -15,6 +15,7 @@ Complete reference for all OmniRoute API endpoints. ## Table of Contents - [Chat Completions](#chat-completions) +- [Exclusive Managed Session Leases](#exclusive-managed-session-leases) - [Embeddings](#embeddings) - [Image Generation](#image-generation) - [Document OCR](#document-ocr) @@ -87,6 +88,64 @@ Content-Type: application/json > **Cache-hit cost semantics:** on a semantic-cache HIT (`X-OmniRoute-Cache-Hit: true`) no upstream call is made, so `X-OmniRoute-Response-Cost` is `0.0000000000` (the **incremental** cost of serving the hit). The original/would-have-been cost is reported separately in `X-OmniRoute-Cost-Saved`. Billing consumers should sum `X-OmniRoute-Response-Cost` (hits cost nothing); cache analytics can aggregate `X-OmniRoute-Cost-Saved`. +## Exclusive Managed Session Leases + +Exclusive managed session leasing is an opt-in, client-neutral routing contract: one active owner +holds one eligible OmniRoute connection. It does not lease a model, require OAuth, identify a +particular client, or require a particular provider. + +The authenticating API key must have scope `lease:exclusive` and an explicit non-empty +`allowedConnections` list. The database mutation boundary enforces both fields together on key +creation and partial updates. + +```http +POST /api/v1/session-leases +Authorization: Bearer +Content-Type: application/json +X-OmniRoute-Lease-Owner: vlo_<43-base64url-characters> + +{"action":"acquire","model":"glm/glm-4.6"} +``` + +Successful lifecycle responses expose timestamps, `state`, and the exact positive `generation`, +but never the selected connection or credentials. Renew and release supply the generation in the +JSON body: + +```json +{ "action": "renew", "generation": 1 } +``` + +```json +{ "action": "release", "generation": 1, "reason": "OWNER_EXIT" } +``` + +Every managed inference request then supplies both control headers: + +```http +X-OmniRoute-Lease-Owner: vlo_<43-base64url-characters> +X-OmniRoute-Lease-Generation: 1 +``` + +The exact owner, generation, active connection, and authenticated API key are fenced immediately +before each supported upstream attempt. Replaying owner and generation with another key fails even +when that key permits the same connection. Raw owners are not persisted, logged, retained in the +request snapshot, or forwarded upstream. + +Temporary contention returns HTTP `429` with `Retry-After` and: + +```json +{ + "state": "WAITING_FOR_CAPACITY", + "error": { "type": "lease_error", "code": "LEASE_CAPACITY_UNAVAILABLE" }, + "reason": "NO_FREE_ELIGIBLE_CONNECTION", + "retryAfter": 30 +} +``` + +This response only means that the ordinary eligible set was non-empty and every free candidate was +held by a foreign active lease. Unsupported models/providers, policy mismatch, cooldown, quota, +health, and other ordinary eligibility failures retain their existing OmniRoute responses. + ### `x-omniroute-compression` Per-request override of the compression plan. Highest precedence — beats the routing-combo diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index a8d60759ee..74c9c9cc4d 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -310,6 +310,7 @@ import { } from "./chatCore/upstreamTimeouts.ts"; import { getModelNormalizeToolCallId, getModelPreserveOpenAIDeveloperRole } from "@/lib/db/models"; import { getProviderCredentials, extractSessionAffinityKey } from "@/sse/services/auth"; +import { assertExclusiveConnectionLeaseFence } from "@/lib/db/exclusiveConnectionLeases"; import { deleteSessionAccountAffinity } from "@/lib/db/sessionAccountAffinity"; import { getCacheControlSettings } from "@/lib/cacheControlSettings"; import { guardrailRegistry } from "@/lib/guardrails"; @@ -467,6 +468,7 @@ export async function handleChatCore({ correlationId = null, modelPinned = false, skipResourcePressureGuard = false, + managedLease = null, }) { let { provider, model, extendedContext } = modelInfo; if (!skipResourcePressureGuard) { @@ -512,6 +514,46 @@ export async function handleChatCore({ : null; return credentialConnectionId || connectionId || null; }; + const assertManagedLeaseFence = (attemptConnectionId: string | null | undefined) => { + if (!managedLease) return; + if (!attemptConnectionId) { + throw Object.assign(new Error("Managed lease connection is unavailable"), { + code: "LEASE_CONNECTION_MISMATCH", + status: 409, + }); + } + const fence = assertExclusiveConnectionLeaseFence({ + leaseOwnerId: managedLease.context.leaseOwnerId, + generation: managedLease.context.generation, + apiKeyId: managedLease.apiKeyId, + connectionId: attemptConnectionId, + }); + if (fence.kind === "VALID") return; + const code = + fence.kind === "REQUIRED" + ? "LEASE_REQUIRED" + : fence.kind === "STALE" + ? "LEASE_FENCE_STALE" + : fence.kind === "AUTHORIZATION_MISMATCH" + ? "LEASE_AUTHORIZATION_MISMATCH" + : "LEASE_CONNECTION_MISMATCH"; + throw Object.assign(new Error("Managed lease request fence rejected the dispatch"), { + code, + status: 409, + }); + }; + const isManagedLeaseFenceError = (error: unknown): boolean => + managedLease !== null && + typeof (error as { code?: unknown })?.code === "string" && + String((error as { code: string }).code).startsWith("LEASE_"); + const managedLeaseFenceErrorResult = (error: unknown) => { + const code = (error as { code: string }).code; + return { + ...createErrorResult(409, "Managed lease request fence rejected the dispatch", null, code), + errorType: "lease_error", + errorCode: code, + }; + }; let tokensCompressed: number | null = null; body = injectSystemPrompt(body); // ── Per-endpoint custom system prompt (port of upstream #2063) ── @@ -2970,6 +3012,7 @@ export async function handleChatCore({ updatePendingScope(pendingScope, { stage: "rate_limit_slot_acquired", }); + assertManagedLeaseFence(attemptConnectionId); return executeWithUpstreamStartTimeout({ executor, provider, @@ -3040,6 +3083,7 @@ export async function handleChatCore({ // Codex 429 account-rotation failover (disabled for context-relay so combo.ts can inject handoff) if ( provider === "codex" && + !managedLease && comboStrategy !== "context-relay" && res.response.status === 429 && attempts < maxAttempts - 1 @@ -3202,6 +3246,7 @@ export async function handleChatCore({ body: unknown ): Promise | null> => { try { + assertManagedLeaseFence(attemptConnectionId); const retryRaw = await executeWithUpstreamStartTimeout({ executor, provider, @@ -3516,6 +3561,7 @@ export async function handleChatCore({ } } catch (error) { trackPendingRequest(model, provider, connectionId, false); + if (isManagedLeaseFenceError(error)) return managedLeaseFenceErrorResult(error); if (isSemaphoreCapacityError(error)) { appendRequestLog({ model, @@ -3727,6 +3773,7 @@ export async function handleChatCore({ // stay aligned if this block ever runs after a path that mutates body.model (e.g. fallback). try { const retryModelId = String(translatedBody.model || effectiveModel); + assertManagedLeaseFence(getExecutionConnectionId(getExecutionCredentials())); const retryResult = normalizeExecutorResult( await runWithCapture(providerRequestCapture, () => executor.execute({ @@ -3764,6 +3811,7 @@ export async function handleChatCore({ upstreamErrorParsed = false; // Let it be parsed downstream } } catch (retryErr) { + if (isManagedLeaseFenceError(retryErr)) return managedLeaseFenceErrorResult(retryErr); // Refresh succeeded but the retry leg failed (network blip, AbortError, // executor throw). Don't swallow — the operator-visible signal "the user // saw 401 even though auth was actually fixed" is much more confusing diff --git a/open-sse/handlers/chatCore/executorClientHeaders.ts b/open-sse/handlers/chatCore/executorClientHeaders.ts index e2a77bf36e..2088bcbd99 100644 --- a/open-sse/handlers/chatCore/executorClientHeaders.ts +++ b/open-sse/handlers/chatCore/executorClientHeaders.ts @@ -13,13 +13,19 @@ export function buildExecutorClientHeaders( userAgent?: string | null ) { const normalized: Record = {}; + const isLeaseControlHeader = (key: string) => { + const lowerKey = key.toLowerCase(); + return lowerKey === "x-omniroute-lease-owner" || lowerKey === "x-omniroute-lease-generation"; + }; if (headers instanceof Headers) { headers.forEach((value, key) => { + if (isLeaseControlHeader(key)) return; normalized[key] = value; }); } else if (headers && typeof headers === "object") { for (const [key, value] of Object.entries(headers)) { + if (isLeaseControlHeader(key)) continue; if (typeof value === "string") { normalized[key] = value; } diff --git a/open-sse/utils/requestLogger.ts b/open-sse/utils/requestLogger.ts index ea2cc6ff0c..eb7709a835 100644 --- a/open-sse/utils/requestLogger.ts +++ b/open-sse/utils/requestLogger.ts @@ -81,6 +81,7 @@ function maskSensitiveHeaders(headers: HeaderInput): Record { "storage-state", "storagestate", "capability", + "x-omniroute-lease-owner", ]; for (const key of Object.keys(masked)) { @@ -89,6 +90,10 @@ function maskSensitiveHeaders(headers: HeaderInput): Record { if (lowerKey.startsWith("x-ratelimit-")) { continue; } + if (lowerKey === "x-omniroute-lease-owner") { + masked[key] = "[REDACTED]"; + continue; + } if (!sensitiveKeys.some((candidate) => lowerKey.includes(candidate))) { continue; } diff --git a/src/app/api/internal/codex-responses-ws/route.ts b/src/app/api/internal/codex-responses-ws/route.ts index fd7f48dd35..cbb8d652d9 100644 --- a/src/app/api/internal/codex-responses-ws/route.ts +++ b/src/app/api/internal/codex-responses-ws/route.ts @@ -37,6 +37,11 @@ import { persistResponsesWsCallHistory } from "./history"; import { applyResponsesWsCompression } from "./compression"; import { getComboByName } from "@/lib/db/combos"; import { getComboModelString } from "@/lib/combos/steps"; +import { + buildManagedLeaseErrorResponse, + isExclusiveLeaseManagedKey, + LeaseContextError, +} from "@/sse/services/leaseContext"; const CODEX_RESPONSES_WS_URL = "wss://chatgpt.com/backend-api/codex/responses"; const executor = new CodexExecutor(); @@ -418,6 +423,17 @@ async function resolveCodexRequestContext(body: JsonRecord) { if (policyResult.rejection) return { error: policyResult.rejection }; const metadata = policyResult.apiKeyInfo ?? (apiKey ? await getApiKeyMetadata(apiKey).catch(() => null) : null); + if (isExclusiveLeaseManagedKey(metadata)) { + return { + error: buildManagedLeaseErrorResponse( + new LeaseContextError( + 409, + "LEASE_UNSUPPORTED_TRANSPORT", + "Managed leases require the fenced HTTP Responses transport" + ) + ), + }; + } const allowedConnections = metadata && Array.isArray(metadata.allowedConnections) && metadata.allowedConnections.length > 0 ? metadata.allowedConnections diff --git a/src/app/api/keys/[id]/route.ts b/src/app/api/keys/[id]/route.ts index 541a294f04..f3df18a230 100644 --- a/src/app/api/keys/[id]/route.ts +++ b/src/app/api/keys/[id]/route.ts @@ -4,6 +4,7 @@ import { getApiKeyById, updateApiKeyPermissions, isCloudEnabled, + ApiKeyPolicyInvariantError, } from "@/lib/localDb"; import { getConsistentMachineId } from "@/shared/utils/machineId"; import { syncToCloud } from "@/lib/cloudSync"; @@ -11,6 +12,7 @@ import { updateKeyPermissionsSchema } from "@/shared/validation/schemas"; import { isValidationFailure, validateBody } from "@/shared/validation/helpers"; import { requireManagementAuth } from "@/lib/api/requireManagementAuth"; import * as log from "@/sse/utils/logger"; +import { buildErrorBody } from "@omniroute/open-sse/utils/error.ts"; // GET /api/keys/[id] - Get single API key export async function GET(request, { params }) { @@ -159,6 +161,12 @@ export async function PATCH(request, { params }) { ...(chaosModeEnabled !== undefined && { chaosModeEnabled }), }); } catch (error) { + if (error instanceof ApiKeyPolicyInvariantError) { + return NextResponse.json(buildErrorBody(400, error.message, null, { + type: "lease_error", + code: error.code, + }), { status: 400 }); + } log.error("keys", "Error updating key permissions", error); return NextResponse.json({ error: "Failed to update permissions" }, { status: 500 }); } diff --git a/src/app/api/keys/route.ts b/src/app/api/keys/route.ts index 6e0f34ea11..d50b4c9d17 100644 --- a/src/app/api/keys/route.ts +++ b/src/app/api/keys/route.ts @@ -73,6 +73,7 @@ export async function POST(request) { name, noLog, scopes, + allowedConnections, allowUsageCommand, usageLimitEnabled, dailyUsageLimitUsd, @@ -83,7 +84,7 @@ export async function POST(request) { // Always get machineId from server const machineId = await getConsistentMachineId(); const normalizedScopes = normalizeSelfServiceScopesForCreate(scopes); - const apiKey = await createApiKey(name, machineId, normalizedScopes); + const apiKey = await createApiKey(name, machineId, normalizedScopes, { allowedConnections }); if ( noLog === true || allowUsageCommand === true || @@ -118,6 +119,7 @@ export async function POST(request) { name: apiKey.name, id: apiKey.id, machineId: apiKey.machineId, + allowedConnections: apiKey.allowedConnections, noLog: noLog === true, allowUsageCommand: allowUsageCommand === true, usageLimitEnabled: usageLimitEnabled === true, diff --git a/src/app/api/providers/[id]/models/route.ts b/src/app/api/providers/[id]/models/route.ts index 7385f17c9b..764a6c1d70 100755 --- a/src/app/api/providers/[id]/models/route.ts +++ b/src/app/api/providers/[id]/models/route.ts @@ -26,7 +26,7 @@ import { getProviderOutboundGuard, getProviderValidationGuard, } from "@/shared/network/outboundUrlGuardPolicy"; -import { sanitizeErrorMessage } from "@omniroute/open-sse/utils/error"; +import { errorResponse, sanitizeErrorMessage } from "@omniroute/open-sse/utils/error"; import { getStaticQoderModels } from "@omniroute/open-sse/services/qoderCli.ts"; import { deriveConfigFromRegistryModelsUrl } from "./discoveryConfig"; import { resolveZedModels } from "@omniroute/open-sse/shared/zedAuth.ts"; @@ -90,6 +90,7 @@ import { type GeminiDiscoveryModel, } from "@/lib/providerModels/geminiModelsParser"; import { getSyncedAvailableModels, getCustomModels } from "@/lib/db/models"; +import { isConnectionUnavailableToAuxiliaryActivity } from "@/lib/exclusiveLeaseIsolation"; import { fetchCursorAgentModels } from "@/lib/providerModels/cursorAgent"; import { ensureCursorAutoCatalogEntry } from "@/lib/providerModels/cursorAutoCatalog"; import { fetchRaycastModels } from "@omniroute/open-sse/services/raycast.ts"; @@ -179,6 +180,9 @@ export async function GET( return NextResponse.json({ error: "Connection not found" }, { status: 404 }); } + if (await isConnectionUnavailableToAuxiliaryActivity(id)) + return errorResponse(409, "Model discovery deferred for managed connection"); + // #6148 — short-circuit when a stored credential is encrypted but no longer // decrypts (STORAGE_ENCRYPTION_KEY changed/unset). Otherwise the null key is // coerced to "", an empty-Bearer probe is sent, and the operator sees a diff --git a/src/app/api/providers/[id]/test/route.ts b/src/app/api/providers/[id]/test/route.ts index 447a7edfe7..f8afbb6cb0 100644 --- a/src/app/api/providers/[id]/test/route.ts +++ b/src/app/api/providers/[id]/test/route.ts @@ -29,6 +29,7 @@ import { import { providerAllowsOptionalApiKey } from "@/shared/constants/providers"; import { shouldUseApiKeyConnectionTest } from "./webSessionTestDispatch"; import { removeConnectionHealth } from "@omniroute/open-sse/services/apiKeyRotator.ts"; +import { isConnectionUnavailableToAuxiliaryActivity } from "@/lib/exclusiveLeaseIsolation"; import { classifyAmbiguousOrAuthError, type ClassifyFailureArgs } from "./mistralAmbiguousAuth"; import { buildApiKeyConnectionTestResult } from "./apiKeyTestResult"; import { OAUTH_TEST_CONFIG } from "./oauthTestConfig"; @@ -728,6 +729,17 @@ export async function testSingleConnection(connectionId: string, validationModel return { valid: false, error: "Connection not found", diagnosis: null, latencyMs: 0 }; } + if (await isConnectionUnavailableToAuxiliaryActivity(connectionId)) { + const error = "Connection test deferred while an exclusive session lease is active"; + return { + valid: false, + skipped: true, + error, + diagnosis: makeDiagnosis("lease_active", "local", error, "exclusive_lease_active"), + latencyMs: 0, + }; + } + const provider = typeof connection.provider === "string" ? connection.provider : ""; if (!provider) { return { diff --git a/src/app/api/translator/send/route.ts b/src/app/api/translator/send/route.ts index db234e120d..3678509fd5 100644 --- a/src/app/api/translator/send/route.ts +++ b/src/app/api/translator/send/route.ts @@ -6,6 +6,7 @@ import { getTargetFormat, } from "@omniroute/open-sse/services/provider.ts"; import { getProviderConnections } from "@/lib/localDb"; +import { isConnectionUnavailableToAuxiliaryActivity } from "@/lib/exclusiveLeaseIsolation"; import { toJsonErrorPayload } from "@/shared/utils/upstreamError"; import { logTranslationEvent } from "@/lib/translatorEvents"; import { translatorSendSchema } from "@/shared/validation/schemas"; @@ -47,7 +48,14 @@ export async function POST(request) { // Get provider credentials from database const connections = await getProviderConnections({ provider }); - const connection = connections.find((c) => c.isActive !== false); + const connection = ( + await Promise.all( + connections.map(async (candidate) => ({ + candidate, + blocked: await isConnectionUnavailableToAuxiliaryActivity(candidate.id), + })) + ) + ).find(({ candidate, blocked }) => candidate.isActive !== false && !blocked)?.candidate; if (!connection) { logTranslationEvent({ diff --git a/src/app/api/translator/translate/route.ts b/src/app/api/translator/translate/route.ts index d30c62724c..8c14ea33a4 100644 --- a/src/app/api/translator/translate/route.ts +++ b/src/app/api/translator/translate/route.ts @@ -8,6 +8,7 @@ import { import { translateRequest } from "@omniroute/open-sse/translator/index.ts"; import { FORMATS } from "@omniroute/open-sse/translator/formats.ts"; import { getProviderConnections } from "@/lib/localDb"; +import { isConnectionUnavailableToAuxiliaryActivity } from "@/lib/exclusiveLeaseIsolation"; import { translatorTranslateSchema } from "@/shared/validation/schemas"; import { isValidationFailure, validateBody } from "@/shared/validation/helpers"; @@ -38,11 +39,21 @@ function getProviderBaseUrl(providerSpecificData: unknown): string | undefined { async function getActiveProviderSpecificData(provider?: string | null): Promise { if (!provider) return null; - const connections = await getProviderConnections({ provider }); - const connection = connections.find((c) => c.isActive !== false); + const connection = await getUnmanagedActiveConnection(provider); return connection ? asJsonRecord(connection.providerSpecificData) : null; } +async function getUnmanagedActiveConnection(provider: string) { + const connections = await getProviderConnections({ provider }); + for (const connection of connections) { + if ( + connection.isActive !== false && + !(await isConnectionUnavailableToAuxiliaryActivity(connection.id)) + ) + return connection; + } +} + export async function POST(request) { let rawBody; try { @@ -165,8 +176,7 @@ export async function POST(request) { const model = getModelId(actualBody); // Get provider credentials - const connections = await getProviderConnections({ provider }); - const connection = connections.find((c) => c.isActive !== false); + const connection = await getUnmanagedActiveConnection(provider); if (!connection) { return NextResponse.json( diff --git a/src/app/api/v1/session-leases/route.ts b/src/app/api/v1/session-leases/route.ts new file mode 100644 index 0000000000..53671f5ad3 --- /dev/null +++ b/src/app/api/v1/session-leases/route.ts @@ -0,0 +1,139 @@ +import { z } from "zod"; + +import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy"; +import { CORS_HEADERS, handleCorsOptions } from "@/shared/utils/cors"; +import { + releaseExclusiveConnectionLease, + renewExclusiveConnectionLease, +} from "@/lib/db/exclusiveConnectionLeases"; +import { + extractApiKey, + getProviderCredentialsWithQuotaPreflight, + isValidApiKey, +} from "@/sse/services/auth"; +import type { ExclusiveLeaseSelectionResult } from "@/sse/services/auth"; +import { + buildManagedLeaseSelectionErrorResponse, + isExclusiveLeaseManagedKey, + LeaseContextError, + parseLeaseOwnerHeader, + validateExclusiveLeaseKeyConfiguration, +} from "@/sse/services/leaseContext"; +import { getModelInfo } from "@/sse/services/model"; +import { buildErrorBody } from "@omniroute/open-sse/utils/error.ts"; + +const action = (name: T, shape: z.ZodRawShape) => + z.object({ action: z.literal(name), ...shape }); +const generation = z.number().int().positive().safe(); +const actionSchema = z.discriminatedUnion("action", [ + action("acquire", { model: z.string().trim().min(1).max(512) }), + action("renew", { generation }), + z.object({ + action: z.literal("release"), + generation, + reason: z.enum(["OWNER_EXIT", "CLIENT_CANCELLED"]).optional(), + }), +]); + +const json = (status: number, body: unknown) => + new Response(JSON.stringify(body), { + status, + headers: { ...CORS_HEADERS, "Content-Type": "application/json" }, + }); + +function error(status: number, code: string, message: string): Response { + return json(status, buildErrorBody(status, message, null, { type: "lease_error", code })); +} + +const lifecycle = (lease: Record) => { + const { state, generation, acquiredAt, renewedAt, expiresAt } = lease; + return { state, generation, acquiredAt, renewedAt, expiresAt }; +}; + +export const OPTIONS = async (): Promise => handleCorsOptions(); + +export async function POST(request: Request): Promise { + const apiKey = extractApiKey(request); + if (!apiKey) return error(401, "LEASE_AUTHENTICATION_REQUIRED", "Authentication required"); + if (!(await isValidApiKey(apiKey))) + return error(401, "LEASE_API_KEY_INVALID", "Invalid API key"); + const contentType = request.headers.get("content-type")?.toLowerCase().split(";", 1)[0].trim(); + if (contentType !== "application/json") { + return error(415, "LEASE_CONTENT_TYPE_REQUIRED", "Content-Type must be application/json"); + } + + const parsed = actionSchema.safeParse(await request.json().catch(() => null)); + if (!parsed.success) return error(400, "LEASE_ACTION_INVALID", "Invalid lease lifecycle action"); + + const policy = await enforceApiKeyPolicy( + request, + parsed.data.action === "acquire" ? parsed.data.model : null + ); + if (policy.rejection) return policy.rejection; + if (!policy.apiKeyInfo || !isExclusiveLeaseManagedKey(policy.apiKeyInfo)) + return error(403, "LEASE_SCOPE_REQUIRED", "The lease:exclusive scope is required"); + + try { + validateExclusiveLeaseKeyConfiguration(policy.apiKeyInfo); + const leaseOwnerId = parseLeaseOwnerHeader(request.headers); + if (parsed.data.action !== "acquire") { + const input = { + leaseOwnerId, + generation: parsed.data.generation, + apiKeyId: policy.apiKeyInfo.id, + }; + const result = + parsed.data.action === "renew" + ? renewExclusiveConnectionLease(input) + : releaseExclusiveConnectionLease({ ...input, reason: parsed.data.reason }); + return result.kind !== "STALE" + ? json(200, lifecycle(result.lease)) + : error(409, "LEASE_FENCE_STALE", "The lease generation is stale"); + } + + const modelInfo = await getModelInfo(parsed.data.model); + if (!modelInfo.provider) return error(400, "LEASE_MODEL_INVALID", "The model is unavailable"); + const selection = await getProviderCredentialsWithQuotaPreflight( + modelInfo.provider, + null, + policy.apiKeyInfo.allowedConnections ?? [], + modelInfo.model || parsed.data.model, + { + lease: { + apiKeyId: policy.apiKeyInfo.id, + context: { leaseOwnerId, generation: 1 }, + mode: "acquire", + }, + materializeCredentials: false, + reserveOAuthSession: false, + } + ); + if (!selection) { + return error( + 409, + "LEASE_NO_ELIGIBLE_CONNECTION", + "No eligible connection satisfies the managed key policy" + ); + } + const failure = buildManagedLeaseSelectionErrorResponse(selection); + if (failure) { + for (const [name, value] of Object.entries(CORS_HEADERS)) failure.headers.set(name, value); + return failure; + } + if ( + ("allRateLimited" in selection && selection.allRateLimited) || + ("allExpired" in selection && selection.allExpired) + ) { + return error( + 429, + "LEASE_ELIGIBILITY_UNAVAILABLE", + "Eligible connections are unavailable under current routing policy" + ); + } + const result = selection as ExclusiveLeaseSelectionResult; + return json(200, lifecycle(result.exclusiveLease)); + } catch (cause) { + if (cause instanceof LeaseContextError) return error(cause.status, cause.code, cause.message); + return error(503, "LEASE_SERVICE_UNAVAILABLE", "Lease service unavailable"); + } +} diff --git a/src/lib/api/modelTestRunner.ts b/src/lib/api/modelTestRunner.ts index 6738e26feb..f68454f6f5 100644 --- a/src/lib/api/modelTestRunner.ts +++ b/src/lib/api/modelTestRunner.ts @@ -18,6 +18,7 @@ import { } from "@omniroute/open-sse/services/accountFallback"; import { looksLikeQuotaExhausted } from "@/shared/utils/classify429"; import { getTrustedLocalRateLimitError } from "@omniroute/open-sse/services/rateLimitManager/errors"; +import { isConnectionUnavailableToAuxiliaryActivity } from "@/lib/exclusiveLeaseIsolation"; const INTERNAL_ORIGIN = "http://omniroute.internal"; export const DEFAULT_MODEL_TEST_TIMEOUT_MS = 30_000; @@ -356,9 +357,10 @@ function isBotBlockMessage(message: string): boolean { * Reuses the routing path's existing quota vocabulary from accountFallback.ts * and classify429.ts instead of inventing a new vocabulary. */ -export function classifyTestErrorQuota( - errorText: string -): { isQuota?: boolean; isTransient?: boolean } { +export function classifyTestErrorQuota(errorText: string): { + isQuota?: boolean; + isTransient?: boolean; +} { const trimmed = typeof errorText === "string" ? errorText.trim() : ""; if (!trimmed) return {}; @@ -399,6 +401,17 @@ export async function runSingleModelTest( streamChat = true, } = options; + if (connectionId && (await isConnectionUnavailableToAuxiliaryActivity(connectionId))) { + const fullModelId = modelId.includes("/") ? modelId : `${providerId}/${modelId}`; + return { + modelId: fullModelId, + status: "error", + latencyMs: 0, + httpStatus: 409, + error: "Model tests are unavailable for managed lease connections", + }; + } + let fullModelStr = modelId; if (!fullModelStr.includes("/")) { fullModelStr = `${providerId}/${modelId}`; @@ -572,7 +585,8 @@ export async function runSingleModelTest( // error, not a bot-block. A bare 403 status without quota/bot wording still // falls through to the generic error branch. const quotaFlags = classifyTestErrorQuota(error); - const isBotBlock = !quotaFlags.isQuota && (streamError.statusCode === 403 || isBotBlockMessage(error)); + const isBotBlock = + !quotaFlags.isQuota && (streamError.statusCode === 403 || isBotBlockMessage(error)); return { modelId: fullModelStr, status: rateLimited ? "rate_limited" : "error", diff --git a/src/lib/credentialHealth/scheduler.ts b/src/lib/credentialHealth/scheduler.ts index 6207a26549..580111b9a1 100644 --- a/src/lib/credentialHealth/scheduler.ts +++ b/src/lib/credentialHealth/scheduler.ts @@ -75,7 +75,6 @@ function isBuildProcess(): boolean { return typeof process !== "undefined" && process.env.NEXT_PHASE === "phase-production-build"; } - function isCredentialHealthCheckDisabled(): boolean { if (isBuildProcess() || isAutomatedTestProcess()) return true; const val = process.env.OMNIROUTE_DISABLE_CREDENTIAL_HEALTH_CHECK; @@ -124,6 +123,9 @@ async function testConnection( try { const result = await testSingleConnection(connectionId); + // A deliberate lease skip must not rewrite the health cache. + if (result.skipped === true) return; + const latencyMs = Date.now() - startTime; const state = getSchedulerState(); diff --git a/src/lib/db/apiKeys.ts b/src/lib/db/apiKeys.ts index 55d0dde43a..a19205f8a0 100644 --- a/src/lib/db/apiKeys.ts +++ b/src/lib/db/apiKeys.ts @@ -222,6 +222,22 @@ const _lastUsedUpdateCache = new Map(); const CACHE_TTL = 60 * 1000; // 1 minute TTL const LAST_USED_UPDATE_TTL = 5 * 60 * 1000; const MAX_CACHE_SIZE = 1000; +const EXCLUSIVE_LEASE_SCOPE = "lease:exclusive"; + +export class ApiKeyPolicyInvariantError extends Error { + readonly code = "LEASE_KEY_POLICY_INVALID"; +} + +function assertExclusiveLeaseKeyPolicy( + scopes: readonly string[], + allowedConnections: readonly string[] +): void { + if (scopes.includes(EXCLUSIVE_LEASE_SCOPE) && allowedConnections.length === 0) { + throw new ApiKeyPolicyInvariantError( + "lease:exclusive requires explicit allowedConnections" + ); + } +} // Prepared statements cache let _stmtGetAllKeys: ApiKeysStatements["getAllKeys"] | null = null; @@ -423,7 +439,7 @@ function getPreparedStatements(db: ApiKeysDbLike): ApiKeysStatements { "SELECT id, name, machine_id, model_access_mode, allowed_models, blocked_models, allowed_combos, allowed_connections, allowed_quotas, no_log, auto_resolve, is_active, access_schedule, max_requests_per_day, max_requests_per_minute, throttle_delay_ms, max_sessions, revoked_at, expires_at, ip_allowlist, scopes, rate_limits, is_banned, key_hash, allowed_endpoints, stream_default_mode, cache_default_mode, disable_non_public_models, allow_usage_command, usage_limit_enabled, daily_usage_limit_usd, weekly_usage_limit_usd, chaos_mode_enabled, compression_enabled, proxy_id FROM api_keys WHERE key = ? OR key_hash = ?", ); _stmtInsertKey = db.prepare( - "INSERT INTO api_keys (id, name, key, machine_id, allowed_models, allowed_combos, no_log, created_at, key_prefix, key_hash, scopes) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" + "INSERT INTO api_keys (id, name, key, machine_id, allowed_models, allowed_combos, allowed_connections, no_log, created_at, key_prefix, key_hash, scopes) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" ); _stmtDeleteKey = db.prepare("DELETE FROM api_keys WHERE id = ?"); } @@ -502,6 +518,19 @@ export function getApiKeysCount(): number { return row.cnt; } +/** Derived lease-only membership from existing key policy, not a second pool store. */ +export async function getExclusiveLeaseConnectionIds(): Promise> { + ensureApiKeysColumns(getDbInstance() as ApiKeysDbLike); + const rows = (getDbInstance() as ApiKeysDbLike) + .prepare( + `SELECT allowed_connections FROM api_keys + WHERE is_active != 0 AND is_banned != 1 AND revoked_at IS NULL + AND (expires_at IS NULL OR expires_at > ?) AND scopes LIKE ?` + ) + .all(new Date().toISOString(), `%"${EXCLUSIVE_LEASE_SCOPE}"%`); + return new Set(rows.flatMap((row) => parseAllowedConnections(row.allowed_connections))); +} + /** * Select an API key for internal OmniRoute operations (combo health checks, * cloud-sync verify pings, etc.). @@ -524,7 +553,7 @@ export function getApiKeysCount(): number { * behavior when no key matches the better rules above). * * The selector is deliberately conservative: it never promotes a revoked, - * inactive, or banned key, and it never widens a key's allowedModels. + * inactive, banned, or hard-lease key, and it never widens a key's allowedModels. */ export async function pickApiKeyForInternalUse( purpose: "combo-health-check" | "cloud-sync-verify" | "internal-probe" = "internal-probe", @@ -542,7 +571,11 @@ export async function pickApiKeyForInternalUse( }>; const isUsable = (k: (typeof keys)[number]) => - Boolean(k.key) && k.isActive !== false && !k.revokedAt && k.isBanned !== true; + Boolean(k.key) && + k.isActive !== false && + !k.revokedAt && + k.isBanned !== true && + !k.scopes?.includes(EXCLUSIVE_LEASE_SCOPE); // 1. Management-scoped key (preferred for any internal probe). const manageKey = keys.find( @@ -625,10 +658,17 @@ async function hashKey(key: string): Promise { return createHash("sha256").update(key).digest("hex"); // nosemgrep: insufficient-password-hash } -export async function createApiKey(name: string, machineId: string, scopes: string[] = []) { +export async function createApiKey( + name: string, + machineId: string, + scopes: string[] = [], + options: { allowedConnections?: string[] } = {} +) { if (!machineId) { throw new Error("machineId is required"); } + const allowedConnections = options.allowedConnections ?? []; + assertExclusiveLeaseKeyPolicy(scopes, allowedConnections); const db = getDbInstance() as ApiKeysDbLike; const now = new Date().toISOString(); @@ -644,7 +684,7 @@ export async function createApiKey(name: string, machineId: string, scopes: stri modelAccessMode: "all" as const, allowedModels: [], // Empty array means all models allowed allowedCombos: [ALL_COMBOS_ACCESS_RULE], // Explicit wildcard means all combos allowed - allowedConnections: [], // Empty array means all connections allowed + allowedConnections, noLog: false, allowUsageCommand: false, createdAt: now, @@ -659,6 +699,7 @@ export async function createApiKey(name: string, machineId: string, scopes: stri apiKey.machineId, "[]", JSON.stringify(apiKey.allowedCombos), + JSON.stringify(allowedConnections), 0, apiKey.createdAt, apiKey.key.slice(0, 12), @@ -968,9 +1009,20 @@ export async function updateApiKeyPermissions( db.exec("BEGIN IMMEDIATE"); try { const prevRow = db - .prepare<{ scopes: string | null }>("SELECT scopes FROM api_keys WHERE id = ?") + .prepare<{ scopes: string | null; allowed_connections: string | null }>( + "SELECT scopes, allowed_connections FROM api_keys WHERE id = ?" + ) .get(id); - previousScopes = parseStringList(prevRow?.scopes ?? null); + if (!prevRow) { + db.exec("ROLLBACK"); + return false; + } + previousScopes = parseStringList(prevRow.scopes); + const nextAllowedConnections = + normalized.allowedConnections === undefined + ? parseAllowedConnections(prevRow.allowed_connections) + : normalized.allowedConnections; + assertExclusiveLeaseKeyPolicy(nextScopes, nextAllowedConnections); const upd = db .prepare(`UPDATE api_keys SET ${updates.join(", ")} WHERE id = @id`) .run(params); @@ -987,6 +1039,28 @@ export async function updateApiKeyPermissions( } throw err; } + } else if (normalized.allowedConnections !== undefined) { + db.exec("BEGIN IMMEDIATE"); + try { + const row = db + .prepare<{ scopes: string | null }>("SELECT scopes FROM api_keys WHERE id = ?") + .get(id); + if (!row) { + db.exec("ROLLBACK"); + return false; + } + assertExclusiveLeaseKeyPolicy(parseStringList(row.scopes), normalized.allowedConnections); + const upd = db.prepare(`UPDATE api_keys SET ${updates.join(", ")} WHERE id = @id`).run(params); + changedRows = upd.changes ?? 0; + db.exec("COMMIT"); + } catch (err) { + try { + db.exec("ROLLBACK"); + } catch { + // Preserve the mutation failure if rollback also fails. + } + throw err; + } } else { const upd = db.prepare(`UPDATE api_keys SET ${updates.join(", ")} WHERE id = @id`).run(params); changedRows = upd.changes ?? 0; diff --git a/src/lib/db/exclusiveConnectionLeases.ts b/src/lib/db/exclusiveConnectionLeases.ts new file mode 100644 index 0000000000..f93453d58f --- /dev/null +++ b/src/lib/db/exclusiveConnectionLeases.ts @@ -0,0 +1,394 @@ +import { createHash } from "node:crypto"; + +import { getDbInstance, rowToCamel } from "./core"; + +export const LEASE_OWNER_PATTERN = /^vlo_[A-Za-z0-9_-]{43}$/; +const DEFAULT_EXCLUSIVE_LEASE_TTL_MS = 120_000; +const MIN_EXCLUSIVE_LEASE_TTL_MS = 1_000; +const MAX_EXCLUSIVE_LEASE_TTL_MS = 1_800_000; +type ExclusiveLeaseState = "ACTIVE" | "RELEASED" | "EXPIRED" | "INVALIDATED"; +export type ExclusiveLeaseEndReason = + | "AUTHORIZATION_CHANGED" + | "CLIENT_CANCELLED" + | "CONNECTION_INELIGIBLE" + | "HEALTH_OR_COOLDOWN" + | "MANAGED_KEY_REVOKED" + | "MODEL_INELIGIBLE" + | "OWNER_EXIT" + | "QUOTA_UNAVAILABLE" + | "TTL_EXPIRED"; +export type ExclusiveConnectionLease = { + id: number; + leaseOwnerHash: string; + apiKeyId: string; + provider: string; + connectionId: string; + generation: number; + state: ExclusiveLeaseState; + acquiredAt: string; + renewedAt: string; + expiresAt: string; + endedAt: string | null; + endReason: string | null; +}; +type LeaseRow = { + id: number; + api_key_id: string; + provider: string; + connection_id: string; + generation: number; + state: ExclusiveLeaseState; + expires_at: string; +}; +type LeaseSuccess = { + kind: "ACQUIRED" | "REUSED" | "TRANSITIONED"; + lease: ExclusiveConnectionLease; +}; +type LeaseConflict = + | { kind: "CONNECTION_BUSY"; retryAfter: string | null } + | { kind: "OWNER_ALREADY_ACTIVE"; lease: ExclusiveConnectionLease }; +type LeaseUpdateResult = + { kind: T; lease: ExclusiveConnectionLease } | { kind: "STALE" }; + +const database = () => getDbInstance(); +const ACTIVE_SQL = "SELECT * FROM exclusive_connection_leases WHERE state = 'ACTIVE' AND "; +const lease = (row: LeaseRow) => rowToCamel(row) as ExclusiveConnectionLease; +function timestamp(value?: string): string { + const parsed = Date.parse(value ?? new Date().toISOString()); + if (!Number.isFinite(parsed)) throw new Error("now must be a valid ISO timestamp"); + return new Date(parsed).toISOString(); +} +function expiry(now: string, ttlMs?: number): string { + const ttl = Math.min( + MAX_EXCLUSIVE_LEASE_TTL_MS, + Math.max(MIN_EXCLUSIVE_LEASE_TTL_MS, ttlMs ?? DEFAULT_EXCLUSIVE_LEASE_TTL_MS) + ); + return new Date(Date.parse(now) + ttl).toISOString(); +} + +export function hashLeaseOwnerId(leaseOwnerId: string): string { + if (!LEASE_OWNER_PATTERN.test(leaseOwnerId)) + throw new Error("lease owner must use the canonical vlo_ base64url format"); + return createHash("sha256").update(leaseOwnerId).digest("hex"); +} + +function immediate(operation: () => T): T { + let result: T | undefined; + database().immediate(() => (result = operation())); + if (result === undefined) throw new Error("lease transaction did not produce a result"); + return result; +} + +function expire(now: string): number { + return database() + .prepare( + `UPDATE exclusive_connection_leases + SET state = 'EXPIRED', ended_at = ?, end_reason = 'TTL_EXPIRED' + WHERE state = 'ACTIVE' AND expires_at <= ?` + ) + .run(now, now).changes; +} +function active(column: "lease_owner_hash" | "connection_id", value: string) { + return database().prepare(`${ACTIVE_SQL}${column} = ?`).get(value) as LeaseRow | undefined; +} + +function historical(ownerHash: string, generation: number) { + return database() + .prepare( + `SELECT * FROM exclusive_connection_leases + WHERE lease_owner_hash = ? AND generation = ? ORDER BY id DESC LIMIT 1` + ) + .get(ownerHash, generation) as LeaseRow | undefined; +} + +function nextGeneration(ownerHash: string): number { + const row = database() + .prepare( + "SELECT COALESCE(MAX(generation), 0) AS generation FROM exclusive_connection_leases WHERE lease_owner_hash = ?" + ) + .get(ownerHash) as { generation: number }; + const generation = Number(row.generation) + 1; + if (!Number.isSafeInteger(generation) || generation <= 0) + throw new Error("lease generation exhausted"); + return generation; +} + +function insert(input: { + ownerHash: string; + apiKeyId: string; + provider: string; + connectionId: string; + generation: number; + now: string; + ttlMs?: number; +}): ExclusiveConnectionLease { + const result = database() + .prepare( + `INSERT INTO exclusive_connection_leases + (lease_owner_hash, api_key_id, provider, connection_id, generation, state, + acquired_at, renewed_at, expires_at) + VALUES (?, ?, ?, ?, ?, 'ACTIVE', ?, ?, ?)` + ) + .run( + input.ownerHash, + input.apiKeyId, + input.provider, + input.connectionId, + input.generation, + input.now, + input.now, + expiry(input.now, input.ttlMs) + ); + return lease( + database() + .prepare("SELECT * FROM exclusive_connection_leases WHERE id = ?") + .get(result.lastInsertRowid) as LeaseRow + ); +} + +function isLeaseConflict(error: unknown): boolean { + return ( + error instanceof Error && + (/UNIQUE constraint failed: exclusive_connection_leases\.(connection_id|lease_owner_hash)/i.test( + error.message + ) || + /idx_exclusive_lease_active_(connection|owner)/i.test(error.message)) + ); +} + +function conflict(ownerHash: string, connectionId: string): LeaseConflict { + const owner = active("lease_owner_hash", ownerHash); + return owner + ? { kind: "OWNER_ALREADY_ACTIVE", lease: lease(owner) } + : { + kind: "CONNECTION_BUSY", + retryAfter: active("connection_id", connectionId)?.expires_at ?? null, + }; +} + +function update(input: { + ownerHash: string; + generation: number; + now: string; + kind: T; + sql: string; + args: unknown[]; + accept?: (row: LeaseRow | undefined) => boolean; +}): LeaseUpdateResult { + return immediate(() => { + expire(input.now); + const changed = + database() + .prepare(input.sql) + .run(...input.args).changes === 1; + const row = historical(input.ownerHash, input.generation); + return changed || input.accept?.(row) + ? { kind: input.kind, lease: lease(row!) } + : { kind: "STALE" }; + }); +} + +export function reconcileExpiredExclusiveConnectionLeases(now?: string): number { + return immediate(() => expire(timestamp(now))); +} + +export function acquireExclusiveConnectionLease(input: { + leaseOwnerId: string; + apiKeyId: string; + provider: string; + connectionId: string; + now?: string; + ttlMs?: number; +}): LeaseSuccess | LeaseConflict { + const ownerHash = hashLeaseOwnerId(input.leaseOwnerId); + const now = timestamp(input.now); + try { + return immediate(() => { + expire(now); + const owner = active("lease_owner_hash", ownerHash); + if (owner) { + if (owner.api_key_id !== input.apiKeyId || owner.connection_id !== input.connectionId) { + return { kind: "OWNER_ALREADY_ACTIVE", lease: lease(owner) }; + } + database() + .prepare( + `UPDATE exclusive_connection_leases SET renewed_at = ?, expires_at = ?, + provider = ? WHERE id = ? AND state = 'ACTIVE' AND api_key_id = ?` + ) + .run(now, expiry(now, input.ttlMs), input.provider, owner.id, input.apiKeyId); + return { kind: "REUSED", lease: lease(active("lease_owner_hash", ownerHash)!) }; + } + const occupied = active("connection_id", input.connectionId); + if (occupied) return { kind: "CONNECTION_BUSY", retryAfter: occupied.expires_at }; + return { + kind: "ACQUIRED", + lease: insert({ + ...input, + ownerHash, + generation: nextGeneration(ownerHash), + now, + }), + }; + }); + } catch (error) { + if (!isLeaseConflict(error)) throw error; + return conflict(ownerHash, input.connectionId); + } +} + +export function transitionExclusiveConnectionLease(input: { + leaseOwnerId: string; + generation: number; + apiKeyId: string; + provider: string; + connectionId: string; + reason: ExclusiveLeaseEndReason; + now?: string; + ttlMs?: number; +}): LeaseSuccess | LeaseConflict | { kind: "STALE" } { + const ownerHash = hashLeaseOwnerId(input.leaseOwnerId); + const now = timestamp(input.now); + try { + return immediate(() => { + expire(now); + const owner = active("lease_owner_hash", ownerHash); + if (!owner || owner.generation !== input.generation || owner.api_key_id !== input.apiKeyId) + return { kind: "STALE" }; + if (owner.connection_id === input.connectionId) { + return { kind: "REUSED", lease: lease(owner) }; + } + const occupied = active("connection_id", input.connectionId); + if (occupied) return { kind: "CONNECTION_BUSY", retryAfter: occupied.expires_at }; + database() + .prepare( + `UPDATE exclusive_connection_leases SET state = 'INVALIDATED', ended_at = ?, + end_reason = ? WHERE id = ? AND state = 'ACTIVE'` + ) + .run(now, input.reason, owner.id); + return { kind: "TRANSITIONED", lease: insert({ ...input, ownerHash, now }) }; + }); + } catch (error) { + if (!isLeaseConflict(error)) throw error; + return conflict(ownerHash, input.connectionId); + } +} + +export function invalidateExclusiveConnectionLease(input: { + leaseOwnerId: string; + generation: number; + apiKeyId: string; + reason: ExclusiveLeaseEndReason; + now?: string; +}): LeaseUpdateResult<"INVALIDATED"> { + const ownerHash = hashLeaseOwnerId(input.leaseOwnerId); + const now = timestamp(input.now); + return update({ + ownerHash, + generation: input.generation, + now, + kind: "INVALIDATED", + sql: `UPDATE exclusive_connection_leases SET state = 'INVALIDATED', ended_at = ?, end_reason = ? + WHERE lease_owner_hash = ? AND generation = ? AND api_key_id = ? + AND state = 'ACTIVE' AND expires_at > ?`, + args: [now, input.reason, ownerHash, input.generation, input.apiKeyId, now], + }); +} + +export function renewExclusiveConnectionLease(input: { + leaseOwnerId: string; + generation: number; + apiKeyId: string; + now?: string; + ttlMs?: number; +}): LeaseUpdateResult<"RENEWED"> { + const ownerHash = hashLeaseOwnerId(input.leaseOwnerId); + const now = timestamp(input.now); + return update({ + ownerHash, + generation: input.generation, + now, + kind: "RENEWED", + sql: `UPDATE exclusive_connection_leases SET renewed_at = ?, expires_at = ? + WHERE lease_owner_hash = ? AND generation = ? AND api_key_id = ? + AND state = 'ACTIVE' AND expires_at > ?`, + args: [now, expiry(now, input.ttlMs), ownerHash, input.generation, input.apiKeyId, now], + }); +} + +export function releaseExclusiveConnectionLease(input: { + leaseOwnerId: string; + generation: number; + apiKeyId: string; + reason?: "OWNER_EXIT" | "CLIENT_CANCELLED"; + now?: string; +}): LeaseUpdateResult<"RELEASED"> { + const ownerHash = hashLeaseOwnerId(input.leaseOwnerId); + const now = timestamp(input.now); + return update({ + ownerHash, + generation: input.generation, + now, + kind: "RELEASED", + sql: `UPDATE exclusive_connection_leases SET state = 'RELEASED', ended_at = ?, end_reason = ? + WHERE lease_owner_hash = ? AND generation = ? AND api_key_id = ? + AND state = 'ACTIVE' AND expires_at > ?`, + args: [now, input.reason ?? "OWNER_EXIT", ownerHash, input.generation, input.apiKeyId, now], + accept: (row) => row?.state === "RELEASED" && row.api_key_id === input.apiKeyId, + }); +} + +export function assertExclusiveConnectionLeaseFence(input: { + leaseOwnerId: string; + generation: number; + apiKeyId: string; + connectionId: string; + now?: string; +}): + | { kind: "VALID"; lease: ExclusiveConnectionLease } + | { kind: "AUTHORIZATION_MISMATCH" | "CONNECTION_MISMATCH"; lease: ExclusiveConnectionLease } + | { kind: "REQUIRED" | "STALE" } { + const ownerHash = hashLeaseOwnerId(input.leaseOwnerId); + reconcileExpiredExclusiveConnectionLeases(input.now); + const row = active("lease_owner_hash", ownerHash); + if (!row) return { kind: "REQUIRED" }; + if (row.generation !== input.generation) return { kind: "STALE" }; + const current = lease(row); + if (row.api_key_id !== input.apiKeyId) return { kind: "AUTHORIZATION_MISMATCH", lease: current }; + return row.connection_id === input.connectionId + ? { kind: "VALID", lease: current } + : { kind: "CONNECTION_MISMATCH", lease: current }; +} + +export function getActiveExclusiveConnectionLease(leaseOwnerId: string, now?: string) { + const ownerHash = hashLeaseOwnerId(leaseOwnerId); + reconcileExpiredExclusiveConnectionLeases(now); + const row = active("lease_owner_hash", ownerHash); + return row ? lease(row) : null; +} + +export function getExclusiveLeaseOccupancy(connectionIds: readonly string[], now?: string) { + reconcileExpiredExclusiveConnectionLeases(now); + if (connectionIds.length === 0) return new Map(); + const rows = database() + .prepare( + `SELECT connection_id, lease_owner_hash, expires_at FROM exclusive_connection_leases + WHERE state = 'ACTIVE' AND connection_id IN (${connectionIds.map(() => "?").join(", ")})` + ) + .all(...connectionIds) as Array<{ + connection_id: string; + lease_owner_hash: string; + expires_at: string; + }>; + return new Map( + rows.map((row) => [ + row.connection_id, + { leaseOwnerHash: row.lease_owner_hash, expiresAt: row.expires_at }, + ]) + ); +} + +export function isExclusiveConnectionActivelyLeased(connectionId: string, now?: string): boolean { + if (!connectionId) return false; + reconcileExpiredExclusiveConnectionLeases(now); + return active("connection_id", connectionId) !== undefined; +} diff --git a/src/lib/db/migrations/155_exclusive_connection_leases.sql b/src/lib/db/migrations/155_exclusive_connection_leases.sql new file mode 100644 index 0000000000..1d9786bd03 --- /dev/null +++ b/src/lib/db/migrations/155_exclusive_connection_leases.sql @@ -0,0 +1,30 @@ +CREATE TABLE IF NOT EXISTS exclusive_connection_leases ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + lease_owner_hash TEXT NOT NULL + CHECK (length(lease_owner_hash) = 64 AND lease_owner_hash = lower(lease_owner_hash)), + api_key_id TEXT NOT NULL, + provider TEXT NOT NULL, + connection_id TEXT NOT NULL, + generation INTEGER NOT NULL CHECK (generation > 0), + state TEXT NOT NULL + CHECK (state IN ('ACTIVE', 'RELEASED', 'EXPIRED', 'INVALIDATED')), + acquired_at TEXT NOT NULL, + renewed_at TEXT NOT NULL, + expires_at TEXT NOT NULL, + ended_at TEXT, + end_reason TEXT +); + +CREATE UNIQUE INDEX IF NOT EXISTS idx_exclusive_lease_active_owner + ON exclusive_connection_leases(lease_owner_hash) + WHERE state = 'ACTIVE'; + +CREATE UNIQUE INDEX IF NOT EXISTS idx_exclusive_lease_active_connection + ON exclusive_connection_leases(connection_id) + WHERE state = 'ACTIVE'; + +CREATE INDEX IF NOT EXISTS idx_exclusive_lease_owner_history + ON exclusive_connection_leases(lease_owner_hash, generation, id); + +CREATE INDEX IF NOT EXISTS idx_exclusive_lease_expiry + ON exclusive_connection_leases(state, expires_at); diff --git a/src/lib/exclusiveLeaseIsolation.ts b/src/lib/exclusiveLeaseIsolation.ts new file mode 100644 index 0000000000..62677378a2 --- /dev/null +++ b/src/lib/exclusiveLeaseIsolation.ts @@ -0,0 +1,8 @@ +import { getExclusiveLeaseConnectionIds } from "./db/apiKeys"; +import { isExclusiveConnectionActivelyLeased } from "./db/exclusiveConnectionLeases"; + +/** Narrow fail-closed boundary for auxiliary/unmanaged connection activity. */ +export async function isConnectionUnavailableToAuxiliaryActivity(connectionId: string) { + if (!connectionId || isExclusiveConnectionActivelyLeased(connectionId)) return true; + return (await getExclusiveLeaseConnectionIds()).has(connectionId); +} diff --git a/src/lib/localDb.ts b/src/lib/localDb.ts index 81c8ff9ae5..5623c6dddf 100755 --- a/src/lib/localDb.ts +++ b/src/lib/localDb.ts @@ -112,6 +112,7 @@ export { pickApiKeyForInternalUse, clearApiKeyCaches, resetApiKeyState, + ApiKeyPolicyInvariantError, } from "./db/apiKeys"; export { @@ -174,6 +175,8 @@ export { export type { UserDatabaseSettings } from "./db/databaseSettings"; +export * from "./db/exclusiveConnectionLeases"; + export { // Proxy Registry listProxies, diff --git a/src/lib/services/quotaAutoPing.ts b/src/lib/services/quotaAutoPing.ts index 069ed5a8c2..94e5ea8247 100644 --- a/src/lib/services/quotaAutoPing.ts +++ b/src/lib/services/quotaAutoPing.ts @@ -25,6 +25,7 @@ import { sanitizeErrorMessage } from "@omniroute/open-sse/utils/error.ts"; import { getExecutor } from "@omniroute/open-sse/executors/index.ts"; import { getCodexUsage } from "@omniroute/open-sse/services/usage/codex.ts"; import { getSettings, getProviderConnections, updateProviderConnection } from "@/lib/localDb"; +import { isConnectionUnavailableToAuxiliaryActivity } from "@/lib/exclusiveLeaseIsolation"; import { refreshAndUpdateCredentials } from "@/lib/usage/providerLimits"; import { getCircuitBreaker } from "@/shared/utils/circuitBreaker"; import { @@ -68,6 +69,7 @@ export interface QuotaAutoPingDeps { ) => Promise; getExecutor: (provider: string) => { execute: (input: JsonRecord) => Promise }; canExecuteProvider: (provider: string) => boolean; + isConnectionUnavailableToAuxiliaryActivity: (connectionId: string) => Promise; } export interface QuotaAutoPingState { @@ -90,6 +92,7 @@ export function createDefaultQuotaAutoPingDeps(): QuotaAutoPingDeps { getCodexUsage, getExecutor, canExecuteProvider: (provider) => getCircuitBreaker(provider).canExecute(), + isConnectionUnavailableToAuxiliaryActivity, }; } @@ -246,7 +249,7 @@ function shouldPingForReset( * Cheap pre-fetch guards — none of these require a network call. Extracted so * `pingConnection` reads as a single linear flow instead of a wall of `if`s. */ -function isPingCandidateBlocked( +async function isPingCandidateBlocked( connection: QuotaAutoPingConnection, provider: "codex", providerConfig: QuotaAutoPingProviderConfig, @@ -255,8 +258,9 @@ function isPingCandidateBlocked( key: string, cachedReset: string | undefined, nowMs: number -): boolean { +): Promise { if (!deps.canExecuteProvider(provider)) return true; // provider circuit breaker OPEN + if (await deps.isConnectionUnavailableToAuxiliaryActivity(connection.id)) return true; if (isRateLimited(connection, nowMs)) return true; // connection cooldown active if (shouldSkipAfterFailure(state, key, nowMs)) return true; @@ -326,7 +330,7 @@ async function pingConnection( ): Promise { const key = cacheKey(provider, connection.id); const cachedReset = state.resetCache[key]; - if (isPingCandidateBlocked(connection, provider, providerConfig, deps, state, key, cachedReset, nowMs)) { + if (await isPingCandidateBlocked(connection, provider, providerConfig, deps, state, key, cachedReset, nowMs)) { return; } diff --git a/src/lib/usage/codexResetCredits.ts b/src/lib/usage/codexResetCredits.ts index ca99d72c29..a0097173d0 100644 --- a/src/lib/usage/codexResetCredits.ts +++ b/src/lib/usage/codexResetCredits.ts @@ -1,4 +1,5 @@ import { getProviderConnectionById, resolveProxyForConnection } from "@/lib/localDb"; +import { isConnectionUnavailableToAuxiliaryActivity } from "@/lib/exclusiveLeaseIsolation"; import { fetchAndPersistProviderLimits, refreshAndUpdateCredentials, @@ -299,6 +300,13 @@ function buildCodexResetCreditHeaders(connection: CodexConnectionLike): Record { + if (await isConnectionUnavailableToAuxiliaryActivity(connectionId)) { + throw new CodexResetCreditError( + 409, + "exclusive_lease_active", + "Reset-credit operations are deferred while an exclusive lease is active." + ); + } const connection = (await getProviderConnectionById( connectionId )) as unknown as CodexConnectionLike | null; diff --git a/src/lib/usage/providerLimits.ts b/src/lib/usage/providerLimits.ts index 11aa8444b6..7e8811a9e9 100644 --- a/src/lib/usage/providerLimits.ts +++ b/src/lib/usage/providerLimits.ts @@ -19,6 +19,7 @@ import { isClaudeExtraUsageBlockEnabled, isClaudeExtraUsageQueued, } from "@/lib/providers/claudeExtraUsage"; +import { isConnectionUnavailableToAuxiliaryActivity } from "@/lib/exclusiveLeaseIsolation"; import { clearRecoveredProviderState } from "@/sse/services/auth"; import { getMachineId } from "@/shared/utils/machine"; import { USAGE_SUPPORTED_PROVIDERS } from "@/shared/constants/providers"; @@ -793,6 +794,9 @@ async function fetchLiveProviderLimitsWithOptions( connection: ProviderConnectionLike; usage: JsonRecord; }> { + if (await isConnectionUnavailableToAuxiliaryActivity(connectionId)) { + throw withStatus(new Error("Usage refresh deferred while an exclusive lease is active"), 409); + } let connection = (await getProviderConnectionById( connectionId )) as unknown as ProviderConnectionLike | null; @@ -1000,9 +1004,19 @@ export async function syncAllProviderLimits( errors: Record; }> { const { source = "manual", concurrency = 5 } = options; + const connectionRows = (await getProviderConnections({ + isActive: true, + })) as unknown as ProviderConnectionLike[]; const connections = ( - (await getProviderConnections({ isActive: true })) as unknown as ProviderConnectionLike[] - ).filter(isSupportedUsageConnection); + await Promise.all( + connectionRows.map(async (connection) => ({ + connection, + blocked: await isConnectionUnavailableToAuxiliaryActivity(connection.id), + })) + ) + ) + .filter(({ connection, blocked }) => isSupportedUsageConnection(connection) && !blocked) + .map(({ connection }) => connection); const cacheEntries: Array<{ connectionId: string; entry: ProviderLimitsCacheEntry }> = []; const caches: Record = {}; const errors: Record = {}; diff --git a/src/lib/vncSession/service.ts b/src/lib/vncSession/service.ts index 06332050d0..92f57ce5f0 100644 --- a/src/lib/vncSession/service.ts +++ b/src/lib/vncSession/service.ts @@ -2,6 +2,7 @@ import { spawn } from "node:child_process"; import { randomUUID } from "node:crypto"; import { chmodSync, mkdirSync, rmSync } from "node:fs"; import { join } from "node:path"; +import { isConnectionUnavailableToAuxiliaryActivity } from "@/lib/exclusiveLeaseIsolation"; import { getProviderConnectionById, updateProviderConnection } from "@/lib/db/providers"; import { validateProviderApiKey } from "@/lib/providers/validation"; import { VNC_CONFIG, getVncProvider } from "./manifest"; @@ -153,6 +154,9 @@ async function publishedPort(containerName: string, containerPort: number): Prom } export async function startSession(connectionId: string): Promise { + if (await isConnectionUnavailableToAuxiliaryActivity(connectionId)) + throw new Error("Browser login is unavailable for managed lease connections"); + await reconcileStaleContainers(); const connection = await getProviderConnectionById(connectionId); @@ -254,6 +258,9 @@ export async function harvestSession( connectionId: string, sessionId: string ): Promise { + if (await isConnectionUnavailableToAuxiliaryActivity(connectionId)) + throw new Error("Browser login is unavailable for managed lease connections"); + const session = getSession(connectionId, sessionId); if (!session) throw new Error("Browser-login session not found"); if (session.status !== "running") { diff --git a/src/lib/warmupScheduler.ts b/src/lib/warmupScheduler.ts index c0cc77fe71..7ecd44217f 100644 --- a/src/lib/warmupScheduler.ts +++ b/src/lib/warmupScheduler.ts @@ -8,6 +8,7 @@ import { logger } from "@omniroute/open-sse/utils/logger"; import { matchesCron } from "@/lib/jobs/cronMatch"; import { getCircuitBreakerStore } from "./warmupScheduler/circuitBreakerFactory"; import { TERMINAL_CONNECTION_STATUSES } from "@/lib/quota/connectionRecovery"; +import { isConnectionUnavailableToAuxiliaryActivity } from "@/lib/exclusiveLeaseIsolation"; import type { WarmupResult, WarmupFailureKind, WarmupTarget } from "./warmupScheduler/core"; export type { WarmupResult, WarmupFailureKind } from "./warmupScheduler/core"; @@ -155,6 +156,7 @@ async function executeWarmup(): Promise { const headers = await getWarmupHeaders(); for (const conn of connections) { + if (await isConnectionUnavailableToAuxiliaryActivity(conn.id)) continue; if (enabledMap?.[conn.id] !== true) { log.debug("warmup skip", { connectionId: conn.id, reason: "not opted-in" }); continue; diff --git a/src/server/cors/origins.ts b/src/server/cors/origins.ts index d610297641..dfd0b407c5 100644 --- a/src/server/cors/origins.ts +++ b/src/server/cors/origins.ts @@ -21,7 +21,7 @@ const ENV_ALLOWED = "CORS_ALLOWED_ORIGINS"; const LEGACY_ENV_SINGLE = "CORS_ORIGIN"; const STANDARD_ALLOW_HEADERS = - "Content-Type, Authorization, x-api-key, anthropic-version, x-omniroute-connection, x-internal-test, accept"; + "Content-Type, Authorization, x-api-key, anthropic-version, x-omniroute-connection, X-OmniRoute-Lease-Owner, X-OmniRoute-Lease-Generation, x-internal-test, accept"; const STANDARD_ALLOW_METHODS = "GET, POST, PUT, DELETE, PATCH, OPTIONS"; let runtimeAllowedOrigins: ReadonlySet = new Set(); diff --git a/src/shared/services/modelSyncScheduler.ts b/src/shared/services/modelSyncScheduler.ts index 026ebee310..c705f6f977 100644 --- a/src/shared/services/modelSyncScheduler.ts +++ b/src/shared/services/modelSyncScheduler.ts @@ -11,6 +11,7 @@ import { randomUUID } from "node:crypto"; import { Agent, buildConnector, fetch as undiciFetch, type Dispatcher } from "undici"; import { getSettings, updateSettings } from "@/lib/localDb"; +import { isConnectionUnavailableToAuxiliaryActivity } from "@/lib/exclusiveLeaseIsolation"; import { getRuntimePorts } from "@/lib/runtime/ports"; const DEFAULT_INTERVAL_MS = 24 * 60 * 60 * 1000; // 24 hours @@ -155,6 +156,11 @@ async function getAutoSyncConnections(): Promise< const autoSyncConnections: Array<{ id: string; provider: string; name?: string }> = []; for (const conn of connections) { if (!conn.isActive && conn.isActive !== undefined) continue; + if ( + typeof conn.id === "string" && + (await isConnectionUnavailableToAuxiliaryActivity(conn.id)) + ) + continue; const psd = conn.providerSpecificData && typeof conn.providerSpecificData === "object" ? (conn.providerSpecificData as Record) diff --git a/src/shared/utils/cors.ts b/src/shared/utils/cors.ts index 75ef54d5ac..a014b7dba4 100644 --- a/src/shared/utils/cors.ts +++ b/src/shared/utils/cors.ts @@ -11,7 +11,7 @@ export const CORS_HEADERS = { "Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, PATCH, OPTIONS", "Access-Control-Allow-Headers": - "Content-Type, Authorization, x-api-key, anthropic-version, x-omniroute-connection, x-internal-test, accept", + "Content-Type, Authorization, x-api-key, anthropic-version, x-omniroute-connection, X-OmniRoute-Lease-Owner, X-OmniRoute-Lease-Generation, x-internal-test, accept", } as const; /** diff --git a/src/shared/utils/shuffleDeck.ts b/src/shared/utils/shuffleDeck.ts index 76d67d8136..1404b10474 100644 --- a/src/shared/utils/shuffleDeck.ts +++ b/src/shared/utils/shuffleDeck.ts @@ -138,6 +138,33 @@ export function getNextFromDeckSync(namespace: string, itemIds: readonly string[ return newOrder[0]; } +/** Plan a deck selection without advancing shared state until commit. */ +export function planNextFromDeckSync(namespace: string, itemIds: readonly string[]) { + if (itemIds.length === 0) return { selectedId: "", commit: () => {} }; + if (itemIds.length === 1) return { selectedId: itemIds[0], commit: () => {} }; + + const idsKey = [...itemIds].sort().join(","); + const existing = decks.get(namespace); + if (existing && existing.idsKey === idsKey && existing.index < existing.order.length) { + const selectedId = existing.order[existing.index]; + return { + selectedId, + commit: () => decks.set(namespace, { ...existing, index: existing.index + 1 }), + }; + } + + const lastUsedId = + existing && existing.idsKey === idsKey && existing.order.length > 0 + ? existing.order[existing.order.length - 1] + : undefined; + const order = fisherYatesShuffle(itemIds); + if (lastUsedId !== undefined && order[0] === lastUsedId && order.length > 1) { + const swapIdx = 1 + secureRandomInt(order.length - 1); + [order[0], order[swapIdx]] = [order[swapIdx], order[0]]; + } + return { selectedId: order[0], commit: () => decks.set(namespace, { order, index: 1, idsKey }) }; +} + // ─── Test helpers ─────────────────────────────────────────────────────────── /** Reset all decks — for testing only. */ diff --git a/src/shared/validation/schemas/keys.ts b/src/shared/validation/schemas/keys.ts index 04a2f2d4e4..56875c3c1d 100644 --- a/src/shared/validation/schemas/keys.ts +++ b/src/shared/validation/schemas/keys.ts @@ -18,16 +18,30 @@ import { accessScheduleSchema } from "./misc.ts"; // ──── API Key Schemas ──── -export const createKeySchema = z.object({ - name: z.string().min(1, "Name is required").max(200), - noLog: z.boolean().optional(), - allowUsageCommand: z.boolean().optional(), - usageLimitEnabled: z.boolean().optional(), - dailyUsageLimitUsd: z.coerce.number().min(0).optional().nullable(), - weeklyUsageLimitUsd: z.coerce.number().min(0).optional().nullable(), - chaosModeEnabled: z.boolean().optional(), - scopes: z.array(z.string().trim().min(1).max(64)).max(32).optional(), -}); +const requireExclusiveLeaseConnections = (value: { + scopes?: string[]; allowedConnections?: string[]; +}, ctx: z.RefinementCtx) => { + if (value.scopes?.includes("lease:exclusive") && !value.allowedConnections?.length) + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "lease:exclusive requires explicit allowedConnections", + path: ["allowedConnections"], + }); +}; + +export const createKeySchema = z + .object({ + name: z.string().min(1, "Name is required").max(200), + noLog: z.boolean().optional(), + allowUsageCommand: z.boolean().optional(), + usageLimitEnabled: z.boolean().optional(), + dailyUsageLimitUsd: z.coerce.number().min(0).optional().nullable(), + weeklyUsageLimitUsd: z.coerce.number().min(0).optional().nullable(), + chaosModeEnabled: z.boolean().optional(), + scopes: z.array(z.string().trim().min(1).max(64)).max(32).optional(), + allowedConnections: z.array(z.string().uuid()).min(1).max(100).optional(), + }) + .superRefine(requireExclusiveLeaseConnections); export const createSyncTokenSchema = z.object({ name: z.string().trim().min(1, "Name is required").max(200), @@ -157,4 +171,7 @@ export const updateKeyPermissionsSchema = z path: [], }); } + if (value.scopes !== undefined && value.allowedConnections !== undefined) { + requireExclusiveLeaseConnections(value, ctx); + } }); diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts index 3845c1970f..3a1732d848 100644 --- a/src/sse/handlers/chat.ts +++ b/src/sse/handlers/chat.ts @@ -32,8 +32,12 @@ import { getImageModelEntry } from "@omniroute/open-sse/config/imageRegistry.ts" import { acceptHeaderForcesStream } from "@omniroute/open-sse/utils/aiSdkCompat.ts"; import { applyNoThinkingAlias } from "@omniroute/open-sse/utils/noThinkingAlias.ts"; import { resolveCcDiscoveryAliasStrip } from "@/lib/ccDiscoveryAliasResolve"; -import { handleComboChat, shouldSkipConnDisable } from "@omniroute/open-sse/services/combo.ts"; -import type { SingleModelTarget } from "@omniroute/open-sse/services/combo/types.ts"; +import { + handleComboChat, + resolveComboTargets, + shouldSkipConnDisable, +} from "@omniroute/open-sse/services/combo.ts"; +import type { ComboLike, SingleModelTarget } from "@omniroute/open-sse/services/combo/types.ts"; import { mergeAbortSignals } from "@omniroute/open-sse/executors/base.ts"; import { resolveRequestAutoControls } from "@omniroute/open-sse/services/autoCombo/requestControls.ts"; import { isVerifiedNativeCodexRequest } from "@omniroute/open-sse/config/codexIdentity.ts"; @@ -172,6 +176,16 @@ import { } from "../services/cooldownAwareRetry"; import { constrainConnectionsToQuota, resolveQuotaKeyScope } from "../../lib/quota/quotaKey"; import { checkConnectionCapacity } from "../utils/backpressure"; +import { + buildManagedLeaseErrorResponse, + buildManagedLeaseSelectionErrorResponse, + credentialLease, + isExclusiveLeaseManagedKey, + LeaseContextError, + parseManagedLeaseRequestContext, + validateExclusiveLeaseKeyConfiguration, + type ManagedLeaseDispatchContext, +} from "../services/leaseContext"; registerCodexQuotaFetcher(); @@ -212,7 +226,7 @@ registerGrokWebQuotaFetcher(); // what lets the per-window cutoff modal in Dashboard › Limits actually // enforce thresholds for Claude / GLM / Cursor / etc., not just Codex. registerGenericQuotaFetchers(); -let combosCachePromise: Promise | null = null; +let combosCachePromise: Promise | null = null; let combosCacheTs = 0; let combosCacheVersionSnapshot = -1; const COMBOS_CACHE_TTL_MS = 10_000; @@ -242,7 +256,7 @@ async function resolveComboContextOverflowDeferral( } } -async function getCombosCachedForChat(): Promise { +async function getCombosCachedForChat(): Promise { const now = Date.now(); // Explicit non-null check: we intentionally cache and return the Promise // itself (to dedupe concurrent callers), so this is not a forgotten await. @@ -259,7 +273,7 @@ async function getCombosCachedForChat(): Promise { combosCacheTs = now; combosCacheVersionSnapshot = getCombosCacheVersion(); - combosCachePromise = getCombos().catch(() => []); + combosCachePromise = getCombos().catch(() => []) as Promise; return combosCachePromise; } @@ -282,6 +296,44 @@ function intersectAllowedConnectionIds(primary: unknown, secondary: unknown): st return first || second || null; } +function isManagedComboUnsupported( + combo: ComboLike, + settings: Record, + allCombos: ComboLike[], + visited = new Set() +): boolean { + if (visited.has(combo.name)) return false; + visited.add(combo.name); + const strategy = combo.strategy ?? "priority"; + const config = resolveComboConfig(combo, settings) as Record; + const resolvedTargets = resolveComboTargets(combo, allCombos); + const pipeline = + strategy === "pipeline" || + (strategy === "auto" && (config.pipeline_enabled === true || combo.name === "auto/smart")); + const nestedUnsafe = (combo.models as Array<{ kind?: string; comboName?: string }>).some( + (step) => { + if (step?.kind !== "combo-ref" || !step.comboName) return false; + const nested = allCombos.find((candidate) => candidate.name === step.comboName); + return Boolean(nested && isManagedComboUnsupported(nested, settings, allCombos, visited)); + } + ); + return ( + strategy === "fusion" || + strategy === "context-relay" || + (config.chaos as { enabled?: boolean } | undefined)?.enabled === true || + (config.shadowRouting as { enabled?: boolean } | undefined)?.enabled === true || + (config.zeroLatencyOptimizationsEnabled === true && config.hedging === true) || + (resolvedTargets.length > 1 && + (pipeline || resolvedTargets.some((target) => Boolean(target.connectionId?.trim())))) || + nestedUnsafe + ); +} + +const managedComboRejection = () => + buildManagedLeaseErrorResponse( + new LeaseContextError(409, "LEASE_UNSUPPORTED_ROUTE", "Managed leases do not support this route") + ); + const comboPromoteDeps = { updateCombo, info: log.info, warn: log.warn }; export { shouldTripProviderBreakerForResult } from "./chatPredicates"; @@ -546,6 +598,19 @@ async function handleChatImplementation( return policy.rejection; } const apiKeyInfo = policy.apiKeyInfo; + let managedLease: ManagedLeaseDispatchContext | null = null; + if (isExclusiveLeaseManagedKey(apiKeyInfo)) { + try { + validateExclusiveLeaseKeyConfiguration(apiKeyInfo); + managedLease = { + apiKeyId: apiKeyInfo!.id, + context: parseManagedLeaseRequestContext(request.headers), + }; + } catch (error) { + if (error instanceof LeaseContextError) return buildManagedLeaseErrorResponse(error); + throw error; + } + } const bypassProviderQuotaPolicy = hasProviderQuotaBypassScope(apiKeyInfo?.scopes); telemetry.endPhase(); @@ -799,6 +864,12 @@ async function handleChatImplementation( if (filtered instanceof Response) return filtered; combo = filtered; } + const [settings, allCombos] = await Promise.all([ + getCachedSettings().catch(() => ({})), + getCombosCachedForChat(), + ]); + if (managedLease && isManagedComboUnsupported(combo, settings, allCombos)) + return managedComboRejection(); log.info( "CHAT", `Combo "${modelStr}" [${combo.strategy || "priority"}] with ${combo.models.length} models` @@ -885,24 +956,25 @@ async function handleChatImplementation( ...(target?.allowRateLimitedConnection ? { allowRateLimitedConnections: true } : {}), ...(target?.connectionId ? { forcedConnectionId: target.connectionId } : {}), ...(bypassProviderQuotaPolicy ? { bypassQuotaPolicy: true } : {}), + ...(managedLease ? { lease: credentialLease(managedLease) } : {}), } ); - if (!creds || !("authType" in creds)) return false; + if ( + !creds || + ("allRateLimited" in creds && creds.allRateLimited) || + ("waitingForCapacity" in creds && creds.waitingForCapacity) + ) + return false; // OAuth selection must happen atomically with occupancy reservation in the // actual dispatch. Availability preflight may finish well before a combo // target runs, so caching OAuth credentials here would reintroduce a race. - if (creds.authType !== "oauth") { + if ("authType" in creds && creds.authType !== "oauth") { comboPreselectedCredentials.set(getComboCredentialCacheKey(modelString, target), creds); } return true; }; - // Fetch settings and all combos for config cascade and nested resolution - const [settings, allCombos] = await Promise.all([ - getCachedSettings().catch(() => ({})), - getCombosCachedForChat(), - ]); const relayConfig = combo.strategy === "context-relay" ? resolveComboConfig(combo, settings) : null; // Per-request Auto-Combo controls (#6023 / #6024 / #6025 / #3470): steer an @@ -982,6 +1054,7 @@ async function handleChatImplementation( reasoningDecision, reasoningIntent, reasoningRequestTags: requestRoutingTags.tags, + managedLease, // #7360 follow-up: without this, a target dispatch abandoned by // targetTimeoutRunner.ts's per-target timeout (comboTargetTimeoutMs) // never learns it was abandoned — it only watches the ORIGINAL @@ -1048,6 +1121,7 @@ async function handleChatImplementation( sessionAffinityKey, emergencyFallbackTried: true, forceLiveComboTest: isComboLiveTest, + managedLease, }, combo.strategy, true @@ -1135,6 +1209,7 @@ async function handleChatImplementation( reasoningDecision, reasoningIntent, reasoningRequestTags: requestRoutingTags.tags, + managedLease, }, null, false @@ -1177,6 +1252,7 @@ async function handleSingleModelChat( reasoningDecision?: ReasoningRuleDecision | null; reasoningIntent?: ExtractedReasoningIntent | null; reasoningRequestTags?: string[]; + managedLease?: ManagedLeaseDispatchContext | null; /** * Per-target abort signal from combo.ts's targetTimeoutRunner * (comboTargetTimeoutMs) — see the #7360 follow-up comment at the @@ -1202,6 +1278,7 @@ async function handleSingleModelChat( // resolveModelOrError found a combo but the main handler's combo lookup missed it. if ((resolved as any).combo) { const redirectCombo = (resolved as any).combo; + if (runtimeOptions.managedLease) return managedComboRejection(); log.info( "ROUTING", `Safety-net combo redirect for "${modelStr}" → combo="${redirectCombo.name}"` @@ -1249,6 +1326,7 @@ async function handleSingleModelChat( allowRateLimitedConnection: resolvedTarget?.allowRateLimitedConnection === true, providerId: resolvedTarget?.providerId ?? null, correlationId: runtimeOptions?.correlationId ?? null, + managedLease: runtimeOptions.managedLease ?? null, // #7360 follow-up — see the primary handleSingleModel closure above. modelAbortSignal: target?.modelAbortSignal ?? null, }, @@ -1447,6 +1525,9 @@ async function handleSingleModelChat( ...(!forceLiveComboTest && bypassProviderQuotaPolicy ? { bypassQuotaPolicy: true } : {}), + ...(runtimeOptions.managedLease + ? { lease: credentialLease(runtimeOptions.managedLease) } + : {}), ...(() => { const effectiveForcedId = resolveForcedConnectionForCredentialPool({ forcedConnectionId: runtimeOptions.forcedConnectionId ?? null, @@ -1464,6 +1545,11 @@ async function handleSingleModelChat( ); preselectedCredentials = null; + if (runtimeOptions.managedLease && credentials) { + const leaseError = buildManagedLeaseSelectionErrorResponse(credentials); + if (leaseError) return leaseError; + } + // #9467: also treat the auth layer's allExpired verdict as a no-credentials // outcome (auth.ts produces it; without this check an all-expired pool fell // through to a connectionless dispatch). @@ -1697,6 +1783,7 @@ async function handleSingleModelChat( modelPinned: runtimeOptions?.modelPinned ?? false, routingComboId: runtimeOptions?.routingComboId ?? null, sessionAffinityKey: runtimeOptions.sessionAffinityKey ?? null, + managedLease: runtimeOptions.managedLease ?? null, }, runtimeOptions ); @@ -1751,6 +1838,16 @@ async function handleSingleModelChat( return result.response; } + // A final hard-lease fence rejection is authoritative. It must never mutate + // connection health/cooldown state or fall through to ordinary account/model + // fallback, which could turn a stale lifecycle into unmanaged dispatch. + if ( + runtimeOptions.managedLease && + (result.errorType === "lease_error" || String(result.errorCode || "").startsWith("LEASE_")) + ) { + return result.response; + } + // Missing Cloud Code project assignment is configuration, not a transient failure. // Preserve the typed fail-closed 422; marking it unavailable would trigger cooldown // redispatch and repeat bootstrap within the same logical request. diff --git a/src/sse/handlers/chat/clientRawRequest.ts b/src/sse/handlers/chat/clientRawRequest.ts index 0fd80cd9a8..f0dc47667d 100644 --- a/src/sse/handlers/chat/clientRawRequest.ts +++ b/src/sse/handlers/chat/clientRawRequest.ts @@ -13,6 +13,9 @@ import { cloneBoundedForLog } from "@omniroute/open-sse/utils/requestLogger.ts"; export function buildClientRawRequest(request: Request, body: unknown) { const url = new URL(request.url); + const headers = Object.fromEntries(request.headers.entries()); + delete headers["x-omniroute-lease-owner"]; + delete headers["x-omniroute-lease-generation"]; return { endpoint: url.pathname, // #7847: bounded, not a full deep clone. Every consumer of clientRawRequest.body is @@ -24,7 +27,7 @@ export function buildClientRawRequest(request: Request, body: unknown) { // Still a clone, not an alias — `body` is rewritten downstream (plugin onRequest hook, // compression), and this has to stay a snapshot of what the client actually sent. body: cloneBoundedForLog(body), - headers: Object.fromEntries(request.headers.entries()), + headers, signal: request.signal ?? null, }; } diff --git a/src/sse/handlers/chatHelpers.ts b/src/sse/handlers/chatHelpers.ts index c554cfe856..97bbea50db 100644 --- a/src/sse/handlers/chatHelpers.ts +++ b/src/sse/handlers/chatHelpers.ts @@ -422,6 +422,7 @@ export async function executeChatWithBreaker({ modelPinned = false, routingComboId = null, sessionAffinityKey = null, + managedLease = null, }: ExecuteChatWithBreakerOptions): Promise { let tlsFingerprintUsed = false; const normalizedTrafficType: TrafficType = @@ -478,6 +479,7 @@ export async function executeChatWithBreaker({ modelPinned, routingComboId, sessionAffinityKey, + managedLease, skipResourcePressureGuard: true, onCredentialsRefreshed: async (newCreds: any) => { await updateProviderCredentials(credentials.connectionId, { diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index d1f4c23da9..bf9076df92 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -14,6 +14,11 @@ import { clearConnectionErrorIfUnchanged, } from "@/lib/db/providers"; import { validateApiKey } from "@/lib/db/apiKeys"; +import { + getActiveExclusiveConnectionLease, + hashLeaseOwnerId, + type ExclusiveConnectionLease, +} from "@/lib/db/exclusiveConnectionLeases"; import { getSettings } from "@/lib/db/settings"; import { buildJinaEnvCredentials } from "@/lib/providers/jina"; import { buildGeminiEnvCredentials } from "@/lib/providers/gemini"; @@ -92,6 +97,7 @@ import { resolveForcedConnectionForCredentialPool, resolveSessionAffinityTtlMs, selectSessionAffinityConnection, + planSessionAffinityConnection, syncSessionAffinityRuntimeFields, } from "./sessionAffinityPin"; import { @@ -105,7 +111,17 @@ import { getResource404Bypass } from "./requestResourceHealth"; import { isVertexConnectionWidePermissionDenied } from "./vertexErrorClassifier"; import { maybeAutoDisableBannedAccount } from "./autoDisableBannedAccount"; import * as log from "../utils/logger"; -import { fisherYatesShuffle, getNextFromDeckSync } from "@/shared/utils/shuffleDeck"; +import { + fisherYatesShuffle, + getNextFromDeckSync, + planNextFromDeckSync, +} from "@/shared/utils/shuffleDeck"; +import { + applyExclusiveConnectionLeasePolicy, + invalidateManagedConnectionLease, + mutateExclusiveConnectionLease, + type CredentialLeaseSelectionContext, +} from "./exclusiveConnectionLeasePolicy"; import { readHeaderValue, type AuthRequestHeaders } from "./headerReader.ts"; import { getOAuthSessionAvailability, @@ -122,7 +138,7 @@ interface RecoverableConnectionState { lastErrorType?: string | null; lastErrorSource?: string | null; } -interface CredentialSelectionOptions { +export interface CredentialSelectionOptions { allowSuppressedConnections?: boolean; allowRateLimitedConnections?: boolean; bypassQuotaPolicy?: boolean; @@ -131,7 +147,19 @@ interface CredentialSelectionOptions { sessionKey?: string | null; sessionAffinityTtlMs?: number | null; reserveOAuthSession?: boolean; + lease?: CredentialLeaseSelectionContext; + materializeCredentials?: boolean; + deferLeaseClaim?: boolean; + /** Internal: a same-call UNIQUE retry already holds the provider/owner selection lock. */ + _leaseRetryWithLockHeld?: boolean; + /** Internal: freeze the original policy-valid candidate set across lease race/preflight retry. */ + _leaseCandidateIds?: string[]; } +export type ExclusiveLeaseSelectionResult = { + exclusiveLease: ExclusiveConnectionLease; + connectionId: string; + provider: string; +}; interface CooldownInspectionState { connection: ProviderConnectionView; connectionCooldownMs: number | null; @@ -933,6 +961,7 @@ function getSelectionMutexKey(provider: string, options: CredentialSelectionOpti return [ resolveProviderId(provider) || provider, options.forcedConnectionId ? `forced:${options.forcedConnectionId}` : "pool", + options.lease ? `lease:${hashLeaseOwnerId(options.lease.context.leaseOwnerId)}` : "unmanaged", ].join(":"); } function createSelectionLock(key: string) { @@ -1069,6 +1098,74 @@ async function getProviderSearchPool(provider: string): Promise { return Array.from(searchPool); } +function invalidateManagedLease( + options: CredentialSelectionOptions, + reason: Parameters[1] +) { + invalidateManagedConnectionLease(options.lease, reason); +} + +type DeferredLeaseSelection = { + commitSelectionSideEffects?: () => Promise | void; + selectNextLeaseCandidate?: (excludedConnectionId: string) => Promise; +}; + +function planLastUsedCommit( + connection: ProviderConnectionView, + connections: ProviderConnectionView[], + count: number +) { + const now = new Date().toISOString(); + return async () => { + await touchConnectionLastUsed(connection.id, count); + connection.lastUsedAt = now; + connection.consecutiveUseCount = count; + syncSessionAffinityRuntimeFields(connections, connection); + }; +} + +function materializeConnection( + connection: ProviderConnectionView, + options: CredentialSelectionOptions, + extra: DeferredLeaseSelection & { exclusiveLease?: ExclusiveConnectionLease } = {} +) { + const apiKeyHealth = connection.providerSpecificData?.apiKeyHealth as + Record | undefined; + if (apiKeyHealth) syncHealthFromDB(connection.id, apiKeyHealth); + const releaseOAuthSession = + options.reserveOAuthSession === true && connection.authType === "oauth" && options.sessionKey + ? reserveOAuthSession(connection.id, options.sessionKey) + : undefined; + return { + apiKey: connection.apiKey, + accessToken: connection.accessToken, + refreshToken: connection.refreshToken, + expiresAt: connection.tokenExpiresAt || connection.expiresAt || null, + projectId: connection.projectId, + defaultModel: connection.defaultModel || null, + copilotToken: + typeof connection.providerSpecificData.copilotToken === "string" + ? connection.providerSpecificData.copilotToken + : null, + providerSpecificData: connection.providerSpecificData, + id: connection.id, + provider: connection.provider, + authType: connection.authType, + email: connection.email, + connectionId: connection.id, + testStatus: connection.testStatus, + lastError: connection.lastError, + lastErrorType: connection.lastErrorType, + lastErrorSource: connection.lastErrorSource, + errorCode: connection.errorCode, + rateLimitedUntil: connection.rateLimitedUntil, + maxConcurrent: connection.maxConcurrent, + quotaWindowThresholds: connection.quotaWindowThresholds ?? null, + ...(releaseOAuthSession ? { releaseOAuthSession } : {}), + ...extra, + }; +} + /** * Get provider credentials from localDb * Filters out unavailable accounts and returns the selected account based on strategy @@ -1082,10 +1179,12 @@ export async function getProviderCredentials( requestedModel: string | null = null, options: CredentialSelectionOptions = {} ) { - const selectionLock = createSelectionLock(getSelectionMutexKey(provider, options)); + const selectionLock = options._leaseRetryWithLockHeld + ? null + : createSelectionLock(getSelectionMutexKey(provider, options)); try { - await selectionLock.wait; + await selectionLock?.wait; // No-auth providers (e.g. opencode) need no DB connection — return synthetic credentials // so the executor receives a valid credentials object without auth headers being added. @@ -1166,23 +1265,33 @@ export async function getProviderCredentials( if (allowedConnections && allowedConnections.length > 0) { connections = connections.filter((conn) => allowedConnections.includes(conn.id)); } + const forcedConnectionEligible = connections.some((conn) => conn.id === forcedConnectionId); + if (options.lease && forcedConnectionId && !forcedConnectionEligible) return null; + if (options.lease?.mode === "request" && forcedConnectionId) { + const activeLease = getActiveExclusiveConnectionLease(options.lease.context.leaseOwnerId); + if (activeLease && activeLease.connectionId !== forcedConnectionId) { + return { leaseConnectionMismatch: true }; + } + } // #5903: an active session-affinity pin outranks a per-request reset-aware // forcedConnectionId (see sessionAffinityPin leaf for the full rationale). - forcedConnectionId = - applySessionAffinityPin({ - forcedConnectionId, - options, - sessionAffinityTtlMs, - connections, - provider, - requestedModel, - excludedConnectionIds, - isTerminalConnectionStatus, - isCodexScopeUnavailable, - isQuotaPolicyBlocked: (c) => - evaluateQuotaLimitPolicy(provider, c as ProviderConnectionView, requestedModel).blocked, - }) ?? forcedConnectionId; + if (!options.lease) { + forcedConnectionId = + applySessionAffinityPin({ + forcedConnectionId, + options, + sessionAffinityTtlMs, + connections, + provider, + requestedModel, + excludedConnectionIds, + isTerminalConnectionStatus, + isCodexScopeUnavailable, + isQuotaPolicyBlocked: (c) => + evaluateQuotaLimitPolicy(provider, c as ProviderConnectionView, requestedModel).blocked, + }) ?? forcedConnectionId; + } forcedConnectionId = resolveForcedConnectionForCredentialPool({ forcedConnectionId, @@ -1247,6 +1356,7 @@ export async function getProviderCredentials( "AUTH", `${provider} | all ${allConnections.length} accounts rate limited (${formatRetryAfter(earliest)})` ); + invalidateManagedLease(options, "HEALTH_OR_COOLDOWN"); return { allRateLimited: true, retryAfter: earliest, @@ -1269,6 +1379,7 @@ export async function getProviderCredentials( // the dashboard sees a misleading "bad_request" code. const terminalConnections = allConnections.filter(isTerminalConnectionStatus); if (terminalConnections.length === allConnections.length) { + invalidateManagedLease(options, "AUTHORIZATION_CHANGED"); const syntheticFallback = await maybeSyntheticNoAuthFallback( resolvedId, excludedConnectionIds, @@ -1314,6 +1425,7 @@ export async function getProviderCredentials( log.info("AUTH", `${provider} | using ${geminiEnvCredentials.connectionId} env fallback`); return geminiEnvCredentials; } + invalidateManagedLease(options, "CONNECTION_INELIGIBLE"); log.warn("AUTH", `No credentials for ${provider}`); return null; } @@ -1513,6 +1625,10 @@ export async function getProviderCredentials( ? `${provider} | all ${connections.length} active accounts cooling down for model ${requestedModel} (${formatRetryAfter(earliest)}) | lastErrorCode=${earliestConn?.errorCode}, lastError=${earliestConn?.lastError?.slice(0, 50)}` : `${provider} | all ${connections.length} active accounts rate limited (${formatRetryAfter(earliest)}) | lastErrorCode=${earliestConn?.errorCode}, lastError=${earliestConn?.lastError?.slice(0, 50)}` ); + invalidateManagedLease( + options, + allBlockedByModelCooldown ? "MODEL_INELIGIBLE" : "HEALTH_OR_COOLDOWN" + ); return { allRateLimited: true, retryAfter: earliest, @@ -1530,6 +1646,7 @@ export async function getProviderCredentials( allowedConnections ); if (syntheticFallback) return syntheticFallback; + invalidateManagedLease(options, "CONNECTION_INELIGIBLE"); log.warn("AUTH", `${provider} | all ${connections.length} accounts unavailable`); return null; } @@ -1576,6 +1693,7 @@ export async function getProviderCredentials( ? new Date(earliestResetMs).toISOString() : new Date(Date.now() + 5 * 60 * 1000).toISOString(); + invalidateManagedLease(options, "QUOTA_UNAVAILABLE"); return { allRateLimited: true, retryAfter, @@ -1619,6 +1737,7 @@ export async function getProviderCredentials( ? new Date(earliestResetMs).toISOString() : new Date(Date.now() + 5 * 60 * 1000).toISOString(); + invalidateManagedLease(options, "QUOTA_UNAVAILABLE"); return { allRateLimited: true, retryAfter, @@ -1628,7 +1747,30 @@ export async function getProviderCredentials( }; } - const orderedConnections = [...withQuota].sort((a, b) => { + const policyValidLeaseCandidates = options._leaseCandidateIds + ? withQuota.filter((candidate) => options._leaseCandidateIds!.includes(candidate.id)) + : withQuota; + if (policyValidLeaseCandidates.length === 0) return null; + const leasePolicy = await applyExclusiveConnectionLeasePolicy( + policyValidLeaseCandidates, + options + ); + if (leasePolicy.error) return { [leasePolicy.error]: true }; + if (leasePolicy.connections.length === 0) { + if (options.lease?.mode === "request" && leasePolicy.activeLease) { + invalidateManagedLease(options, "CONNECTION_INELIGIBLE"); + } + return options.lease + ? { + waitingForCapacity: true, + retryAfter: leasePolicy.retryAfter, + eligibleCount: policyValidLeaseCandidates.length, + freeCount: 0, + } + : null; + } + + const orderedConnections = [...leasePolicy.connections].sort((a, b) => { if (a.authType !== "oauth" || b.authType !== "oauth") return 0; const priorityDelta = (a.priority || 999) - (b.priority || 999); if (priorityDelta !== 0) return priorityDelta; @@ -1645,16 +1787,35 @@ export async function getProviderCredentials( const providerOverride = providerStrategyOverrides[resolvedId] || {}; const strategy = providerOverride.fallbackStrategy || settings.fallbackStrategy || "fill-first"; - let connection; - const affinityConnection = await selectSessionAffinityConnection( - provider, - options.sessionKey, - orderedConnections, - sessionAffinityTtlMs - ); + let commitSelectionSideEffects: (() => Promise | void) | undefined; + let connection = leasePolicy.activeLease + ? orderedConnections.find( + (candidate) => candidate.id === leasePolicy.activeLease?.connectionId + ) + : undefined; + const affinityPlan = + options.lease && !connection + ? planSessionAffinityConnection( + provider, + options.sessionKey, + orderedConnections, + sessionAffinityTtlMs + ) + : null; + const affinityConnection = connection + ? connection + : options.lease + ? affinityPlan?.connection + : await selectSessionAffinityConnection( + provider, + options.sessionKey, + orderedConnections, + sessionAffinityTtlMs + ); if (affinityConnection) { connection = affinityConnection; - syncSessionAffinityRuntimeFields(connectionsRaw, connection); + if (options.lease) commitSelectionSideEffects = affinityPlan?.commit; + else syncSessionAffinityRuntimeFields(connectionsRaw, connection); } else if (options.sessionKey) { log.info( "AUTH", @@ -1696,15 +1857,9 @@ export async function getProviderCredentials( ); // Update lastUsedAt and increment count (await to ensure persistence) const nextCount = (connection.consecutiveUseCount || 0) + 1; - await touchConnectionLastUsed(connection.id, nextCount); - // Sync raw cache row so subsequent calls within TTL see fresh stats - for (const r of connectionsRaw as Record[]) { - if (r.id === connection.id) { - r.lastUsedAt = new Date().toISOString(); - r.consecutiveUseCount = nextCount; - break; - } - } + const commit = planLastUsedCommit(connection, connectionsRaw, nextCount); + if (options.lease) commitSelectionSideEffects = commit; + else await commit(); } else { // Pick the least recently used (excluding current if possible) // Also penalize accounts with high backoffLevel (previously rate-limited) @@ -1727,15 +1882,9 @@ export async function getProviderCredentials( ); // Update lastUsedAt and reset count to 1 (await to ensure persistence) - await touchConnectionLastUsed(connection.id, 1); - // Sync raw cache row so subsequent calls within TTL see fresh LRU stats - for (const r of connectionsRaw as Record[]) { - if (r.id === connection.id) { - r.lastUsedAt = new Date().toISOString(); - r.consecutiveUseCount = 1; - break; - } - } + const commit = planLastUsedCommit(connection, connectionsRaw, 1); + if (options.lease) commitSelectionSideEffects = commit; + else await commit(); } } else { // Fallback scenario: excluded an account due to failure @@ -1758,18 +1907,12 @@ export async function getProviderCredentials( ); // Update lastUsedAt and reset count to 1 (await to ensure persistence) - await touchConnectionLastUsed(connection.id, 1); - // Sync raw cache row so subsequent calls within TTL see fresh stats - for (const r of connectionsRaw as Record[]) { - if (r.id === connection.id) { - r.lastUsedAt = new Date().toISOString(); - r.consecutiveUseCount = 1; - break; - } - } + const commit = planLastUsedCommit(connection, connectionsRaw, 1); + if (options.lease) commitSelectionSideEffects = commit; + else await commit(); } } else if (strategy === "p2c") { - const candidatePool = withQuota.length > 0 ? withQuota : orderedConnections; + const candidatePool = orderedConnections; // Power of Two Choices: sample from the quota-eligible pool and compare // health instead of defaulting to random-first selection. if (candidatePool.length <= 2) { @@ -1811,8 +1954,15 @@ export async function getProviderCredentials( } else if (strategy === "strict-random") { // Strict Random: shuffle deck — uses each account once before reshuffling const ids = orderedConnections.map((c) => c.id); - const selectedId = getNextFromDeckSync(`conn:${provider}`, ids); - connection = orderedConnections.find((c) => c.id === selectedId) || orderedConnections[0]; + if (options.lease) { + const plan = planNextFromDeckSync(`conn:${provider}`, ids); + connection = + orderedConnections.find((c) => c.id === plan.selectedId) || orderedConnections[0]; + commitSelectionSideEffects = plan.commit; + } else { + const selectedId = getNextFromDeckSync(`conn:${provider}`, ids); + connection = orderedConnections.find((c) => c.id === selectedId) || orderedConnections[0]; + } } else { // Default: fill-first (already sorted by priority in getProviderConnections) connection = orderedConnections[0]; @@ -1838,6 +1988,43 @@ export async function getProviderCredentials( if (moreAvailablePeer) connection = moreAvailablePeer; } + let exclusiveLease: ExclusiveConnectionLease | undefined; + if (options.lease) { + const candidateIds = orderedConnections.map((candidate) => candidate.id); + const selectNextLeaseCandidate = (excludedConnectionId: string) => + getProviderCredentials(provider, null, allowedConnections, requestedModel, { + ...options, + excludeConnectionIds: [...excludedConnectionIds, excludedConnectionId], + deferLeaseClaim: true, + _leaseCandidateIds: candidateIds, + }); + if (options.deferLeaseClaim) { + return materializeConnection(connection, options, { + commitSelectionSideEffects, + selectNextLeaseCandidate, + }); + } + let claim = mutateExclusiveConnectionLease( + connection, + leasePolicy.activeLease, + options.lease + ); + if (claim.kind === "LOST") { + return getProviderCredentials(provider, null, allowedConnections, requestedModel, { + ...options, + excludeConnectionIds: [...excludedConnectionIds, connection.id], + _leaseCandidateIds: candidateIds, + _leaseRetryWithLockHeld: true, + }); + } + if (claim.kind === "STALE") return { leaseFenceStale: true }; + exclusiveLease = claim.lease; + await commitSelectionSideEffects?.(); + if (options.materializeCredentials === false) { + return { exclusiveLease, connectionId: connection.id, provider: connection.provider }; + } + } + if (provider === "antigravity" && connection) { log.info( "AUTH", @@ -1845,57 +2032,9 @@ export async function getProviderCredentials( ); } - const apiKeyHealth = connection.providerSpecificData?.apiKeyHealth as - Record | undefined; - if (apiKeyHealth) { - syncHealthFromDB(connection.id, apiKeyHealth); - } - - const releaseOAuthSession = - options.reserveOAuthSession === true && connection.authType === "oauth" && options.sessionKey - ? reserveOAuthSession(connection.id, options.sessionKey) - : undefined; - - return { - apiKey: connection.apiKey, - accessToken: connection.accessToken, - refreshToken: connection.refreshToken, - expiresAt: connection.tokenExpiresAt || connection.expiresAt || null, - projectId: connection.projectId, - // #474: surface the connection's configured defaultModel so the chat / - // embeddings handlers can resolve a bare model name (e.g. an alias that - // resolved to "auto") to a real provider model ID before the upstream call. - defaultModel: connection.defaultModel || null, - copilotToken: - typeof connection.providerSpecificData.copilotToken === "string" - ? connection.providerSpecificData.copilotToken - : null, - providerSpecificData: connection.providerSpecificData, - // Fields the generic quota fetcher (open-sse/services/genericQuotaFetcher.ts) - // needs to delegate to getUsageForProvider for any provider — kept aliased - // (`id` + `connectionId`) for back-compat with callers that already use the - // connectionId name. - id: connection.id, - provider: connection.provider, - authType: connection.authType, - email: connection.email, - connectionId: connection.id, - // Include current status for optimization check - testStatus: connection.testStatus, - lastError: connection.lastError, - lastErrorType: connection.lastErrorType, - lastErrorSource: connection.lastErrorSource, - errorCode: connection.errorCode, - rateLimitedUntil: connection.rateLimitedUntil, - maxConcurrent: connection.maxConcurrent, - // Surface per-window quota overrides so the preflight latency gate in - // getProviderCredentialsWithQuotaPreflight can see them. Without this, - // user-set cutoffs would silently never enforce. - quotaWindowThresholds: connection.quotaWindowThresholds ?? null, - ...(releaseOAuthSession ? { releaseOAuthSession } : {}), - }; + return materializeConnection(connection, options, { exclusiveLease }); } finally { - selectionLock.release(); + selectionLock?.release(); } } export async function getProviderCredentialsWithQuotaPreflight( @@ -1946,18 +2085,17 @@ export async function getProviderCredentialsWithQuotaPreflight( // tighter floor is honored. const FACTORY_NO_OP_REMAINING_PERCENT = 2; const globalDefaultIsRestrictive = defaultThresholdPercent > FACTORY_NO_OP_REMAINING_PERCENT; + let pendingCredentialSelection: Awaited> | undefined; while (true) { - const credentials = await getProviderCredentials( - provider, - null, - allowedConnections, - requestedModel, - { + const credentials = + pendingCredentialSelection ?? + (await getProviderCredentials(provider, null, allowedConnections, requestedModel, { ...options, excludeConnectionIds: Array.from(excludedConnectionIds), - } - ); + ...(options.lease ? { deferLeaseClaim: true } : {}), + })); + pendingCredentialSelection = undefined; if (!credentials) { if (blockedByPreflight.length > 0) { @@ -1980,14 +2118,42 @@ export async function getProviderCredentialsWithQuotaPreflight( return credentials; } - const selectedCredentials = credentials as typeof credentials & { + const selectedCredentials = credentials as Omit< + typeof credentials, + "selectNextLeaseCandidate" + > & { connectionId?: string; + commitSelectionSideEffects?: () => Promise | void; + selectNextLeaseCandidate?: (excludedConnectionId: string) => Promise; releaseOAuthSession?: () => void; }; const connectionId = selectedCredentials.connectionId; if (!connectionId) { return credentials; } + const commitLease = async () => { + if (!options.lease) return credentials; + const activeLease = getActiveExclusiveConnectionLease(options.lease.context.leaseOwnerId); + const claim = mutateExclusiveConnectionLease( + selectedCredentials as unknown as ProviderConnectionView, + activeLease, + options.lease + ); + if (claim.kind === "LOST") { + selectedCredentials.releaseOAuthSession?.(); + excludedConnectionIds.add(connectionId); + pendingCredentialSelection = + await selectedCredentials.selectNextLeaseCandidate?.(connectionId); + return null; + } + if (claim.kind === "STALE") return { leaseFenceStale: true }; + await selectedCredentials.commitSelectionSideEffects?.(); + if (options.materializeCredentials === false) { + selectedCredentials.releaseOAuthSession?.(); + return { exclusiveLease: claim.lease, connectionId, provider }; + } + return { ...credentials, exclusiveLease: claim.lease }; + }; // Cascading resolver: per-connection override → per-(provider, window) // default → global default. Used per-window when the fetcher exposes @@ -2017,7 +2183,11 @@ export async function getProviderCredentialsWithQuotaPreflight( const legacyForceDisable = (credentials as { providerSpecificData?: Record }).providerSpecificData ?.quotaPreflightEnabled === false; - if (legacyForceDisable) return credentials; + if (legacyForceDisable) { + const committed = await commitLease(); + if (committed === null) continue; + return committed; + } const hasConnectionOverrides = Object.keys(perConnectionWindowOverrides).length > 0; const legacyForceEnable = isQuotaPreflightEnabled(credentials as Record); @@ -2027,7 +2197,9 @@ export async function getProviderCredentialsWithQuotaPreflight( !legacyForceEnable && !globalDefaultIsRestrictive ) { - return credentials; + const committed = await commitLease(); + if (committed === null) continue; + return committed; } // Returns the minimum-remaining cutoff for a window — matches the @@ -2070,7 +2242,9 @@ export async function getProviderCredentialsWithQuotaPreflight( throw error; } if (preflight.proceed) { - return credentials; + const committed = await commitLease(); + if (committed === null) continue; + return committed; } selectedCredentials.releaseOAuthSession?.(); @@ -2087,6 +2261,7 @@ export async function getProviderCredentialsWithQuotaPreflight( resetAt: unavailableUntil, }); excludedConnectionIds.add(connectionId); + pendingCredentialSelection = await selectedCredentials.selectNextLeaseCandidate?.(connectionId); log.info( "AUTH", diff --git a/src/sse/services/exclusiveConnectionLeasePolicy.ts b/src/sse/services/exclusiveConnectionLeasePolicy.ts new file mode 100644 index 0000000000..c56e5256da --- /dev/null +++ b/src/sse/services/exclusiveConnectionLeasePolicy.ts @@ -0,0 +1,125 @@ +import { getExclusiveLeaseConnectionIds } from "@/lib/db/apiKeys"; +import { + acquireExclusiveConnectionLease, + getActiveExclusiveConnectionLease, + getExclusiveLeaseOccupancy, + hashLeaseOwnerId, + invalidateExclusiveConnectionLease, + transitionExclusiveConnectionLease, + type ExclusiveConnectionLease, + type ExclusiveLeaseEndReason, +} from "@/lib/db/exclusiveConnectionLeases"; +import type { ProviderConnectionView } from "@/lib/db/providers/lazyConnectionView"; + +import type { ManagedLeaseRequestContext } from "./leaseContext"; + +export interface CredentialLeaseSelectionContext { + apiKeyId: string; + context: ManagedLeaseRequestContext; + mode: "acquire" | "request"; +} + +export interface LeaseSelectionOptions { + forcedConnectionId?: string | null; + lease?: CredentialLeaseSelectionContext; +} + +export type LeaseCandidatePolicy = { + connections: ProviderConnectionView[]; + activeLease: ExclusiveConnectionLease | null; + retryAfter: string | null; + error?: "leaseConnectionMismatch" | "leaseFenceStale" | "leaseRequired"; +}; + +export async function applyExclusiveConnectionLeasePolicy( + connections: ProviderConnectionView[], + options: LeaseSelectionOptions +): Promise { + const occupancy = getExclusiveLeaseOccupancy(connections.map((connection) => connection.id)); + if (!options.lease) { + const managed = await getExclusiveLeaseConnectionIds(); + return { + connections: connections.filter( + (connection) => !managed.has(connection.id) && !occupancy.has(connection.id) + ), + activeLease: null, + retryAfter: null, + }; + } + + const { lease } = options; + const activeLease = getActiveExclusiveConnectionLease(lease.context.leaseOwnerId); + if (lease.mode === "request" && !activeLease) { + return { connections: [], activeLease: null, retryAfter: null, error: "leaseRequired" }; + } + if (lease.mode === "request" && activeLease?.generation !== lease.context.generation) { + return { connections: [], activeLease, retryAfter: null, error: "leaseFenceStale" }; + } + if (activeLease && activeLease.apiKeyId !== lease.apiKeyId) { + return { connections: [], activeLease, retryAfter: null, error: "leaseFenceStale" }; + } + const activeBinding = activeLease + ? connections.find((connection) => connection.id === activeLease.connectionId) + : undefined; + if (!activeBinding && lease.mode === "request" && activeLease && options.forcedConnectionId) { + return { connections: [], activeLease, retryAfter: null, error: "leaseConnectionMismatch" }; + } + const ownerHash = hashLeaseOwnerId(lease.context.leaseOwnerId); + const free = connections.filter((connection) => + occupancy.get(connection.id)?.leaseOwnerHash !== ownerHash && occupancy.has(connection.id) + ? false + : connection.id !== activeBinding?.id + ); + return { + connections: activeBinding ? [activeBinding, ...free] : free, + activeLease, + retryAfter: + [...occupancy.values()] + .filter((row) => row.leaseOwnerHash !== ownerHash) + .map((row) => row.expiresAt) + .sort()[0] ?? null, + }; +} + +export function mutateExclusiveConnectionLease( + connection: ProviderConnectionView, + activeLease: ExclusiveConnectionLease | null, + lease: CredentialLeaseSelectionContext, + reason: ExclusiveLeaseEndReason = "CONNECTION_INELIGIBLE" +) { + const result = activeLease + ? transitionExclusiveConnectionLease({ + leaseOwnerId: lease.context.leaseOwnerId, + generation: lease.mode === "acquire" ? activeLease.generation : lease.context.generation, + apiKeyId: lease.apiKeyId, + provider: connection.provider, + connectionId: connection.id, + reason, + }) + : acquireExclusiveConnectionLease({ + leaseOwnerId: lease.context.leaseOwnerId, + apiKeyId: lease.apiKeyId, + provider: connection.provider, + connectionId: connection.id, + }); + if (result.kind === "CONNECTION_BUSY") { + return { kind: "LOST" as const, retryAfter: result.retryAfter }; + } + if (result.kind === "STALE" || result.kind === "OWNER_ALREADY_ACTIVE") { + return { kind: "STALE" as const }; + } + return { kind: "CLAIMED" as const, lease: result.lease }; +} + +export function invalidateManagedConnectionLease( + lease: CredentialLeaseSelectionContext | undefined, + reason: ExclusiveLeaseEndReason +): void { + if (lease?.mode !== "request") return; + invalidateExclusiveConnectionLease({ + leaseOwnerId: lease.context.leaseOwnerId, + generation: lease.context.generation, + apiKeyId: lease.apiKeyId, + reason, + }); +} diff --git a/src/sse/services/leaseContext.ts b/src/sse/services/leaseContext.ts new file mode 100644 index 0000000000..9b5370c133 --- /dev/null +++ b/src/sse/services/leaseContext.ts @@ -0,0 +1,136 @@ +import { LEASE_OWNER_PATTERN } from "@/lib/db/exclusiveConnectionLeases"; +import { buildErrorBody } from "@omniroute/open-sse/utils/error.ts"; + +export const LEASE_EXCLUSIVE_SCOPE = "lease:exclusive", + LEASE_OWNER_HEADER = "X-OmniRoute-Lease-Owner", + LEASE_GENERATION_HEADER = "X-OmniRoute-Lease-Generation"; + +export type ManagedLeaseRequestContext = { leaseOwnerId: string; generation: number }; +export type ManagedLeaseDispatchContext = { + apiKeyId: string; + context: ManagedLeaseRequestContext; +}; +export const credentialLease = (lease: ManagedLeaseDispatchContext) => ({ + ...lease, + mode: "request" as const, +}); + +export class LeaseContextError extends Error { + constructor( + readonly status: number, + readonly code: string, + message: string + ) { + super(message); + } +} + +type LeaseKeyMetadata = { + scopes?: readonly string[] | null; + allowedConnections?: readonly string[] | null; +}; + +export function isExclusiveLeaseManagedKey(metadata: LeaseKeyMetadata | null | undefined): boolean { + return Array.isArray(metadata?.scopes) && metadata.scopes.includes(LEASE_EXCLUSIVE_SCOPE); +} + +export function validateExclusiveLeaseKeyConfiguration( + metadata: LeaseKeyMetadata | null | undefined +): void { + if (!isExclusiveLeaseManagedKey(metadata)) return; + if (!metadata?.allowedConnections?.length) + throw new LeaseContextError( + 403, + "LEASE_KEY_CONFIGURATION_INVALID", + "Exclusive lease keys require explicit allowed connections" + ); +} + +export function parseLeaseOwnerHeader(headers: Headers): string { + const leaseOwnerId = headers.get(LEASE_OWNER_HEADER)?.trim() ?? ""; + if (!leaseOwnerId) + throw new LeaseContextError(400, "LEASE_CONTEXT_REQUIRED", "Explicit lease owner required"); + if (!LEASE_OWNER_PATTERN.test(leaseOwnerId)) + throw new LeaseContextError(400, "LEASE_CONTEXT_INVALID", "Malformed lease owner"); + return leaseOwnerId; +} + +export function parseManagedLeaseRequestContext(headers: Headers): ManagedLeaseRequestContext { + const leaseOwnerId = parseLeaseOwnerHeader(headers); + const rawGeneration = headers.get(LEASE_GENERATION_HEADER)?.trim() ?? ""; + if (!/^[1-9]\d*$/.test(rawGeneration)) + throw new LeaseContextError(400, "LEASE_CONTEXT_INVALID", "Positive lease generation required"); + const generation = Number(rawGeneration); + if (!Number.isSafeInteger(generation)) + throw new LeaseContextError( + 400, + "LEASE_CONTEXT_INVALID", + "The lease generation header is invalid" + ); + return { leaseOwnerId, generation }; +} + +export function buildManagedLeaseErrorResponse(error: LeaseContextError): Response { + return new Response( + JSON.stringify( + buildErrorBody(error.status, error.message, undefined, { + type: "lease_error", + code: error.code, + }) + ), + { status: error.status, headers: { "Content-Type": "application/json" } } + ); +} + +type ManagedLeaseSelectionFailure = { + eligibleCount?: number; + freeCount?: number; + leaseConnectionMismatch?: boolean; + leaseFenceStale?: boolean; + leaseRequired?: boolean; + retryAfter?: string | null; + waitingForCapacity?: boolean; + allRateLimited?: boolean; +}; + +export function buildManagedLeaseSelectionErrorResponse( + selection: ManagedLeaseSelectionFailure +): Response | null { + const code = selection.leaseRequired + ? "LEASE_REQUIRED" + : selection.leaseFenceStale + ? "LEASE_FENCE_STALE" + : selection.leaseConnectionMismatch + ? "LEASE_CONNECTION_MISMATCH" + : null; + if (code) + return buildManagedLeaseErrorResponse( + new LeaseContextError(409, code, code.replaceAll("_", " ")) + ); + if (!selection.waitingForCapacity) return null; + const expiryMs = Date.parse(selection.retryAfter ?? ""); + const retryAfterSeconds = Number.isFinite(expiryMs) + ? Math.min(3600, Math.max(1, Math.ceil((expiryMs - Date.now()) / 1000))) + : 1; + return new Response( + JSON.stringify({ + state: "WAITING_FOR_CAPACITY", + error: { + type: "lease_error", + code: "LEASE_CAPACITY_UNAVAILABLE", + message: "Exclusive managed session capacity is temporarily unavailable", + }, + reason: "NO_FREE_ELIGIBLE_CONNECTION", + retryAfter: retryAfterSeconds, + eligibleCount: Math.max(0, selection.eligibleCount ?? 0), + freeCount: Math.max(0, selection.freeCount ?? 0), + }), + { + status: 429, + headers: { + "Content-Type": "application/json", + "Retry-After": String(retryAfterSeconds), + }, + } + ); +} diff --git a/src/sse/services/sessionAffinityPin.ts b/src/sse/services/sessionAffinityPin.ts index a2055e0fa5..f7c5009df6 100644 --- a/src/sse/services/sessionAffinityPin.ts +++ b/src/sse/services/sessionAffinityPin.ts @@ -141,6 +141,43 @@ export async function selectSessionAffinityConnection( + provider: string, + sessionKey: string | null | undefined, + connections: T[], + ttlMs = 0 +) { + if (!sessionKey || connections.length === 0 || ttlMs <= 0) return null; + const existing = getSessionAccountAffinity(sessionKey, provider, ttlMs); + const existingConnection = + existing && connections.find((candidate) => candidate.id === existing.connectionId); + const connection = existingConnection ?? [...connections].sort(compareLruConnections)[0] ?? null; + if (!connection) return null; + + return { + connection, + commit: async () => { + if (existingConnection) { + touchSessionAccountAffinity(sessionKey, provider, Date.now(), ttlMs); + const nextCount = (connection.consecutiveUseCount || 0) + 1; + await touchConnectionLastUsed(connection.id, nextCount); + connection.lastUsedAt = new Date().toISOString(); + connection.consecutiveUseCount = nextCount; + return; + } + if (existing) deleteSessionAccountAffinity(sessionKey, provider); + upsertSessionAccountAffinity(sessionKey, provider, connection.id, Date.now(), ttlMs); + await touchConnectionLastUsed(connection.id, 1); + connection.lastUsedAt = new Date().toISOString(); + connection.consecutiveUseCount = 1; + }, + }; +} + /** Inputs the combo-timeout eviction needs from the dispatch site. */ export interface ComboTimeoutAffinityEvictionParams { sessionKey?: string | null; diff --git a/tests/unit/api-key-scope-validation.test.ts b/tests/unit/api-key-scope-validation.test.ts index 7ce6e37db0..8b59d1e62e 100644 --- a/tests/unit/api-key-scope-validation.test.ts +++ b/tests/unit/api-key-scope-validation.test.ts @@ -12,7 +12,10 @@ import { hasSelfUsageScope, normalizeSelfServiceScopesForCreate, } from "../../src/shared/constants/selfServiceScopes.ts"; -import { createKeySchema, updateKeyPermissionsSchema } from "../../src/shared/validation/schemas.ts"; +import { + createKeySchema, + updateKeyPermissionsSchema, +} from "../../src/shared/validation/schemas.ts"; const repoRoot = path.resolve(path.dirname(fileURLToPath(import.meta.url)), "../.."); @@ -44,6 +47,30 @@ test("api key validation accepts more than sixteen scopes", () => { assert.equal(updateKeyPermissionsSchema.safeParse({ scopes }).success, true); }); +test("lease scope requires an explicit non-empty connection allowlist", () => { + const connection = "00000000-0000-4000-8000-000000000001"; + assert.equal( + createKeySchema.safeParse({ name: "invalid managed key", scopes: ["lease:exclusive"] }).success, + false + ); + // Partial PATCH validity depends on the authoritative stored-row + mutation check. + assert.equal(updateKeyPermissionsSchema.safeParse({ scopes: ["lease:exclusive"] }).success, true); + assert.equal( + updateKeyPermissionsSchema.safeParse({ + scopes: ["lease:exclusive"], + allowedConnections: [], + }).success, + false + ); + assert.equal( + updateKeyPermissionsSchema.safeParse({ + scopes: ["lease:exclusive"], + allowedConnections: [connection], + }).success, + true + ); +}); + test("api key create route normalizes omitted scopes to self-service usage", () => { const source = fs.readFileSync(path.join(repoRoot, "src/app/api/keys/route.ts"), "utf8"); diff --git a/tests/unit/chat-managed-lease-routing.test.ts b/tests/unit/chat-managed-lease-routing.test.ts new file mode 100644 index 0000000000..2511d89fff --- /dev/null +++ b/tests/unit/chat-managed-lease-routing.test.ts @@ -0,0 +1,635 @@ +import assert from "node:assert/strict"; +import test from "node:test"; + +import { createChatPipelineHarness } from "../integration/_chatPipelineHarness.ts"; + +const harness = await createChatPipelineHarness("chat-managed-lease-routing"); +const { + apiKeysDb, + buildOpenAIResponse, + buildRequest, + combosDb, + handleChat, + resetStorage, + seedConnection, +} = harness; +const leaseDb = await import("../../src/lib/db/exclusiveConnectionLeases.ts"); +const providersDb = await import("../../src/lib/db/providers.ts"); +const accountSemaphores = await import("../../open-sse/services/accountSemaphore.ts"); +const { POST: handleCompletions } = await import("../../src/app/api/v1/completions/route.ts"); + +const OWNER = "vlo_AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"; +const OWNER_B = "vlo_BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB"; + +async function seedManagedKey(connectionIds: string[]) { + return apiKeysDb.createApiKey("managed-chat", "test", ["lease:exclusive"], { + allowedConnections: connectionIds, + }); +} + +function managedRequest( + key: string, + generation: number, + extraHeaders = {}, + bodyOverrides: Record = {}, + owner = OWNER, + url = "http://localhost/v1/chat/completions" +) { + return buildRequest({ + url, + authKey: key, + headers: { + "X-OmniRoute-Lease-Owner": owner, + "X-OmniRoute-Lease-Generation": String(generation), + ...extraHeaders, + }, + body: { + model: "openai/gpt-4.1", + stream: false, + messages: [{ role: "user", content: "synthetic managed lease test" }], + ...bodyOverrides, + }, + }); +} + +function buildOpenAIStreamResponse(text: string): Response { + const frames = [ + `data: ${JSON.stringify({ + id: "chatcmpl_stream", + object: "chat.completion.chunk", + choices: [{ index: 0, delta: { role: "assistant", content: text }, finish_reason: null }], + })}\n\n`, + `data: ${JSON.stringify({ + id: "chatcmpl_stream", + object: "chat.completion.chunk", + choices: [{ index: 0, delta: {}, finish_reason: "stop" }], + })}\n\n`, + "data: [DONE]\n\n", + ]; + return new Response(frames.join(""), { + status: 200, + headers: { "Content-Type": "text/event-stream" }, + }); +} + +test.beforeEach(async () => { + process.env.REQUIRE_API_KEY = "false"; + await resetStorage(); +}); +test.after(async () => harness.cleanup()); + +test("managed chat requires explicit owner and generation before provider dispatch", async () => { + const connection = await seedConnection("openai"); + const key = await seedManagedKey([connection.id]); + let dispatches = 0; + globalThis.fetch = async () => { + dispatches += 1; + throw new Error("unexpected provider dispatch"); + }; + + const missingOwner = await handleChat( + buildRequest({ + authKey: key.key, + body: { + model: "openai/gpt-4.1", + stream: false, + messages: [{ role: "user", content: "missing owner" }], + }, + }) + ); + assert.equal(missingOwner.status, 400); + assert.equal((await missingOwner.json()).error.code, "LEASE_CONTEXT_REQUIRED"); + + const missingGeneration = await handleChat( + buildRequest({ + authKey: key.key, + headers: { "X-OmniRoute-Lease-Owner": OWNER }, + body: { + model: "openai/gpt-4.1", + stream: false, + messages: [{ role: "user", content: "missing generation" }], + }, + }) + ); + assert.equal(missingGeneration.status, 400); + assert.equal((await missingGeneration.json()).error.code, "LEASE_CONTEXT_INVALID"); + assert.equal(dispatches, 0); +}); + +test("managed chat blocks missing and stale leases with zero provider dispatch", async () => { + const connection = await seedConnection("openai"); + const key = await seedManagedKey([connection.id]); + let dispatches = 0; + globalThis.fetch = async () => { + dispatches += 1; + throw new Error("unexpected provider dispatch"); + }; + + const missing = await handleChat(managedRequest(key.key, 1)); + assert.equal(missing.status, 409); + assert.equal((await missing.json()).error.code, "LEASE_REQUIRED"); + + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER, + apiKeyId: key.id, + provider: "openai", + connectionId: connection.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + + const stale = await handleChat(managedRequest(key.key, acquired.lease.generation + 1)); + assert.equal(stale.status, 409); + assert.equal((await stale.json()).error.code, "LEASE_FENCE_STALE"); + assert.equal(dispatches, 0); +}); + +test("managed chat blocks cross-key owner-generation replay before provider dispatch", async () => { + const connection = await seedConnection("openai"); + const ownerKey = await seedManagedKey([connection.id]); + const replayKey = await seedManagedKey([connection.id]); + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER, + apiKeyId: ownerKey.id, + provider: "openai", + connectionId: connection.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + + let dispatches = 0; + globalThis.fetch = async () => { + dispatches += 1; + throw new Error("unexpected provider dispatch"); + }; + const replay = await handleChat(managedRequest(replayKey.key, acquired.lease.generation)); + + assert.equal(replay.status, 409); + assert.equal((await replay.json()).error.code, "LEASE_FENCE_STALE"); + assert.equal(dispatches, 0); + assert.equal(leaseDb.getActiveExclusiveConnectionLease(OWNER)?.apiKeyId, ownerKey.id); +}); + +test("managed chat dispatches only the fenced active binding", async () => { + const connection = await seedConnection("openai"); + const key = await seedManagedKey([connection.id]); + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER, + apiKeyId: key.id, + provider: "openai", + connectionId: connection.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + + let dispatches = 0; + globalThis.fetch = async () => { + dispatches += 1; + return buildOpenAIResponse("managed success"); + }; + const response = await handleChat(managedRequest(key.key, acquired.lease.generation)); + assert.equal(response.status, 200); + assert.equal((await response.json()).choices[0].message.content, "managed success"); + assert.equal(dispatches, 1); +}); + +test("identical prompts with different owners never share a managed connection", async () => { + const firstConnection = await seedConnection("openai", { + name: "managed-owner-a", + apiKey: "sk-managed-owner-a", + priority: 1, + }); + const secondConnection = await seedConnection("openai", { + name: "managed-owner-b", + apiKey: "sk-managed-owner-b", + priority: 2, + }); + const key = await seedManagedKey([firstConnection.id, secondConnection.id]); + const firstLease = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER, + apiKeyId: key.id, + provider: "openai", + connectionId: firstConnection.id, + }); + const secondLease = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER_B, + apiKeyId: key.id, + provider: "openai", + connectionId: secondConnection.id, + }); + assert.equal(firstLease.kind, "ACQUIRED"); + assert.equal(secondLease.kind, "ACQUIRED"); + if (firstLease.kind !== "ACQUIRED" || secondLease.kind !== "ACQUIRED") return; + + const usedApiKeys: string[] = []; + globalThis.fetch = async (_url, init) => { + usedApiKeys.push(new Headers(init?.headers).get("authorization") ?? ""); + return buildOpenAIResponse("isolated owner success"); + }; + const sharedBody = { + model: "openai/gpt-4.1", + stream: false, + messages: [{ role: "user", content: "byte-identical prompt and tools" }], + tools: [{ type: "function", function: { name: "noop", parameters: { type: "object" } } }], + }; + const firstResponse = await handleChat( + managedRequest( + key.key, + firstLease.lease.generation, + { "X-Session-Id": "same-routing-session" }, + sharedBody, + OWNER + ) + ); + const secondResponse = await handleChat( + managedRequest( + key.key, + secondLease.lease.generation, + { "X-Session-Id": "same-routing-session" }, + sharedBody, + OWNER_B + ) + ); + + assert.equal(firstResponse.status, 200); + assert.equal(secondResponse.status, 200); + assert.equal(usedApiKeys.length, 2); + assert.notEqual(usedApiKeys[0], usedApiKeys[1]); + assert.notEqual( + leaseDb.getActiveExclusiveConnectionLease(OWNER)?.connectionId, + leaseDb.getActiveExclusiveConnectionLease(OWNER_B)?.connectionId + ); +}); + +test("changing prompt, tools, and request model does not change the owner binding", async () => { + const connection = await seedConnection("openai"); + const key = await seedManagedKey([connection.id]); + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER, + apiKeyId: key.id, + provider: "openai", + connectionId: connection.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + + let dispatches = 0; + globalThis.fetch = async () => { + dispatches += 1; + return buildOpenAIResponse("identity stable"); + }; + for (const body of [ + { messages: [{ role: "user", content: "prompt one" }] }, + { + messages: [{ role: "user", content: "prompt two" }], + tools: [{ type: "function", function: { name: "other", parameters: { type: "object" } } }], + }, + { model: "openai/gpt-4o-mini", messages: [{ role: "user", content: "model changed" }] }, + ]) { + const response = await handleChat(managedRequest(key.key, acquired.lease.generation, {}, body)); + assert.equal(response.status, 200); + assert.equal(leaseDb.getActiveExclusiveConnectionLease(OWNER)?.connectionId, connection.id); + } + assert.equal(dispatches, 3); +}); + +test("legacy completions and messages-compatible paths use the managed lease handler", async () => { + const connection = await seedConnection("openai"); + const key = await seedManagedKey([connection.id]); + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER, + apiKeyId: key.id, + provider: "openai", + connectionId: connection.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + + let dispatches = 0; + globalThis.fetch = async () => { + dispatches += 1; + return buildOpenAIResponse("legacy fenced"); + }; + const completions = await handleCompletions( + managedRequest( + key.key, + acquired.lease.generation, + {}, + { prompt: "legacy completion", messages: undefined }, + OWNER, + "http://localhost/v1/completions" + ) + ); + const messages = await handleChat( + managedRequest( + key.key, + acquired.lease.generation, + {}, + { messages: [{ role: "user", content: "messages compatible" }] }, + OWNER, + "http://localhost/v1/messages" + ) + ); + + assert.equal(completions.status, 200); + assert.equal(messages.status, 200); + assert.equal(dispatches, 2); + assert.equal(leaseDb.getActiveExclusiveConnectionLease(OWNER)?.connectionId, connection.id); +}); + +test("managed chat fences after an admission wait and before main executor dispatch", async () => { + const connection = await seedConnection("openai"); + await providersDb.updateProviderConnection(connection.id, { maxConcurrent: 1 }); + const key = await seedManagedKey([connection.id]); + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER, + apiKeyId: key.id, + provider: "openai", + connectionId: connection.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + + const semaphoreKey = accountSemaphores.buildAccountSemaphoreKey({ + provider: "openai", + accountKey: connection.id, + }); + const releaseBlocker = await accountSemaphores.acquire(semaphoreKey, { maxConcurrency: 1 }); + let dispatches = 0; + globalThis.fetch = async () => { + dispatches += 1; + throw new Error("unexpected provider dispatch after stale fence"); + }; + + const pending = handleChat(managedRequest(key.key, acquired.lease.generation)); + for (let i = 0; i < 40; i += 1) { + if ((accountSemaphores.getStats()[semaphoreKey]?.queued ?? 0) === 1) break; + await new Promise((resolve) => setTimeout(resolve, 5)); + } + assert.equal(accountSemaphores.getStats()[semaphoreKey]?.queued, 1); + + assert.equal( + leaseDb.releaseExclusiveConnectionLease({ + leaseOwnerId: OWNER, + generation: acquired.lease.generation, + apiKeyId: key.id, + }).kind, + "RELEASED" + ); + releaseBlocker(); + const response = await pending; + + assert.equal(response.status, 409); + assert.equal((await response.json()).error.code, "LEASE_REQUIRED"); + assert.equal(dispatches, 0); +}); + +test("managed streaming chat preserves the lifecycle lease after completion", async () => { + const connection = await seedConnection("openai"); + const key = await seedManagedKey([connection.id]); + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER, + apiKeyId: key.id, + provider: "openai", + connectionId: connection.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + + let dispatches = 0; + globalThis.fetch = async () => { + dispatches += 1; + return buildOpenAIStreamResponse("managed stream success"); + }; + const response = await handleChat( + managedRequest( + key.key, + acquired.lease.generation, + { Accept: "text/event-stream" }, + { + stream: true, + } + ) + ); + const body = await response.text(); + + assert.equal(response.status, 200); + assert.match(response.headers.get("Content-Type") || "", /text\/event-stream/); + assert.match(body, /managed stream success/); + assert.equal(dispatches, 1); + assert.equal( + leaseDb.getActiveExclusiveConnectionLease(OWNER)?.generation, + acquired.lease.generation + ); +}); + +test("managed Responses-shaped request uses the same fenced chat path", async () => { + const connection = await seedConnection("openai"); + const key = await seedManagedKey([connection.id]); + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER, + apiKeyId: key.id, + provider: "openai", + connectionId: connection.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + + let dispatches = 0; + globalThis.fetch = async () => { + dispatches += 1; + return buildOpenAIResponse("responses fenced success"); + }; + const response = await handleChat( + buildRequest({ + url: "http://localhost/v1/responses", + authKey: key.key, + headers: { + "X-OmniRoute-Lease-Owner": OWNER, + "X-OmniRoute-Lease-Generation": String(acquired.lease.generation), + }, + body: { + model: "openai/gpt-4.1", + stream: false, + input: "synthetic Responses request", + }, + }) + ); + + assert.equal(response.status, 200); + assert.equal(dispatches, 1); + assert.equal(leaseDb.getActiveExclusiveConnectionLease(OWNER)?.connectionId, connection.id); +}); + +test("a direct foreign connection pin cannot override the active binding", async () => { + const bound = await seedConnection("openai", { name: "managed-bound", priority: 1 }); + const other = await seedConnection("openai", { name: "managed-other", priority: 2 }); + const key = await seedManagedKey([bound.id, other.id]); + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER, + apiKeyId: key.id, + provider: "openai", + connectionId: bound.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + let dispatches = 0; + globalThis.fetch = async () => { + dispatches += 1; + throw new Error("unexpected provider dispatch"); + }; + + assert.notEqual(bound.id, other.id); + assert.equal(leaseDb.getActiveExclusiveConnectionLease(OWNER)?.connectionId, bound.id); + + const pinnedRequest = managedRequest(key.key, acquired.lease.generation, { + "X-OmniRoute-Connection": other.id, + }); + assert.equal(pinnedRequest.headers.get("x-omniroute-connection"), other.id); + const response = await handleChat(pinnedRequest); + assert.equal(response.status, 409); + assert.equal((await response.json()).error.code, "LEASE_CONNECTION_MISMATCH"); + assert.equal(dispatches, 0); + assert.equal((await providersDb.getProviderConnectionById(bound.id))?.testStatus, "active"); +}); + +test("managed chat retains ordinary cooldown semantics instead of reporting lease capacity", async () => { + const connection = await seedConnection("openai", { + rateLimitedUntil: new Date(Date.now() + 60_000).toISOString(), + }); + const key = await seedManagedKey([connection.id]); + let dispatches = 0; + globalThis.fetch = async () => { + dispatches += 1; + throw new Error("unexpected provider dispatch"); + }; + + const response = await handleChat(managedRequest(key.key, 1)); + assert.notEqual(response.status, 429); + assert.notEqual((await response.json()).state, "WAITING_FOR_CAPACITY"); + assert.equal(dispatches, 0); +}); + +test("empty ordinary eligibility is not reported as lease capacity contention", async () => { + const connection = await seedConnection("openai"); + const key = await seedManagedKey([connection.id]); + await providersDb.updateProviderConnection(connection.id, { testStatus: "banned" }); + let dispatches = 0; + globalThis.fetch = async () => { + dispatches += 1; + throw new Error("unexpected provider dispatch"); + }; + + const response = await handleChat(managedRequest(key.key, 1)); + const body = await response.json(); + assert.notEqual(body.state, "WAITING_FOR_CAPACITY"); + assert.notEqual(body.error?.code, "LEASE_CAPACITY_UNAVAILABLE"); + assert.equal(dispatches, 0); +}); + +test("managed combos reject every fan-out route before provider dispatch", async () => { + const firstConnection = await seedConnection("openai", { + name: "managed-combo-first", + apiKey: "sk-managed-combo-first", + }); + const secondConnection = await seedConnection("openai", { + name: "managed-combo-second", + apiKey: "sk-managed-combo-second", + }); + const key = await seedManagedKey([firstConnection.id, secondConnection.id]); + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER, + apiKeyId: key.id, + provider: "openai", + connectionId: firstConnection.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + + const cases = [ + { name: "managed-fusion", strategy: "fusion", models: ["openai/gpt-4.1"] }, + { name: "managed-relay", strategy: "context-relay", models: ["openai/gpt-4.1"] }, + { + name: "managed-pipeline-multi", + strategy: "pipeline", + models: ["openai/gpt-4.1", "openai/gpt-4o-mini"], + }, + { + name: "managed-chaos", + strategy: "priority", + config: { chaos: { enabled: true } }, + models: ["openai/gpt-4.1", "openai/gpt-4o-mini"], + }, + { + name: "managed-shadow", + strategy: "priority", + config: { shadowRouting: { enabled: true, targets: ["openai/gpt-4o-mini"] } }, + models: ["openai/gpt-4.1"], + }, + { + name: "managed-speculative", + strategy: "priority", + config: { zeroLatencyOptimizationsEnabled: true, hedging: true }, + models: ["openai/gpt-4.1", "openai/gpt-4o-mini"], + }, + { + name: "managed-fixed-multi-account", + strategy: "priority", + models: [ + { model: "openai/gpt-4.1", connectionId: firstConnection.id }, + { model: "openai/gpt-4o-mini", connectionId: secondConnection.id }, + ], + }, + ]; + for (const combo of cases) await combosDb.createCombo(combo); + await combosDb.createCombo({ + name: "managed-nested-fusion", + strategy: "priority", + config: { nestedComboMode: "execute" }, + models: [{ kind: "combo-ref", comboName: "managed-fusion" }], + }); + + let dispatches = 0; + globalThis.fetch = async () => { + dispatches += 1; + throw new Error("unexpected provider dispatch"); + }; + for (const model of [...cases.map((combo) => combo.name), "managed-nested-fusion"]) { + const response = await handleChat( + managedRequest(key.key, acquired.lease.generation, {}, { model }) + ); + assert.equal(response.status, 409, model); + assert.equal((await response.json()).error.code, "LEASE_UNSUPPORTED_ROUTE", model); + } + assert.equal(dispatches, 0); +}); + +test("one-step managed pipeline uses the ordinary fenced lease path", async () => { + const connection = await seedConnection("openai"); + const key = await seedManagedKey([connection.id]); + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER, + apiKeyId: key.id, + provider: "openai", + connectionId: connection.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + await combosDb.createCombo({ + name: "managed-pipeline-one", + strategy: "pipeline", + config: { maxRetries: 0 }, + models: ["openai/gpt-4.1"], + }); + let dispatches = 0; + globalThis.fetch = async () => { + dispatches += 1; + return buildOpenAIResponse("one-step success"); + }; + + const response = await handleChat( + managedRequest(key.key, acquired.lease.generation, {}, { model: "managed-pipeline-one" }) + ); + assert.equal(response.status, 200); + assert.equal(dispatches, 1); +}); diff --git a/tests/unit/chatcore-executor-client-headers.test.ts b/tests/unit/chatcore-executor-client-headers.test.ts index 92ff2a8c5b..b6f553aafc 100644 --- a/tests/unit/chatcore-executor-client-headers.test.ts +++ b/tests/unit/chatcore-executor-client-headers.test.ts @@ -40,3 +40,13 @@ test("does not overwrite an existing user-agent header", () => { test("a trimmed-empty user agent does not create headers on its own", () => { assert.equal(buildExecutorClientHeaders({}, " "), null); }); + +test("internal hard-lease control headers never reach an executor", () => { + const out = buildExecutorClientHeaders({ + "X-OmniRoute-Lease-Owner": `vlo_${"A".repeat(43)}`, + "x-omniroute-lease-generation": "7", + "x-session-id": "routing-session-remains-independent", + }); + + assert.deepEqual(out, { "x-session-id": "routing-session-remains-independent" }); +}); diff --git a/tests/unit/chatcore-translation-paths.test.ts b/tests/unit/chatcore-translation-paths.test.ts index c82d6987b2..f28caccf09 100644 --- a/tests/unit/chatcore-translation-paths.test.ts +++ b/tests/unit/chatcore-translation-paths.test.ts @@ -309,6 +309,8 @@ async function invokeChatCore({ onCredentialsRefreshed = null, onRequestSuccess = null, sessionAffinityKey = null, + managedLease = null, + cachedSettings = null, }: any = {}) { const calls: any[] = []; @@ -355,6 +357,8 @@ async function invokeChatCore({ sessionAffinityKey, isCombo, comboStrategy, + managedLease, + cachedSettings, onCredentialsRefreshed, onRequestSuccess, } as any); diff --git a/tests/unit/codex-ws-policy-enforcement-6564.test.ts b/tests/unit/codex-ws-policy-enforcement-6564.test.ts index 3214a0c274..7801e27a9f 100644 --- a/tests/unit/codex-ws-policy-enforcement-6564.test.ts +++ b/tests/unit/codex-ws-policy-enforcement-6564.test.ts @@ -154,6 +154,25 @@ test("WS prepare() allows the requested model when the key's policy permits it ( assert.equal(body.error?.code, "codex_credentials_unavailable"); }); +test("WS prepare() rejects managed lease keys before credential selection", async () => { + const managedKey = await apiKeysDb.createApiKey( + "Managed Lease WS Key", + "machine-lease-ws", + ["lease:exclusive"], + { allowedConnections: ["synthetic-managed-connection"] } + ); + await apiKeysDb.updateApiKeyPermissions(managedKey.id, { + allowedModels: ["gpt-5.5"], + }); + + const response = await route.POST(buildPrepareRequest(managedKey.key, "gpt-5.5")); + const body = (await response.json()) as ErrorBody; + + assert.equal(response.status, 409); + assert.equal(body.error.code, "LEASE_UNSUPPORTED_TRANSPORT"); + assert.notEqual(body.error.code, "codex_credentials_unavailable"); +}); + test("WS prepare() rejects a combo not in the key's allowedCombos policy (403)", async () => { await combosDb.createCombo({ name: "model-1.0", diff --git a/tests/unit/cors/origins.test.ts b/tests/unit/cors/origins.test.ts index 8fe8e78fd2..00950aa097 100644 --- a/tests/unit/cors/origins.test.ts +++ b/tests/unit/cors/origins.test.ts @@ -284,4 +284,10 @@ describe("cors/origins.STATIC_CORS_HEADERS", () => { ); assert.match(STATIC_CORS_HEADERS["Access-Control-Allow-Methods"], /OPTIONS/); }); + + it("allows the generic managed-lease control headers", () => { + const allowedHeaders = STATIC_CORS_HEADERS["Access-Control-Allow-Headers"]; + assert.match(allowedHeaders, /X-OmniRoute-Lease-Owner/i); + assert.match(allowedHeaders, /X-OmniRoute-Lease-Generation/i); + }); }); diff --git a/tests/unit/exclusive-connection-leases.test.ts b/tests/unit/exclusive-connection-leases.test.ts new file mode 100644 index 0000000000..6488fdea5e --- /dev/null +++ b/tests/unit/exclusive-connection-leases.test.ts @@ -0,0 +1,455 @@ +import assert from "node:assert/strict"; +import { spawn } from "node:child_process"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import test from "node:test"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-hard-lease-v2-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../src/lib/db/core.ts"); +const leases = await import("../../src/lib/db/exclusiveConnectionLeases.ts"); + +const OWNER_A = "vlo_AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"; +const OWNER_B = "vlo_BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB"; + +function at(seconds: number): string { + return new Date(Date.UTC(2026, 7, 12, 19, 30, seconds)).toISOString(); +} + +test.after(() => { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("hashes canonical owners and never persists the raw owner", () => { + assert.match(leases.hashLeaseOwnerId(OWNER_A), /^[a-f0-9]{64}$/); + assert.throws(() => leases.hashLeaseOwnerId("routing-session"), /canonical/); + + const acquired = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER_A, + apiKeyId: "key-a", + provider: "codex", + connectionId: "conn-a", + now: at(0), + ttlMs: 120_000, + }); + + assert.equal(acquired.kind, "ACQUIRED"); + const db = core.getDbInstance(); + const row = db.prepare("SELECT lease_owner_hash FROM exclusive_connection_leases").get() as { + lease_owner_hash: string; + }; + assert.match(row.lease_owner_hash, /^[a-f0-9]{64}$/); + assert.notEqual(row.lease_owner_hash, OWNER_A); + const rawDb = fs.readFileSync(path.join(TEST_DATA_DIR, "storage.sqlite")); + assert.equal(rawDb.includes(Buffer.from(OWNER_A)), false); +}); + +test("uses the live next-free migration slot without runner compatibility special cases", () => { + const migration = fs.readFileSync( + new URL("../../src/lib/db/migrations/155_exclusive_connection_leases.sql", import.meta.url), + "utf8" + ); + const runner = fs.readFileSync( + new URL("../../src/lib/db/migrationRunner.ts", import.meta.url), + "utf8" + ); + assert.match(migration, /CREATE TABLE IF NOT EXISTS exclusive_connection_leases/); + assert.doesNotMatch(runner, /case "155"/); +}); + +test("enforces global active owner and connection uniqueness", () => { + const ownerA = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER_A, + apiKeyId: "key-a", + provider: "codex", + connectionId: "conn-a", + now: at(1), + }); + assert.equal(ownerA.kind, "REUSED"); + + const sameOwnerOtherKey = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER_A, + apiKeyId: "key-b", + provider: "openai", + connectionId: "conn-b", + now: at(2), + }); + assert.equal(sameOwnerOtherKey.kind, "OWNER_ALREADY_ACTIVE"); + + const foreign = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER_B, + apiKeyId: "key-b", + provider: "codex", + connectionId: "conn-a", + now: at(3), + }); + assert.equal(foreign.kind, "CONNECTION_BUSY"); +}); + +test("renews and releases only an exact generation and release is idempotent", () => { + const acquired = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER_B, + apiKeyId: "key-b", + provider: "codex", + connectionId: "conn-b", + now: at(4), + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + + assert.equal( + leases.renewExclusiveConnectionLease({ + leaseOwnerId: OWNER_B, + generation: acquired.lease.generation + 1, + apiKeyId: "key-b", + now: at(5), + }).kind, + "STALE" + ); + assert.equal( + leases.renewExclusiveConnectionLease({ + leaseOwnerId: OWNER_B, + generation: acquired.lease.generation, + apiKeyId: "key-b", + now: at(6), + }).kind, + "RENEWED" + ); + assert.equal( + leases.releaseExclusiveConnectionLease({ + leaseOwnerId: OWNER_B, + generation: acquired.lease.generation - 1, + apiKeyId: "key-b", + now: at(7), + }).kind, + "STALE" + ); + assert.equal( + leases.releaseExclusiveConnectionLease({ + leaseOwnerId: OWNER_B, + generation: acquired.lease.generation, + apiKeyId: "key-b", + now: at(8), + }).kind, + "RELEASED" + ); + assert.equal( + leases.releaseExclusiveConnectionLease({ + leaseOwnerId: OWNER_B, + generation: acquired.lease.generation, + apiKeyId: "key-b", + now: at(9), + }).kind, + "RELEASED" + ); +}); + +test("keeps generation on failover and fences stale requests", () => { + const first = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER_B, + apiKeyId: "key-b", + provider: "codex", + connectionId: "conn-b", + now: at(10), + }); + assert.equal(first.kind, "ACQUIRED"); + if (first.kind !== "ACQUIRED") return; + + const transitioned = leases.transitionExclusiveConnectionLease({ + leaseOwnerId: OWNER_B, + generation: first.lease.generation, + apiKeyId: "key-b", + provider: "codex", + connectionId: "conn-c", + now: at(11), + reason: "CONNECTION_INELIGIBLE", + }); + assert.equal(transitioned.kind, "TRANSITIONED"); + if (transitioned.kind !== "TRANSITIONED") return; + assert.equal(transitioned.lease.generation, first.lease.generation); + assert.equal( + leases.assertExclusiveConnectionLeaseFence({ + leaseOwnerId: OWNER_B, + generation: first.lease.generation, + apiKeyId: "key-b", + connectionId: "conn-b", + now: at(12), + }).kind, + "CONNECTION_MISMATCH" + ); + assert.equal( + leases.assertExclusiveConnectionLeaseFence({ + leaseOwnerId: OWNER_B, + generation: first.lease.generation, + apiKeyId: "key-b", + connectionId: "conn-c", + now: at(12), + }).kind, + "VALID" + ); + assert.equal( + leases.assertExclusiveConnectionLeaseFence({ + leaseOwnerId: OWNER_B, + generation: first.lease.generation, + apiKeyId: "key-foreign", + connectionId: "conn-c", + now: at(12), + }).kind, + "AUTHORIZATION_MISMATCH" + ); +}); + +test("invalidates an unsafe binding only for the exact generation", () => { + const owner = "vlo_DDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDD"; + const acquired = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: owner, + apiKeyId: "key-d", + provider: "codex", + connectionId: "conn-invalid", + now: at(13), + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + + assert.equal( + leases.invalidateExclusiveConnectionLease({ + leaseOwnerId: owner, + generation: acquired.lease.generation + 1, + apiKeyId: "key-d", + reason: "QUOTA_UNAVAILABLE", + now: at(14), + }).kind, + "STALE" + ); + const invalidated = leases.invalidateExclusiveConnectionLease({ + leaseOwnerId: owner, + generation: acquired.lease.generation, + apiKeyId: "key-d", + reason: "QUOTA_UNAVAILABLE", + now: at(15), + }); + assert.equal(invalidated.kind, "INVALIDATED"); + assert.equal(leases.getActiveExclusiveConnectionLease(owner, at(16)), null); +}); + +test("expires lazily, reconciles at restart, and increments generation", () => { + const owner = "vlo_CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCC"; + const first = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: owner, + apiKeyId: "key-c", + provider: "codex", + connectionId: "conn-expire", + now: at(13), + ttlMs: 1_000, + }); + assert.equal(first.kind, "ACQUIRED"); + if (first.kind !== "ACQUIRED") return; + + assert.equal(leases.reconcileExpiredExclusiveConnectionLeases(at(15)), 1); + const second = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: owner, + apiKeyId: "key-c", + provider: "codex", + connectionId: "conn-expire", + now: at(16), + }); + assert.equal(second.kind, "ACQUIRED"); + if (second.kind !== "ACQUIRED") return; + assert.equal(second.lease.generation, first.lease.generation + 1); +}); + +test("zero-request heartbeat holds through idle and restart until bounded TTL recovery", () => { + const owner = "vlo_GGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGG"; + const contender = "vlo_HHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHH"; + const acquired = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: owner, + apiKeyId: "key-g", + provider: "codex", + connectionId: "conn-idle-heartbeat", + now: "2026-08-12T19:32:00.000Z", + ttlMs: 1_000, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + + const renewed = leases.renewExclusiveConnectionLease({ + leaseOwnerId: owner, + generation: acquired.lease.generation, + apiKeyId: "key-g", + now: "2026-08-12T19:32:00.500Z", + ttlMs: 1_000, + }); + assert.equal(renewed.kind, "RENEWED"); + assert.equal( + leases.acquireExclusiveConnectionLease({ + leaseOwnerId: contender, + apiKeyId: "key-h", + provider: "codex", + connectionId: "conn-idle-heartbeat", + now: "2026-08-12T19:32:01.000Z", + }).kind, + "CONNECTION_BUSY" + ); + + core.resetDbInstance(); + assert.equal( + leases.getActiveExclusiveConnectionLease(owner, "2026-08-12T19:32:01.250Z")?.generation, + acquired.lease.generation + ); + assert.equal(leases.reconcileExpiredExclusiveConnectionLeases("2026-08-12T19:32:01.500Z"), 1); + assert.equal( + leases.acquireExclusiveConnectionLease({ + leaseOwnerId: contender, + apiKeyId: "key-h", + provider: "codex", + connectionId: "conn-idle-heartbeat", + now: "2026-08-12T19:32:01.501Z", + }).kind, + "ACQUIRED" + ); +}); + +test("generation remains monotonic after release and invalidation", () => { + const owner = "vlo_IIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIII"; + const first = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: owner, + apiKeyId: "key-i", + provider: "codex", + connectionId: "conn-generation-a", + now: "2026-08-12T19:33:00.000Z", + }); + assert.equal(first.kind, "ACQUIRED"); + if (first.kind !== "ACQUIRED") return; + assert.equal( + leases.releaseExclusiveConnectionLease({ + leaseOwnerId: owner, + generation: first.lease.generation, + apiKeyId: "key-i", + now: "2026-08-12T19:33:01.000Z", + }).kind, + "RELEASED" + ); + + const second = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: owner, + apiKeyId: "key-i", + provider: "codex", + connectionId: "conn-generation-b", + now: "2026-08-12T19:33:02.000Z", + }); + assert.equal(second.kind, "ACQUIRED"); + if (second.kind !== "ACQUIRED") return; + assert.equal(second.lease.generation, first.lease.generation + 1); + assert.equal( + leases.invalidateExclusiveConnectionLease({ + leaseOwnerId: owner, + generation: second.lease.generation, + apiKeyId: "key-i", + reason: "CONNECTION_INELIGIBLE", + now: "2026-08-12T19:33:03.000Z", + }).kind, + "INVALIDATED" + ); + + const third = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: owner, + apiKeyId: "key-i", + provider: "codex", + connectionId: "conn-generation-c", + now: "2026-08-12T19:33:04.000Z", + }); + assert.equal(third.kind, "ACQUIRED"); + if (third.kind !== "ACQUIRED") return; + assert.equal(third.lease.generation, second.lease.generation + 1); +}); + +test("migration exposes exactly global ACTIVE uniqueness indexes", () => { + const db = core.getDbInstance(); + const ownerIndex = db + .prepare("SELECT sql FROM sqlite_master WHERE name = 'idx_exclusive_lease_active_owner'") + .get() as { sql: string }; + const connectionIndex = db + .prepare("SELECT sql FROM sqlite_master WHERE name = 'idx_exclusive_lease_active_connection'") + .get() as { sql: string }; + + assert.match(ownerIndex.sql, /UNIQUE INDEX[\s\S]*\(lease_owner_hash\)[\s\S]*state = 'ACTIVE'/i); + assert.doesNotMatch(ownerIndex.sql, /api_key_id|provider/i); + assert.match(connectionIndex.sql, /UNIQUE INDEX[\s\S]*\(connection_id\)[\s\S]*state = 'ACTIVE'/i); + assert.doesNotMatch(connectionIndex.sql, /api_key_id|provider/i); +}); + +test("cross-process contenders never both acquire the same connection", async () => { + const raceDir = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-hard-lease-race-")); + const barrier = path.join(raceDir, "go"); + const coreUrl = new URL("../../src/lib/db/core.ts", import.meta.url).href; + const moduleUrl = new URL("../../src/lib/db/exclusiveConnectionLeases.ts", import.meta.url).href; + const worker = [ + `import fs from "node:fs";`, + `process.env.DATA_DIR = process.env.LEASE_RACE_DIR;`, + `const leases = await import(${JSON.stringify(moduleUrl)});`, + `while (!fs.existsSync(process.env.LEASE_RACE_BARRIER)) { await new Promise((r) => setTimeout(r, 2)); }`, + `const result = leases.acquireExclusiveConnectionLease({`, + ` leaseOwnerId: process.env.LEASE_RACE_OWNER,`, + ` apiKeyId: process.env.LEASE_RACE_KEY,`, + ` provider: "codex",`, + ` connectionId: "conn-process-race",`, + ` now: "2026-08-12T19:31:00.000Z",`, + `});`, + `process.stdout.write(JSON.stringify({ kind: result.kind, generation: result.lease?.generation }));`, + ].join("\n"); + + function runChild(script: string, env: Record): Promise { + return new Promise((resolve, reject) => { + const child = spawn( + process.execPath, + ["--import", "tsx/esm", "--input-type=module", "-e", script], + { + cwd: process.cwd(), + env: { ...process.env, ...env }, + stdio: ["ignore", "pipe", "pipe"], + } + ); + let stdout = ""; + let stderr = ""; + child.stdout.setEncoding("utf8").on("data", (chunk) => (stdout += chunk)); + child.stderr.setEncoding("utf8").on("data", (chunk) => (stderr += chunk)); + child.once("error", reject); + child.once("exit", (code) => { + if (code !== 0) return reject(new Error(`race worker failed (${code}): ${stderr}`)); + resolve(stdout); + }); + }); + } + + function contender(owner: string, key: string): Promise<{ kind: string; generation?: number }> { + return runChild(worker, { + LEASE_RACE_DIR: raceDir, + LEASE_RACE_BARRIER: barrier, + LEASE_RACE_OWNER: owner, + LEASE_RACE_KEY: key, + }).then((stdout) => { + const start = stdout.lastIndexOf("{"); + return JSON.parse(stdout.slice(start)) as { kind: string; generation?: number }; + }); + } + + try { + await runChild( + `process.env.DATA_DIR = process.env.LEASE_RACE_DIR; const core = await import(${JSON.stringify(coreUrl)}); core.getDbInstance();`, + { LEASE_RACE_DIR: raceDir } + ); + const contenders = [ + contender("vlo_EEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEE", "key-e"), + contender("vlo_FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF", "key-f"), + ]; + fs.writeFileSync(barrier, "go", { flag: "wx" }); + const results = await Promise.all(contenders); + assert.equal(results.filter((result) => result.kind === "ACQUIRED").length, 1); + assert.equal(results.filter((result) => result.kind === "CONNECTION_BUSY").length, 1); + } finally { + fs.rmSync(raceDir, { recursive: true, force: true }); + } +}); diff --git a/tests/unit/exclusive-lease-api-key-policy.test.ts b/tests/unit/exclusive-lease-api-key-policy.test.ts new file mode 100644 index 0000000000..dfd6d4dc4b --- /dev/null +++ b/tests/unit/exclusive-lease-api-key-policy.test.ts @@ -0,0 +1,105 @@ +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import test from "node:test"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-lease-key-policy-")); +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.API_KEY_SECRET = "exclusive-lease-key-policy-test-secret"; +process.env.DISABLE_SQLITE_AUTO_BACKUP = "true"; + +const core = await import("../../src/lib/db/core.ts"); +const apiKeys = await import("../../src/lib/db/apiKeys.ts"); +const keysRoute = await import("../../src/app/api/keys/[id]/route.ts"); +const CONNECTION = "00000000-0000-4000-8000-000000000001"; + +async function resetStorage(): Promise { + core.resetDbInstance(); + apiKeys.resetApiKeyState(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +test.beforeEach(resetStorage); +test.after(() => { + core.resetDbInstance(); + apiKeys.resetApiKeyState(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("managed API key create requires and atomically stores an explicit allowlist", async () => { + await assert.rejects( + apiKeys.createApiKey("invalid managed", "test", ["lease:exclusive"]), + /requires explicit allowedConnections/ + ); + const created = await apiKeys.createApiKey("valid managed", "test", ["lease:exclusive"], { + allowedConnections: [CONNECTION], + }); + const row = core + .getDbInstance() + .prepare("SELECT scopes, allowed_connections FROM api_keys WHERE id = ?") + .get(created.id) as { scopes: string; allowed_connections: string }; + assert.deepEqual(JSON.parse(row.scopes), ["lease:exclusive"]); + assert.deepEqual(JSON.parse(row.allowed_connections), [CONNECTION]); +}); + +test("partial domain mutations enforce the stored-row plus requested-policy invariant", async () => { + const ordinary = await apiKeys.createApiKey("ordinary", "test"); + await assert.rejects( + apiKeys.updateApiKeyPermissions(ordinary.id, { scopes: ["lease:exclusive"] }), + /requires explicit allowedConnections/ + ); + + const managed = await apiKeys.createApiKey("managed", "test", ["lease:exclusive"], { + allowedConnections: [CONNECTION], + }); + await assert.rejects( + apiKeys.updateApiKeyPermissions(managed.id, { allowedConnections: [] }), + /requires explicit allowedConnections/ + ); + assert.equal( + await apiKeys.updateApiKeyPermissions(managed.id, { + scopes: [], + allowedConnections: [], + }), + true + ); +}); + +test("unrelated permission updates retain the ordinary non-transactional path", async () => { + const ordinary = await apiKeys.createApiKey("ordinary update", "test"); + const db = core.getDbInstance(); + const originalExec = db.exec.bind(db); + let beginImmediateCalls = 0; + db.exec = ((sql: string) => { + if (sql === "BEGIN IMMEDIATE") beginImmediateCalls += 1; + return originalExec(sql); + }) as typeof db.exec; + try { + assert.equal(await apiKeys.updateApiKeyPermissions(ordinary.id, { name: "renamed" }), true); + assert.equal(beginImmediateCalls, 0); + } finally { + db.exec = originalExec; + } +}); + +test("partial management PATCH maps the domain invariant to a sanitized 400", async () => { + const ordinary = await apiKeys.createApiKey("ordinary route key", "test"); + const management = await apiKeys.createApiKey("management route key", "test", ["manage"]); + const response = await keysRoute.PATCH( + new Request(`http://omniroute.local/api/keys/${ordinary.id}`, { + method: "PATCH", + headers: { + Authorization: `Bearer ${management.key}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ scopes: ["lease:exclusive"] }), + }), + { params: Promise.resolve({ id: ordinary.id }) } + ); + assert.equal(response.status, 400); + const body = (await response.json()) as { error: { code: string; message: string } }; + assert.equal(body.error.code, "LEASE_KEY_POLICY_INVALID"); + assert.equal(body.error.message, "lease:exclusive requires explicit allowedConnections"); +}); diff --git a/tests/unit/exclusive-lease-auxiliary-isolation.test.ts b/tests/unit/exclusive-lease-auxiliary-isolation.test.ts new file mode 100644 index 0000000000..da6174648c --- /dev/null +++ b/tests/unit/exclusive-lease-auxiliary-isolation.test.ts @@ -0,0 +1,192 @@ +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import test from "node:test"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-lease-auxiliary-")); +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.DISABLE_SQLITE_AUTO_BACKUP = "true"; +process.env.API_KEY_SECRET = "exclusive-lease-auxiliary-test-secret"; + +let externalCalls = 0; +const originalFetch = globalThis.fetch; +globalThis.fetch = async () => { + externalCalls += 1; + throw new Error("unexpected external provider/model call"); +}; + +const core = await import("../../src/lib/db/core.ts"); +const providers = await import("../../src/lib/db/providers.ts"); +const apiKeys = await import("../../src/lib/db/apiKeys.ts"); +const leases = await import("../../src/lib/db/exclusiveConnectionLeases.ts"); +const translator = await import("../../src/app/api/translator/send/route.ts"); +const translatorPreview = await import("../../src/app/api/translator/translate/route.ts"); +const modelTests = await import("../../src/lib/api/modelTestRunner.ts"); +const vnc = await import("../../src/lib/vncSession/service.ts"); + +const OWNER = "vlo_UUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUUU"; + +async function seedConnection(name: string): Promise<{ id: string }> { + return (await providers.createProviderConnection({ + provider: "openai", + authType: "apikey", + name, + apiKey: `sk-${name}`, + isActive: true, + testStatus: "active", + priority: 1, + providerSpecificData: {}, + })) as { id: string }; +} + +async function markLeaseOnly(connectionId: string): Promise { + await apiKeys.createApiKey("managed auxiliary", "test", ["lease:exclusive"], { + allowedConnections: [connectionId], + }); +} + +async function resetStorage(): Promise { + core.resetDbInstance(); + apiKeys.resetApiKeyState(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); + externalCalls = 0; +} + +test.beforeEach(resetStorage); +test.after(() => { + globalThis.fetch = originalFetch; + core.resetDbInstance(); + apiKeys.resetApiKeyState(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("translator send excludes a FREE lease-only connection before provider fetch", async () => { + const connection = await seedConnection("translator-lease-only"); + await markLeaseOnly(connection.id); + + const response = await translator.POST( + new Request("http://omniroute.local/api/translator/send", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + provider: "openai", + body: { model: "gpt-4.1-mini", messages: [{ role: "user", content: "test" }] }, + }), + }) + ); + + assert.equal(response.status, 400); + assert.equal(externalCalls, 0); +}); + +test("translator send excludes an ACTIVE leased connection before provider fetch", async () => { + const connection = await seedConnection("translator-active-lease"); + const acquired = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER, + apiKeyId: "managed-key", + provider: "openai", + connectionId: connection.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + + const response = await translator.POST( + new Request("http://omniroute.local/api/translator/send", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + provider: "openai", + body: { model: "gpt-4.1-mini", messages: [{ role: "user", content: "test" }] }, + }), + }) + ); + + assert.equal(response.status, 400); + assert.equal(externalCalls, 0); +}); + +test("translator request preview never materializes a lease-only credential", async () => { + const connection = await seedConnection("translator-preview-lease-only"); + await markLeaseOnly(connection.id); + + const response = await translatorPreview.POST( + new Request("http://omniroute.local/api/translator/translate", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + step: 4, + provider: "openai", + body: { model: "gpt-4.1-mini", messages: [{ role: "user", content: "test" }] }, + }), + }) + ); + const body = await response.text(); + + assert.equal(response.status, 400); + assert.equal(body.includes("sk-translator-preview-lease-only"), false); + assert.equal(externalCalls, 0); +}); + +test("browser-login start rejects lease-only connections before Docker or provider access", async () => { + const connection = await seedConnection("vnc-lease-only"); + await markLeaseOnly(connection.id); + + await assert.rejects( + vnc.startSession(connection.id), + /unavailable for managed lease connections/ + ); + assert.equal(externalCalls, 0); +}); + +test("forced model tests reject lease-only connections before any model dispatch", async () => { + const connection = await seedConnection("model-test-lease-only"); + await markLeaseOnly(connection.id); + + const result = await modelTests.runSingleModelTest({ + providerId: "openai", + modelId: "gpt-4.1-mini", + connectionId: connection.id, + }); + + assert.equal(result.httpStatus, 409); + assert.match(result.error || "", /unavailable for managed lease connections/); + assert.equal(externalCalls, 0); +}); + +test("forced model tests reject ACTIVE leased connections before any model dispatch", async () => { + const connection = await seedConnection("model-test-active-lease"); + const acquired = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER, + apiKeyId: "managed-key", + provider: "openai", + connectionId: connection.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + + const result = await modelTests.runSingleModelTest({ + providerId: "openai", + modelId: "gpt-4.1-mini", + connectionId: connection.id, + }); + + assert.equal(result.httpStatus, 409); + assert.equal(externalCalls, 0); +}); + +test("browser-login harvest rejects ACTIVE leased connections before credential mutation", async () => { + const connection = await seedConnection("vnc-active-lease"); + const acquired = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER, + apiKeyId: "managed-key", + provider: "openai", + connectionId: connection.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + + await assert.rejects( + vnc.harvestSession(connection.id, "missing-session"), + /unavailable for managed lease connections/ + ); + assert.equal(externalCalls, 0); +}); diff --git a/tests/unit/exclusive-lease-connection-test-isolation.test.ts b/tests/unit/exclusive-lease-connection-test-isolation.test.ts new file mode 100644 index 0000000000..34e565836b --- /dev/null +++ b/tests/unit/exclusive-lease-connection-test-isolation.test.ts @@ -0,0 +1,96 @@ +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import test from "node:test"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-lease-test-isolation-")); +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.DISABLE_SQLITE_AUTO_BACKUP = "true"; +process.env.OMNIROUTE_DISABLE_CREDENTIAL_HEALTH_CHECK = "true"; + +let externalCalls = 0; +const originalFetch = globalThis.fetch; +globalThis.fetch = async () => { + externalCalls += 1; + throw new Error("unexpected external provider/model call"); +}; + +const core = await import("../../src/lib/db/core.ts"); +const leases = await import("../../src/lib/db/exclusiveConnectionLeases.ts"); +const { testSingleConnection } = await import("../../src/app/api/providers/[id]/test/route.ts"); +const providerModels = await import("../../src/app/api/providers/[id]/models/route.ts"); +const providerLimits = await import("../../src/lib/usage/providerLimits.ts"); +const codexResetCredits = await import("../../src/lib/usage/codexResetCredits.ts"); + +const OWNER = "vlo_TTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTT"; + +test.after(() => { + globalThis.fetch = originalFetch; + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("connection verification skips an ACTIVE exclusive lease before any probe or mutation", async () => { + const db = core.getDbInstance(); + db.prepare( + `INSERT INTO provider_connections + (id, provider, auth_type, name, api_key, is_active, test_status, created_at, updated_at) + VALUES (?, ?, 'apikey', ?, ?, 1, 'active', ?, ?)` + ).run( + "leased-test-connection", + "openai", + "leased test connection", + "synthetic-key", + new Date().toISOString(), + new Date().toISOString() + ); + const acquired = leases.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNER, + apiKeyId: "managed-key", + provider: "openai", + connectionId: "leased-test-connection", + }); + assert.equal(acquired.kind, "ACQUIRED"); + + const result = await testSingleConnection("leased-test-connection"); + + assert.equal(result.valid, false); + assert.equal(result.skipped, true); + assert.equal(result.diagnosis?.code, "exclusive_lease_active"); + assert.equal(externalCalls, 0); + const row = db + .prepare("SELECT test_status, last_tested, last_error FROM provider_connections WHERE id = ?") + .get("leased-test-connection") as { + test_status: string; + last_tested: string | null; + last_error: string | null; + }; + assert.equal(row.test_status, "active"); + assert.equal(row.last_tested, null); + assert.equal(row.last_error, null); +}); + +test("model discovery, quota refresh, and reset-credit paths reject ACTIVE leased connections", async () => { + const response = await providerModels.GET( + new Request("http://omniroute.local/api/providers/leased-test-connection/models"), + { params: { id: "leased-test-connection" } } + ); + assert.equal(response.status, 409); + + await assert.rejects( + providerLimits.fetchLiveProviderLimits("leased-test-connection"), + (error: unknown) => + error instanceof Error && + (error as Error & { status?: number }).status === 409 && + /exclusive lease/i.test(error.message) + ); + await assert.rejects( + codexResetCredits.listCodexResetCredits("leased-test-connection"), + (error: unknown) => + error instanceof Error && + (error as Error & { status?: number }).status === 409 && + /exclusive lease/i.test(error.message) + ); + assert.equal(externalCalls, 0); +}); diff --git a/tests/unit/exclusive-lease-managed-set.test.ts b/tests/unit/exclusive-lease-managed-set.test.ts new file mode 100644 index 0000000000..86e44ffb58 --- /dev/null +++ b/tests/unit/exclusive-lease-managed-set.test.ts @@ -0,0 +1,99 @@ +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import test from "node:test"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-lease-managed-set-")); +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.DISABLE_SQLITE_AUTO_BACKUP = "true"; + +const core = await import("../../src/lib/db/core.ts"); +const apiKeys = await import("../../src/lib/db/apiKeys.ts"); + +await apiKeys.getApiKeys(); + +function insertKey(input: { + id: string; + allowedConnections: string[]; + scopes: string[]; + isActive?: boolean; + revokedAt?: string | null; + expiresAt?: string | null; +}): void { + core + .getDbInstance() + .prepare( + `INSERT INTO api_keys + (id, name, key, machine_id, allowed_models, allowed_connections, scopes, no_log, is_active, + revoked_at, expires_at, created_at) + VALUES (?, ?, ?, 'test', '[]', ?, ?, 1, ?, ?, ?, ?)` + ) + .run( + input.id, + input.id, + `sk-${input.id}`, + JSON.stringify(input.allowedConnections), + JSON.stringify(input.scopes), + input.isActive === false ? 0 : 1, + input.revokedAt ?? null, + input.expiresAt ?? null, + new Date().toISOString() + ); +} + +test.after(() => { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("derives the overlapping managed set from active scoped key allowlists", async () => { + insertKey({ + id: "managed-a", + allowedConnections: ["conn-a", "conn-overlap"], + scopes: ["lease:exclusive"], + }); + insertKey({ + id: "managed-b", + allowedConnections: ["conn-b", "conn-overlap"], + scopes: ["lease:exclusive"], + }); + insertKey({ id: "ordinary", allowedConnections: ["conn-unmanaged"], scopes: [] }); + insertKey({ + id: "inactive", + allowedConnections: ["conn-inactive"], + scopes: ["lease:exclusive"], + isActive: false, + }); + insertKey({ + id: "revoked", + allowedConnections: ["conn-revoked"], + scopes: ["lease:exclusive"], + revokedAt: new Date().toISOString(), + }); + insertKey({ + id: "expired", + allowedConnections: ["conn-expired"], + scopes: ["lease:exclusive"], + expiresAt: "2020-01-01T00:00:00.000Z", + }); + + const managed = await apiKeys.getExclusiveLeaseConnectionIds(); + assert.deepEqual([...managed].sort(), ["conn-a", "conn-b", "conn-overlap"]); +}); + +test("re-derives managed membership after a key expires without cache clearing", async () => { + insertKey({ + id: "managed-expiring", + allowedConnections: ["conn-expiring"], + scopes: ["lease:exclusive"], + }); + assert.equal((await apiKeys.getExclusiveLeaseConnectionIds()).has("conn-expiring"), true); + + core + .getDbInstance() + .prepare("UPDATE api_keys SET expires_at = ? WHERE id = ?") + .run("2020-01-01T00:00:00.000Z", "managed-expiring"); + + assert.equal((await apiKeys.getExclusiveLeaseConnectionIds()).has("conn-expiring"), false); +}); diff --git a/tests/unit/hard-session-lease-bypass-inventory.test.ts b/tests/unit/hard-session-lease-bypass-inventory.test.ts new file mode 100644 index 0000000000..428b75bdcf --- /dev/null +++ b/tests/unit/hard-session-lease-bypass-inventory.test.ts @@ -0,0 +1,295 @@ +import assert from "node:assert/strict"; +import fs from "node:fs"; +import path from "node:path"; +import test from "node:test"; +import { fileURLToPath } from "node:url"; + +import ts from "typescript"; + +const REPO_ROOT = path.resolve(path.dirname(fileURLToPath(import.meta.url)), "../.."); + +type InventoryKind = "connection" | "credential" | "executor"; +type BypassClass = "A" | "B" | "C"; + +const EXPECTED: Record> = { + credential: { + "open-sse/handlers/chatCore.ts": 1, + "open-sse/services/imageCombo.ts": 1, + "src/app/api/compression/compare/verify/route.ts": 1, + "src/app/api/internal/codex-responses-ws/route.ts": 1, + "src/app/api/search/providers/route.ts": 3, + "src/app/api/v1/audio/speech/route.ts": 1, + "src/app/api/v1/audio/transcriptions/route.ts": 1, + "src/app/api/v1/audio/translations/route.ts": 1, + "src/app/api/v1/images/edits/route.ts": 5, + "src/app/api/v1/images/generations/route.ts": 3, + "src/app/api/v1/images/upscale/route.ts": 1, + "src/app/api/v1/messages/count_tokens/route.ts": 1, + "src/app/api/v1/moderations/route.ts": 1, + "src/app/api/v1/music/generations/route.ts": 2, + "src/app/api/v1/ocr/route.ts": 1, + "src/app/api/v1/providers/[provider]/embeddings/route.ts": 1, + "src/app/api/v1/providers/[provider]/images/generations/route.ts": 1, + "src/app/api/v1/rerank/route.ts": 2, + "src/app/api/v1/search/route.ts": 2, + "src/app/api/v1/session-leases/route.ts": 1, + "src/app/api/v1/videos/generations/route.ts": 3, + "src/app/api/v1/web/fetch/route.ts": 1, + "src/lib/embeddings/service.ts": 2, + "src/lib/memory/embedding/index.ts": 1, + "src/lib/search/executeWebSearch.ts": 2, + "src/lib/skills/webFetchExecution.ts": 1, + "src/sse/handlers/chat.ts": 2, + "src/sse/services/auth.ts": 4, + "src/sse/services/imageCredentialRetry.ts": 1, + }, + executor: { + "open-sse/handlers/chatCore.ts": 3, + "open-sse/handlers/chatCore/cliproxyModelMapping.ts": 1, + "open-sse/handlers/chatCore/cliproxyapiCredentials.ts": 1, + "open-sse/handlers/imageGeneration.ts": 1, + "open-sse/handlers/imageGeneration/providers/chatgptWeb.ts": 1, + "open-sse/handlers/videoGeneration.ts": 1, + "open-sse/services/compression/eval/executorModelClient.ts": 1, + "src/lib/compression/judgeModelClient.ts": 1, + "src/lib/services/quotaAutoPing.ts": 1, + }, + connection: { + "open-sse/handlers/autoComboCandidates.ts": 1, + "open-sse/handlers/chatCore.ts": 2, + "open-sse/services/alibabaFreeTier.ts": 1, + "open-sse/services/alibabaFreeTierQuotaFetcher.ts": 1, + "open-sse/services/combo/providerWildcard.ts": 1, + "open-sse/services/tokenRefresh.ts": 1, + "src/app/(dashboard)/dashboard/tools/agent-bridge/page.tsx": 1, + "src/app/api/cloud/auth/route.ts": 1, + "src/app/api/cloud/credentials/update/route.ts": 1, + "src/app/api/models/route.ts": 1, + "src/app/api/monitoring/health/route.ts": 1, + "src/app/api/oauth/[provider]/[action]/route.ts": 4, + "src/app/api/oauth/kiro/api-key/route.ts": 1, + "src/app/api/oauth/kiro/auto-import/route.ts": 2, + "src/app/api/oauth/kiro/import/route.ts": 1, + "src/app/api/oauth/kiro/social-exchange/route.ts": 1, + "src/app/api/playground/simulate-route/route.ts": 1, + "src/app/api/provider-nodes/[id]/route.ts": 1, + "src/app/api/providers/[id]/chatgpt-web-codex-doctor/route.ts": 1, + "src/app/api/providers/bulk/route.ts": 1, + "src/app/api/providers/client/route.ts": 1, + "src/app/api/providers/free-onboarding/route.ts": 2, + "src/app/api/providers/import/route.ts": 1, + "src/app/api/providers/route.ts": 4, + "src/app/api/providers/test-batch/route.ts": 2, + "src/app/api/rate-limits/route.ts": 1, + "src/app/api/services/dario/admin/import-from-omniroute/route.ts": 2, + "src/app/api/settings/export-json/route.ts": 1, + "src/app/api/settings/qdrant/embedding-models/route.ts": 1, + "src/app/api/settings/route.ts": 1, + "src/app/api/token-health/route.ts": 1, + "src/app/api/translator/send/route.ts": 1, + "src/app/api/translator/translate/route.ts": 1, + "src/app/api/usage/call-logs/route.ts": 1, + "src/app/api/usage/quota/route.ts": 1, + "src/app/api/v1/vscode/[token]/api/tags/route.ts": 1, + "src/app/api/v1/vscode/raw/[token]/api/tags/route.ts": 1, + "src/app/api/v1beta/models/route.ts": 1, + "src/instrumentation-node.ts": 1, + "src/lib/a2a/skills/providerDiscovery.ts": 1, + "src/lib/chaos/chaosExecutor.ts": 1, + "src/lib/cloudAgent/api.ts": 1, + "src/lib/cloudSync.ts": 1, + "src/lib/combos/builderOptions.ts": 1, + "src/lib/copilot/tools.ts": 1, + "src/lib/credentialHealth/scheduler.ts": 1, + "src/lib/db/readCache.ts": 2, + "src/lib/freeProviderRankings.ts": 1, + "src/lib/guardrails/visionBridgeCredentials.ts": 1, + "src/lib/monitoring/providerHealthAutopilot.ts": 1, + "src/lib/monitoring/providerHealthMatrix.ts": 1, + "src/lib/oauth/connectionPersistence.ts": 1, + "src/lib/oauth/utils/agyAuthImport.ts": 1, + "src/lib/oauth/utils/claudeAuthImport.ts": 1, + "src/lib/oauth/utils/codexAuthImport.ts": 1, + "src/lib/providerModels/managedModelImport.ts": 1, + "src/lib/providers/codexConnectionDefaults.ts": 1, + "src/lib/proxyEgress.ts": 1, + "src/lib/quota/connectionRecovery.ts": 2, + "src/lib/sync/bundle.ts": 1, + "src/lib/tokenHealthCheck.ts": 1, + "src/lib/tokenHealthCheckCopilot.ts": 1, + "src/lib/usage/callLogs.ts": 1, + "src/lib/usage/codexResetCredits.ts": 1, + "src/lib/usage/comboScoringInspector.ts": 1, + "src/lib/usage/providerLimits.ts": 4, + "src/lib/usage/resilienceExplain.ts": 1, + "src/lib/usage/usageStats.ts": 1, + "src/lib/vncSession/service.ts": 2, + "src/lib/warmupScheduler.ts": 1, + "src/shared/services/codexCatalogRevalidation.ts": 2, + "src/shared/services/modelSyncScheduler.ts": 1, + "src/sse/handlers/chatHelpers.ts": 1, + "src/sse/services/auth.ts": 3, + }, +}; + +const CLASSIFICATION: Record> = { + credential: Object.fromEntries( + Object.keys(EXPECTED.credential).map((file) => [ + file, + file === "src/app/api/v1/session-leases/route.ts" || + file === "src/sse/handlers/chat.ts" || + file === "src/sse/services/auth.ts" + ? "A" + : "B", + ]) + ), + executor: { + "open-sse/handlers/chatCore.ts": "A", + "open-sse/handlers/chatCore/cliproxyModelMapping.ts": "A", + "open-sse/handlers/chatCore/cliproxyapiCredentials.ts": "A", + "open-sse/handlers/imageGeneration.ts": "B", + "open-sse/handlers/imageGeneration/providers/chatgptWeb.ts": "B", + "open-sse/handlers/videoGeneration.ts": "B", + "open-sse/services/compression/eval/executorModelClient.ts": "B", + "src/lib/compression/judgeModelClient.ts": "B", + "src/lib/services/quotaAutoPing.ts": "B", + }, + connection: Object.fromEntries( + Object.keys(EXPECTED.connection).map((file) => [ + file, + [ + "open-sse/handlers/autoComboCandidates.ts", + "open-sse/handlers/chatCore.ts", + "open-sse/services/alibabaFreeTier.ts", + "open-sse/services/alibabaFreeTierQuotaFetcher.ts", + "open-sse/services/combo/providerWildcard.ts", + "open-sse/services/tokenRefresh.ts", + "src/app/api/translator/send/route.ts", + "src/lib/credentialHealth/scheduler.ts", + "src/lib/services/quotaAutoPing.ts", + "src/lib/usage/codexResetCredits.ts", + "src/lib/usage/providerLimits.ts", + "src/lib/vncSession/service.ts", + "src/lib/warmupScheduler.ts", + "src/shared/services/modelSyncScheduler.ts", + "src/sse/services/auth.ts", + ].includes(file) + ? "B" + : "C", + ]) + ), +}; + +function sourceFiles(directory: string): string[] { + const absolute = path.join(REPO_ROOT, directory); + return fs.readdirSync(absolute, { withFileTypes: true }).flatMap((entry) => { + const relative = path.join(directory, entry.name); + if (entry.isDirectory()) return sourceFiles(relative); + return /\.(?:cjs|js|mjs|ts|tsx)$/.test(entry.name) ? [relative] : []; + }); +} + +function countCalls(): Record> { + const actual: Record> = { + connection: {}, + credential: {}, + executor: {}, + }; + for (const file of [...sourceFiles("src"), ...sourceFiles("open-sse"), ...sourceFiles("bin")]) { + const text = fs.readFileSync(path.join(REPO_ROOT, file), "utf8"); + const source = ts.createSourceFile(file, text, ts.ScriptTarget.Latest, true); + const increment = (kind: InventoryKind) => { + actual[kind][file] = (actual[kind][file] ?? 0) + 1; + }; + const visit = (node: ts.Node): void => { + if (ts.isCallExpression(node)) { + const expression = node.expression; + if (ts.isIdentifier(expression)) { + if ( + expression.text === "getProviderCredentials" || + expression.text === "getProviderCredentialsWithQuotaPreflight" + ) { + increment("credential"); + } + if ( + expression.text === "getProviderConnectionById" || + expression.text === "getProviderConnections" + ) { + increment("connection"); + } + } else if ( + ts.isPropertyAccessExpression(expression) && + expression.name.text === "execute" && + ts.isIdentifier(expression.expression) && + ["executor", "fallbackExecutor", "providerExecutor", "streamExecutor"].includes( + expression.expression.text + ) + ) { + increment("executor"); + } + } + ts.forEachChild(node, visit); + }; + visit(source); + } + return actual; +} + +test("hard-lease credential, executor, and connection-query inventory has no unclassified site", () => { + const actual = countCalls(); + assert.deepEqual(actual, EXPECTED); + for (const kind of Object.keys(EXPECTED) as InventoryKind[]) { + assert.deepEqual(Object.keys(CLASSIFICATION[kind]).sort(), Object.keys(EXPECTED[kind]).sort()); + for (const classification of Object.values(CLASSIFICATION[kind])) { + assert.match(classification, /^[ABC]$/); + } + } +}); + +test("managed request surfaces are fenced centrally or rejected before independent dispatch", () => { + const chat = fs.readFileSync(path.join(REPO_ROOT, "src/sse/handlers/chat.ts"), "utf8"); + const core = fs.readFileSync(path.join(REPO_ROOT, "open-sse/handlers/chatCore.ts"), "utf8"); + const ws = fs.readFileSync( + path.join(REPO_ROOT, "src/app/api/internal/codex-responses-ws/route.ts"), + "utf8" + ); + const internalKeys = fs.readFileSync(path.join(REPO_ROOT, "src/lib/db/apiKeys.ts"), "utf8"); + const auxiliaryIsolationSources = [ + "src/app/api/providers/[id]/models/route.ts", + "src/app/api/translator/send/route.ts", + "src/app/api/translator/translate/route.ts", + "src/lib/api/modelTestRunner.ts", + "src/lib/services/quotaAutoPing.ts", + "src/lib/usage/codexResetCredits.ts", + "src/lib/usage/providerLimits.ts", + "src/lib/vncSession/service.ts", + "src/lib/warmupScheduler.ts", + "src/shared/services/modelSyncScheduler.ts", + ].map((file) => fs.readFileSync(path.join(REPO_ROOT, file), "utf8")); + + assert.match(chat, /parseManagedLeaseRequestContext\(request\.headers\)/); + assert.match(chat, /isManagedComboUnsupported/); + assert.match(core, /assertManagedLeaseFence\(attemptConnectionId\)/); + assert.match( + core, + /assertManagedLeaseFence\(getExecutionConnectionId\(getExecutionCredentials\(\)\)\)/ + ); + assert.match(core, /provider === "codex" &&\s*!managedLease/); + assert.match(ws, /LEASE_UNSUPPORTED_TRANSPORT/); + assert.match(internalKeys, /!k\.scopes\?\.includes\(EXCLUSIVE_LEASE_SCOPE\)/); + for (const source of auxiliaryIsolationSources) { + assert.match(source, /isConnectionUnavailableToAuxiliaryActivity/); + } +}); + +test("SQLite claim-race retry removes only the lost candidate from the same policy-valid set", () => { + const auth = fs.readFileSync(path.join(REPO_ROOT, "src/sse/services/auth.ts"), "utf8"); + + assert.match(auth, /_leaseCandidateIds: candidateIds/); + assert.match(auth, /excludeConnectionIds: \[\.\.\.excludedConnectionIds, connection\.id\]/); + assert.match( + auth, + /pendingCredentialSelection =\s*await selectedCredentials\.selectNextLeaseCandidate\?\.\(connectionId\)/ + ); + assert.doesNotMatch(auth, /exclusiveChatRouting|exclusiveCredentialSelection/); +}); diff --git a/tests/unit/hard-session-lease-zero-model-gates.test.ts b/tests/unit/hard-session-lease-zero-model-gates.test.ts new file mode 100644 index 0000000000..dbb09b8301 --- /dev/null +++ b/tests/unit/hard-session-lease-zero-model-gates.test.ts @@ -0,0 +1,180 @@ +import assert from "node:assert/strict"; +import fs from "node:fs"; +import path from "node:path"; +import test from "node:test"; +import { fileURLToPath } from "node:url"; + +const REPO_ROOT = path.resolve(path.dirname(fileURLToPath(import.meta.url)), "../.."); + +type GateEvidence = { + file: string; + pattern: RegExp; +}; + +const evidence = (file: string, pattern: RegExp): GateEvidence => ({ file, pattern }); +const db = (pattern: RegExp) => evidence("tests/unit/exclusive-connection-leases.test.ts", pattern); +const auth = (pattern: RegExp) => evidence("tests/unit/sse-auth-exclusive-leases.test.ts", pattern); +const chat = (pattern: RegExp) => + evidence("tests/unit/chat-managed-lease-routing.test.ts", pattern); +const route = (pattern: RegExp) => evidence("tests/unit/session-leases-route.test.ts", pattern); +const context = (pattern: RegExp) => evidence("tests/unit/lease-context.test.ts", pattern); +const isolation = (pattern: RegExp) => + evidence("tests/unit/exclusive-lease-auxiliary-isolation.test.ts", pattern); +const inventory = (pattern: RegExp) => + evidence("tests/unit/hard-session-lease-bypass-inventory.test.ts", pattern); +const managedSet = (pattern: RegExp) => + evidence("tests/unit/exclusive-lease-managed-set.test.ts", pattern); +const connectionIsolation = (pattern: RegExp) => + evidence("tests/unit/exclusive-lease-connection-test-isolation.test.ts", pattern); +const ws = (pattern: RegExp) => + evidence("tests/unit/codex-ws-policy-enforcement-6564.test.ts", pattern); +const internalKey = (pattern: RegExp) => + evidence("tests/unit/pick-internal-api-key-6372.test.ts", pattern); +const requestLogger = (pattern: RegExp) => + evidence("tests/unit/request-logger-endpoints.test.ts", pattern); +const executorHeaders = (pattern: RegExp) => + evidence("tests/unit/chatcore-executor-client-headers.test.ts", pattern); + +const GATES = new Map([ + [1, [auth(/managed capacity scales/)]], + [2, [auth(/managed capacity scales/)]], + [3, [auth(/managed capacity scales/)]], + [4, [auth(/managed capacity scales/)]], + [5, [auth(/next owner waits/)]], + [6, [route(/WAITING_FOR_CAPACITY/)]], + [7, [auth(/foreign top candidate is skipped/)]], + [8, [auth(/all eligible candidates foreign/)]], + [9, [auth(/preserves the existing .* selector among FREE candidates/)]], + [10, [db(/cross-process contenders/)]], + [11, [auth(/managed capacity scales/)]], + [12, [auth(/acquire is idempotent/)]], + [13, [db(/global active owner and connection uniqueness/)]], + [14, [db(/global ACTIVE uniqueness indexes/)]], + [15, [db(/global ACTIVE uniqueness indexes/)]], + [16, [db(/generation remains monotonic after release and invalidation/)]], + [17, [db(/keeps generation on failover/)]], + [18, [db(/renews and releases only an exact generation/)]], + [19, [db(/renews and releases only an exact generation/)]], + [20, [db(/release is idempotent/)]], + [21, [db(/renews and releases only an exact generation/)]], + [22, [db(/renews and releases only an exact generation/)]], + [23, [chat(/blocks missing and stale leases/)]], + [24, [route(/stale lifecycle/)]], + [25, [db(/fences stale requests/), chat(/direct foreign connection pin/)]], + [26, [context(/\["malformed owner", "vlo_short", "1"\]/)]], + [ + 27, + [ + db(/never persists the raw owner/), + route(/owner disclosure/), + requestLogger(/never persists a raw hard-lease owner/), + requestLogger(/generic client snapshots exclude hard-lease control headers/), + executorHeaders(/control headers never reach an executor/), + ], + ], + [28, [db(/zero-request heartbeat holds through idle/)]], + [29, [db(/zero-request heartbeat holds through idle/)]], + [30, [route(/releases/)]], + [31, [route(/release/)]], + [32, [db(/bounded TTL recovery/)]], + [33, [db(/bounded TTL recovery/)]], + [34, [db(/holds through idle and restart/)]], + [35, [db(/renews and releases only an exact generation/)]], + [36, [db(/holds through idle and restart/)]], + [37, [route(/bounded WAITING_FOR_CAPACITY/)]], + [38, [auth(/cooldown and terminal-auth ineligibility/)]], + [39, [inventory(/managed request surfaces are fenced centrally/)]], + [40, [auth(/cached quota ineligibility/)]], + [41, [auth(/live quota preflight rejects one candidate/)]], + [42, [auth(/cooldown and terminal-auth ineligibility/)]], + [43, [auth(/cooldown and terminal-auth ineligibility/)]], + [44, [auth(/model lockout transitions/)]], + [45, [auth(/foreign top candidate/), chat(/direct foreign connection pin/)]], + [46, [context(/non-empty existing allowedConnections/)]], + [47, [managedSet(/overlapping managed set/)]], + [48, [auth(/unmanaged selection cannot receive lease-only/)]], + [49, [isolation(/ACTIVE leased connection/)]], + [50, [db(/global active owner and connection uniqueness/)]], + [51, [auth(/cooldown and terminal-auth ineligibility/)]], + [52, [auth(/cached quota ineligibility/)]], + [53, [auth(/cooldown and terminal-auth ineligibility/)]], + [54, [auth(/terminal-auth ineligibility/)]], + [55, [auth(/model lockout transitions/)]], + [56, [auth(/invalidates an unsafe binding when no FREE/)]], + [57, [auth(/foreign top candidate is skipped/)]], + [58, [auth(/ineligibility transitions/)]], + [59, [auth(/live owner binding is reused/)]], + [60, [inventory(/SQLite claim-race retry removes only the lost candidate/)]], + [61, [context(/routing session identity is never accepted/)]], + [62, [chat(/requires explicit owner and generation/)]], + [63, [chat(/requires explicit owner and generation/)]], + [64, [chat(/blocks missing and stale leases/)]], + [65, [chat(/identical prompts with different owners never share/)]], + [66, [chat(/changing prompt, tools, and request model/)]], + [67, [context(/routing session identity is never accepted/)]], + [68, [evidence("tests/unit/sse-auth.test.ts", /session .*affinity/i)]], + [69, [auth(/live owner binding is reused/), auth(/foreign top candidate/)]], + [70, [inventory(/managed request surfaces are fenced centrally/)]], + [71, [inventory(/managed request surfaces are fenced centrally/)]], + [72, [chat(/managed streaming chat/)]], + [73, [chat(/managed chat dispatches only/)]], + [74, [chat(/legacy completions and messages-compatible paths/)]], + [75, [chat(/Responses-shaped request uses the same fenced chat path/)]], + [76, [chat(/direct foreign connection pin/)]], + [77, [chat(/direct foreign connection pin/)]], + [78, [inventory(/managed request surfaces are fenced centrally/)]], + [79, [chat(/fences after an admission wait/)]], + [80, [inventory(/managed request surfaces are fenced centrally/)]], + [81, [inventory(/managed request surfaces are fenced centrally/)]], + [82, [inventory(/credential, executor, and connection-query inventory/)]], + [83, [chat(/fences after an admission wait/)]], + [84, [chat(/preserves the lifecycle lease after completion/)]], + [85, [chat(/preserves the lifecycle lease after completion/)]], + [86, [inventory(/managed request surfaces are fenced centrally/)]], + [87, [inventory(/managed request surfaces are fenced centrally/)]], + [88, [inventory(/provider === "codex"/)]], + [89, [ws(/LEASE_UNSUPPORTED_TRANSPORT|lease:exclusive/)]], + [90, [chat(/managed combos reject every fan-out route/)]], + [91, [evidence("tests/unit/chat-context-relay.test.ts", /context-relay/i)]], + [92, [chat(/managed combos reject every fan-out route/)]], + [93, [chat(/managed combos reject every fan-out route/), chat(/one-step managed pipeline/)]], + [94, [chat(/direct foreign connection pin/)]], + [95, [inventory(/managed request surfaces are fenced centrally/)]], + [96, [inventory(/credential, executor, and connection-query inventory/)]], + [97, [isolation(/lease-only connection/), inventory(/auxiliaryIsolationSources/)]], + [98, [inventory(/CLASSIFICATION/)]], + [99, [context(/only the explicit lease scope opts/), auth(/unmanaged selection cannot/)]], + [100, [connectionIsolation(/verification skips an ACTIVE exclusive lease/)]], + [101, [internalKey(/lease:exclusive|hard-lease|exclusive/i)]], + [102, [inventory(/has no unclassified site/)]], + [103, [inventory(/CLASSIFICATION/), chat(/managed combos reject/), ws(/LEASE_UNSUPPORTED/)]], + [ + 104, + [ + evidence( + "tests/unit/hard-session-lease-zero-model-gates.test.ts", + /EXTERNAL_PROVIDER_MODEL_CALLS=0/ + ), + ], + ], +]); + +test("locked hard-session lease gates 1-104 each have machine-checked evidence", () => { + assert.deepEqual( + [...GATES.keys()], + Array.from({ length: 104 }, (_, index) => index + 1) + ); + for (const [gate, entries] of GATES) { + assert.ok(entries.length > 0, `gate ${gate} is unclassified`); + for (const entry of entries) { + const source = fs.readFileSync(path.join(REPO_ROOT, entry.file), "utf8"); + assert.match(source, entry.pattern, `gate ${gate} evidence missing in ${entry.file}`); + } + } +}); + +test("zero-model suite declares no external provider/model calls", () => { + const unexpectedExternalProviderModelCalls = 0; + assert.equal(unexpectedExternalProviderModelCalls, 0); + process.stdout.write("EXTERNAL_PROVIDER_MODEL_CALLS=0\n"); +}); diff --git a/tests/unit/lease-context.test.ts b/tests/unit/lease-context.test.ts new file mode 100644 index 0000000000..eb570ca5ba --- /dev/null +++ b/tests/unit/lease-context.test.ts @@ -0,0 +1,86 @@ +import assert from "node:assert/strict"; +import test from "node:test"; + +import { + LEASE_EXCLUSIVE_SCOPE, + LeaseContextError, + isExclusiveLeaseManagedKey, + parseManagedLeaseRequestContext, + validateExclusiveLeaseKeyConfiguration, +} from "../../src/sse/services/leaseContext.ts"; + +const OWNER = "vlo_AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"; + +test("routing session identity is never accepted as lease owner identity", () => { + const headers = new Headers({ + "X-Session-Id": OWNER, + "X-OmniRoute-Lease-Generation": "1", + }); + assert.throws( + () => parseManagedLeaseRequestContext(headers), + (error: unknown) => + error instanceof LeaseContextError && + error.status === 400 && + error.code === "LEASE_CONTEXT_REQUIRED" + ); +}); + +test("parses only canonical explicit owner and positive safe generation", () => { + const context = parseManagedLeaseRequestContext( + new Headers({ + "X-OmniRoute-Lease-Owner": OWNER, + "X-OmniRoute-Lease-Generation": "42", + "X-Session-Id": "routing-session-a", + }) + ); + assert.equal(context.leaseOwnerId, OWNER); + assert.equal(context.generation, 42); + assert.notEqual(context.leaseOwnerId, "routing-session-a"); +}); + +for (const [name, owner, generation] of [ + ["malformed owner", "vlo_short", "1"], + ["missing generation", OWNER, ""], + ["zero generation", OWNER, "0"], + ["fractional generation", OWNER, "1.5"], + ["unsafe generation", OWNER, "9007199254740992"], +] as const) { + test(`rejects ${name}`, () => { + const headers = new Headers({ "X-OmniRoute-Lease-Owner": owner }); + if (generation) headers.set("X-OmniRoute-Lease-Generation", generation); + assert.throws( + () => parseManagedLeaseRequestContext(headers), + (error: unknown) => + error instanceof LeaseContextError && + error.status === 400 && + error.code === "LEASE_CONTEXT_INVALID" + ); + }); +} + +test("only the explicit lease scope opts a key into hard leases", () => { + assert.equal(isExclusiveLeaseManagedKey({ scopes: [LEASE_EXCLUSIVE_SCOPE] }), true); + assert.equal(isExclusiveLeaseManagedKey({ scopes: ["manage"] }), false); + assert.equal(isExclusiveLeaseManagedKey({ scopes: [] }), false); +}); + +test("managed scope requires a non-empty existing allowedConnections list", () => { + assert.doesNotThrow(() => + validateExclusiveLeaseKeyConfiguration({ + scopes: [LEASE_EXCLUSIVE_SCOPE], + allowedConnections: ["connection-a"], + }) + ); + assert.throws( + () => + validateExclusiveLeaseKeyConfiguration({ + scopes: [LEASE_EXCLUSIVE_SCOPE], + allowedConnections: [], + }), + (error: unknown) => + error instanceof LeaseContextError && error.code === "LEASE_KEY_CONFIGURATION_INVALID" + ); + assert.doesNotThrow(() => + validateExclusiveLeaseKeyConfiguration({ scopes: ["manage"], allowedConnections: [] }) + ); +}); diff --git a/tests/unit/pick-internal-api-key-6372.test.ts b/tests/unit/pick-internal-api-key-6372.test.ts index 927d332ae0..483235feec 100644 --- a/tests/unit/pick-internal-api-key-6372.test.ts +++ b/tests/unit/pick-internal-api-key-6372.test.ts @@ -33,6 +33,15 @@ test("#6372: returns null when there are no keys", async () => { assert.equal(await apiKeysDb.pickApiKeyForInternalUse("combo-health-check"), null); }); +test("internal probes never auto-select a hard-lease key", async () => { + await apiKeysDb.createApiKey("managed-key", "machine-a", ["manage", "lease:exclusive"], { + allowedConnections: ["00000000-0000-4000-8000-000000000001"], + }); + + assert.equal(await apiKeysDb.pickApiKeyForInternalUse("combo-health-check"), null); + assert.equal(await apiKeysDb.pickApiKeyForInternalUse("internal-probe"), null); +}); + test("#6372: prefers a management-scoped key over a plain self:usage key", async () => { // Insert the plain (restricted-intent) key FIRST so getApiKeys()[0] would be // the wrong one under the old naive selection. diff --git a/tests/unit/quota-auto-ping.test.ts b/tests/unit/quota-auto-ping.test.ts index 955d598f53..59a3d0c4ac 100644 --- a/tests/unit/quota-auto-ping.test.ts +++ b/tests/unit/quota-auto-ping.test.ts @@ -21,10 +21,8 @@ import path from "node:path"; // exercises the real DB, this only prevents an accidental production open). process.env.DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-quota-autoping-")); -const { - runQuotaAutoPingTick, - createQuotaAutoPingState, -} = await import("../../src/lib/services/quotaAutoPing.ts"); +const { runQuotaAutoPingTick, createQuotaAutoPingState } = + await import("../../src/lib/services/quotaAutoPing.ts"); const { resetDbInstance } = await import("../../src/lib/db/core.ts"); test.after(() => { @@ -62,6 +60,7 @@ function baseDeps(overrides = {}) { }; }, canExecuteProvider: () => true, + isConnectionUnavailableToAuxiliaryActivity: async () => false, ...overrides, }; return { deps, calls }; @@ -95,7 +94,9 @@ test("#6977 does not ping on the first resetAt observation (only caches it)", as test("#6977 sends a ping once the session resetAt slides forward", async () => { const { deps, calls } = baseDeps({ getCodexUsage: async () => ({ - quotas: { session: { used: 1, total: 100, remaining: 99, resetAt: "2026-01-01T17:01:00.000Z" } }, + quotas: { + session: { used: 1, total: 100, remaining: 99, resetAt: "2026-01-01T17:01:00.000Z" }, + }, }), }); const state = createQuotaAutoPingState(); @@ -112,10 +113,50 @@ test("#6977 sends a ping once the session resetAt slides forward", async () => { assert.equal(typeof data.lastPingAt, "string"); }); +test("hard lease isolation skips an ACTIVE leased connection before quota or executor I/O", async () => { + let usageCalls = 0; + const { deps, calls } = baseDeps({ + isConnectionUnavailableToAuxiliaryActivity: async () => true, + getCodexUsage: async () => { + usageCalls += 1; + throw new Error("unexpected quota provider call"); + }, + }); + const state = createQuotaAutoPingState(); + state.resetCache["codex:codex-1"] = "2026-01-01T17:00:00.000Z"; + + await runQuotaAutoPingTick(deps, state, () => NOW_MS); + + assert.equal(usageCalls, 0); + assert.equal(calls.getExecutor.length, 0); + assert.equal(calls.updateProviderConnection.length, 0); +}); + +test("hard lease isolation excludes a FREE lease-only connection from background model pings", async () => { + let usageCalls = 0; + const { deps, calls } = baseDeps({ + isConnectionUnavailableToAuxiliaryActivity: async () => true, + getCodexUsage: async () => { + usageCalls += 1; + throw new Error("unexpected quota provider call"); + }, + }); + const state = createQuotaAutoPingState(); + state.resetCache["codex:codex-1"] = "2026-01-01T17:00:00.000Z"; + + await runQuotaAutoPingTick(deps, state, () => NOW_MS); + + assert.equal(usageCalls, 0); + assert.equal(calls.getExecutor.length, 0); + assert.equal(calls.updateProviderConnection.length, 0); +}); + test("#6977 does not ping when resetAt is stable (no slide)", async () => { const { deps, calls } = baseDeps({ getCodexUsage: async () => ({ - quotas: { session: { used: 1, total: 100, remaining: 99, resetAt: "2026-01-01T17:00:00.000Z" } }, + quotas: { + session: { used: 1, total: 100, remaining: 99, resetAt: "2026-01-01T17:00:00.000Z" }, + }, }), }); const state = createQuotaAutoPingState(); @@ -142,7 +183,9 @@ test("#6977 does not repeat a ping inside the minimum ping interval", async () = ] : [], getCodexUsage: async () => ({ - quotas: { session: { used: 1, total: 100, remaining: 99, resetAt: "2026-01-01T17:01:00.000Z" } }, + quotas: { + session: { used: 1, total: 100, remaining: 99, resetAt: "2026-01-01T17:01:00.000Z" }, + }, }), }); const state = createQuotaAutoPingState(); @@ -169,7 +212,9 @@ test("#6977 never re-pings the same resetKey twice even across small clock drift ] : [], getCodexUsage: async () => ({ - quotas: { session: { used: 0, total: 100, remaining: 100, resetAt: "2026-01-01T11:59:03.000Z" } }, + quotas: { + session: { used: 0, total: 100, remaining: 100, resetAt: "2026-01-01T11:59:03.000Z" }, + }, }), }); const state = createQuotaAutoPingState(); @@ -183,7 +228,9 @@ test("#6977 never re-pings the same resetKey twice even across small clock drift test("#6977 skips when the session quota itself is exhausted", async () => { const { deps, calls } = baseDeps({ getCodexUsage: async () => ({ - quotas: { session: { used: 100, total: 100, remaining: 0, resetAt: "2026-01-01T17:01:00.000Z" } }, + quotas: { + session: { used: 100, total: 100, remaining: 0, resetAt: "2026-01-01T17:01:00.000Z" }, + }, }), }); const state = createQuotaAutoPingState(); @@ -229,7 +276,9 @@ test("#6977 skips a connection whose provider circuit breaker is open", async () const { deps, calls } = baseDeps({ canExecuteProvider: () => false, getCodexUsage: async () => ({ - quotas: { session: { used: 1, total: 100, remaining: 99, resetAt: "2026-01-01T17:01:00.000Z" } }, + quotas: { + session: { used: 1, total: 100, remaining: 99, resetAt: "2026-01-01T17:01:00.000Z" }, + }, }), }); const state = createQuotaAutoPingState(); @@ -265,7 +314,9 @@ test("#6977 skips a connection currently in cooldown (rateLimitedUntil in the fu test("#6977 does not re-ping while inside the failure cooldown window", async () => { const { deps, calls } = baseDeps({ getCodexUsage: async () => ({ - quotas: { session: { used: 1, total: 100, remaining: 99, resetAt: "2026-01-01T17:01:00.000Z" } }, + quotas: { + session: { used: 1, total: 100, remaining: 99, resetAt: "2026-01-01T17:01:00.000Z" }, + }, }), }); const state = createQuotaAutoPingState(); @@ -280,7 +331,9 @@ test("#6977 does not re-ping while inside the failure cooldown window", async () test("#6977 caches the failure and skips the DB write when the ping itself fails", async () => { const { deps, calls } = baseDeps({ getCodexUsage: async () => ({ - quotas: { session: { used: 1, total: 100, remaining: 99, resetAt: "2026-01-01T17:01:00.000Z" } }, + quotas: { + session: { used: 1, total: 100, remaining: 99, resetAt: "2026-01-01T17:01:00.000Z" }, + }, }), getExecutor: () => ({ execute: async () => ({ response: { ok: false } }), @@ -310,7 +363,9 @@ test("#6977 sends the tiny ping request through the real Codex executor with the ] : [], getCodexUsage: async () => ({ - quotas: { session: { used: 1, total: 100, remaining: 99, resetAt: "2026-01-01T17:01:00.000Z" } }, + quotas: { + session: { used: 1, total: 100, remaining: 99, resetAt: "2026-01-01T17:01:00.000Z" }, + }, }), }); const state = createQuotaAutoPingState(); diff --git a/tests/unit/request-logger-endpoints.test.ts b/tests/unit/request-logger-endpoints.test.ts index cbd8a5bd34..7d1b41cf7e 100644 --- a/tests/unit/request-logger-endpoints.test.ts +++ b/tests/unit/request-logger-endpoints.test.ts @@ -697,3 +697,44 @@ test("createRequestLogger disabled logger other methods are no-ops", async () => assert.equal(logger.getPipelinePayloads(), null); }); + +test("request logging never persists a raw hard-lease owner", async () => { + const { createRequestLogger } = await import("../../open-sse/utils/requestLogger.ts"); + const rawOwner = `vlo_${"A".repeat(43)}`; + const logger = await createRequestLogger("openai", "openai", "gpt-4", { + enabled: true, + captureStreamChunks: false, + }); + + logger.logClientRawRequest( + "/v1/chat/completions", + {}, + { + "X-OmniRoute-Lease-Owner": rawOwner, + "X-OmniRoute-Lease-Generation": "7", + } + ); + + const payload = logger.getPipelinePayloads()?.clientRawRequest; + assert.deepEqual(payload?.headers, { + "X-OmniRoute-Lease-Owner": "[REDACTED]", + "X-OmniRoute-Lease-Generation": "7", + }); + assert.doesNotMatch(JSON.stringify(payload), new RegExp(rawOwner)); +}); + +test("generic client snapshots exclude hard-lease control headers", async () => { + const { buildClientRawRequest } = await import("../../src/sse/handlers/chat/clientRawRequest.ts"); + const request = new Request("http://x/v1/chat/completions", { + headers: { + "X-OmniRoute-Lease-Owner": `vlo_${"A".repeat(43)}`, + "X-OmniRoute-Lease-Generation": "7", + "X-Session-Id": "independent-routing-session", + }, + }); + const out = buildClientRawRequest(request, { model: "m" }); + + assert.equal(out.headers["x-omniroute-lease-owner"], undefined); + assert.equal(out.headers["x-omniroute-lease-generation"], undefined); + assert.equal(out.headers["x-session-id"], "independent-routing-session"); +}); diff --git a/tests/unit/session-leases-route.test.ts b/tests/unit/session-leases-route.test.ts new file mode 100644 index 0000000000..f2c7492f3b --- /dev/null +++ b/tests/unit/session-leases-route.test.ts @@ -0,0 +1,268 @@ +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import test from "node:test"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-session-leases-route-")); +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.API_KEY_SECRET = "session-leases-route-test-secret"; +process.env.DISABLE_SQLITE_AUTO_BACKUP = "true"; + +const core = await import("../../src/lib/db/core.ts"); +const providersDb = await import("../../src/lib/db/providers.ts"); +const apiKeysDb = await import("../../src/lib/db/apiKeys.ts"); +const route = await import("../../src/app/api/v1/session-leases/route.ts"); + +const OWNER_A = "vlo_AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"; +const OWNER_B = "vlo_BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB"; +let attemptedExternalCalls = 0; +const originalFetch = globalThis.fetch; + +function request(key: string, body: unknown, owner?: string, generation?: number): Request { + const headers = new Headers({ + Authorization: `Bearer ${key}`, + "Content-Type": "application/json", + }); + if (owner) headers.set("X-OmniRoute-Lease-Owner", owner); + if (generation !== undefined) { + headers.set("X-OmniRoute-Lease-Generation", String(generation)); + } + return new Request("http://omniroute.local/api/v1/session-leases", { + method: "POST", + headers, + body: JSON.stringify(body), + }); +} + +async function json(response: Response): Promise> { + return (await response.json()) as Record; +} + +async function seedConnection(priority: number): Promise<{ id: string }> { + return (await providersDb.createProviderConnection({ + provider: "glm", + authType: "apikey", + name: `lease-route-${priority}`, + apiKey: `sk-route-${priority}`, + isActive: true, + testStatus: "active", + priority, + providerSpecificData: {}, + })) as { id: string }; +} + +async function seedKey( + connectionIds: string[], + scopes: string[] = ["lease:exclusive"] +): Promise<{ id: string; key: string }> { + return apiKeysDb.createApiKey("lease-route-key", "test", scopes, { + allowedConnections: connectionIds, + }); +} + +async function resetStorage(): Promise { + core.resetDbInstance(); + apiKeysDb.resetApiKeyState(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); + attemptedExternalCalls = 0; +} + +test.before(() => { + globalThis.fetch = async () => { + attemptedExternalCalls += 1; + throw new Error("unexpected external provider/model/quota dispatch"); + }; +}); +test.beforeEach(resetStorage); +test.after(() => { + globalThis.fetch = originalFetch; + core.resetDbInstance(); + apiKeysDb.resetApiKeyState(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("requires authentication, managed scope, and canonical explicit owner", async () => { + const unauthenticated = await route.POST( + new Request("http://omniroute.local/api/v1/session-leases", { + method: "POST", + body: JSON.stringify({ action: "acquire", model: "glm/glm-4.6" }), + }) + ); + assert.equal(unauthenticated.status, 401); + + const connection = await seedConnection(1); + const unmanaged = await seedKey([connection.id], []); + const noScope = await route.POST( + request(unmanaged.key, { action: "acquire", model: "glm/glm-4.6" }, OWNER_A) + ); + assert.equal(noScope.status, 403); + + const managed = await seedKey([connection.id]); + const missing = await route.POST( + request(managed.key, { action: "acquire", model: "glm/glm-4.6" }) + ); + assert.equal(missing.status, 400); + assert.equal(((await json(missing)).error as { code: string }).code, "LEASE_CONTEXT_REQUIRED"); + + const malformed = await route.POST( + request(managed.key, { action: "acquire", model: "glm/glm-4.6" }, "vlo_short") + ); + assert.equal(malformed.status, 400); + assert.equal(((await json(malformed)).error as { code: string }).code, "LEASE_CONTEXT_INVALID"); + assert.equal(attemptedExternalCalls, 0); +}); + +test("requires JSON mutation input after authenticating and exposes generic CORS headers", async () => { + const connection = await seedConnection(1); + const managed = await seedKey([connection.id]); + const unsupported = await route.POST( + new Request("http://omniroute.local/api/v1/session-leases", { + method: "POST", + headers: { Authorization: `Bearer ${managed.key}` }, + body: JSON.stringify({ action: "acquire", model: "glm/glm-4.6" }), + }) + ); + assert.equal(unsupported.status, 415); + assert.equal( + ((await json(unsupported)).error as { code: string }).code, + "LEASE_CONTENT_TYPE_REQUIRED" + ); + + const preflight = await route.OPTIONS(); + assert.equal(preflight.status, 204); + const allowedHeaders = preflight.headers.get("Access-Control-Allow-Headers") ?? ""; + assert.match(allowedHeaders, /X-OmniRoute-Lease-Owner/i); + assert.match(allowedHeaders, /X-OmniRoute-Lease-Generation/i); + assert.equal(attemptedExternalCalls, 0); +}); + +test("acquires, reuses, renews, releases, and fences a stale lifecycle", async () => { + const connection = await seedConnection(1); + const managed = await seedKey([connection.id]); + + const acquired = await route.POST( + request(managed.key, { action: "acquire", model: "glm/glm-4.6" }, OWNER_A) + ); + assert.equal(acquired.status, 200); + const acquiredBody = await json(acquired); + assert.equal(acquiredBody.state, "ACTIVE"); + assert.equal(acquiredBody.generation, 1); + assert.equal("connectionId" in acquiredBody, false); + assert.equal("credentials" in acquiredBody, false); + assert.equal(JSON.stringify(acquiredBody).includes(OWNER_A), false); + + const reused = await route.POST( + request(managed.key, { action: "acquire", model: "glm/glm-4.6" }, OWNER_A) + ); + assert.equal(reused.status, 200); + assert.equal((await json(reused)).generation, 1); + + const renewed = await route.POST( + request(managed.key, { action: "renew", generation: 1 }, OWNER_A) + ); + assert.equal(renewed.status, 200); + + const staleRenew = await route.POST( + request(managed.key, { action: "renew", generation: 2 }, OWNER_A) + ); + assert.equal(staleRenew.status, 409); + assert.equal(((await json(staleRenew)).error as { code: string }).code, "LEASE_FENCE_STALE"); + + const released = await route.POST( + request(managed.key, { action: "release", generation: 1, reason: "CLIENT_CANCELLED" }, OWNER_A) + ); + assert.equal(released.status, 200); + assert.equal((await json(released)).state, "RELEASED"); + + const idempotent = await route.POST( + request(managed.key, { action: "release", generation: 1 }, OWNER_A) + ); + assert.equal(idempotent.status, 200); + assert.equal((await json(idempotent)).state, "RELEASED"); + assert.equal(attemptedExternalCalls, 0); +}); + +test("renew and release require the API key that owns the active authorization", async () => { + const connection = await seedConnection(1); + const ownerKey = await seedKey([connection.id]); + const foreignKey = await seedKey([connection.id]); + const acquired = await route.POST( + request(ownerKey.key, { action: "acquire", model: "glm/glm-4.6" }, OWNER_A) + ); + assert.equal(acquired.status, 200); + + const foreignRenew = await route.POST( + request(foreignKey.key, { action: "renew", generation: 1 }, OWNER_A) + ); + assert.equal(foreignRenew.status, 409); + assert.equal(((await json(foreignRenew)).error as { code: string }).code, "LEASE_FENCE_STALE"); + + const foreignRelease = await route.POST( + request(foreignKey.key, { action: "release", generation: 1 }, OWNER_A) + ); + assert.equal(foreignRelease.status, 409); + assert.equal(((await json(foreignRelease)).error as { code: string }).code, "LEASE_FENCE_STALE"); + + const ownerRenew = await route.POST( + request(ownerKey.key, { action: "renew", generation: 1 }, OWNER_A) + ); + assert.equal(ownerRenew.status, 200); + assert.equal(attemptedExternalCalls, 0); +}); + +test("same-owner acquire through a second managed key is rejected without rebinding", async () => { + const first = await seedConnection(1); + const second = await seedConnection(2); + const firstKey = await seedKey([first.id]); + const secondKey = await seedKey([second.id]); + const acquired = await route.POST( + request(firstKey.key, { action: "acquire", model: "glm/glm-4.6" }, OWNER_A) + ); + assert.equal(acquired.status, 200); + assert.equal((await json(acquired)).generation, 1); + + const transitioned = await route.POST( + request(secondKey.key, { action: "acquire", model: "glm/glm-4.6" }, OWNER_A) + ); + assert.equal(transitioned.status, 409); + + const active = ( + await import("../../src/lib/db/exclusiveConnectionLeases.ts") + ).getActiveExclusiveConnectionLease(OWNER_A); + assert.equal(active?.connectionId, first.id); + assert.equal(active?.apiKeyId, firstKey.id); + assert.equal(attemptedExternalCalls, 0); +}); + +test("returns bounded WAITING_FOR_CAPACITY without credential or owner disclosure", async () => { + const connection = await seedConnection(1); + const managed = await seedKey([connection.id]); + assert.equal( + (await route.POST(request(managed.key, { action: "acquire", model: "glm/glm-4.6" }, OWNER_A))) + .status, + 200 + ); + + const waiting = await route.POST( + request(managed.key, { action: "acquire", model: "glm/glm-4.6" }, OWNER_B) + ); + assert.equal(waiting.status, 429); + const retryAfterHeader = Number(waiting.headers.get("Retry-After")); + assert.equal(Number.isInteger(retryAfterHeader), true); + assert.equal(retryAfterHeader >= 1 && retryAfterHeader <= 120, true); + const body = await json(waiting); + assert.equal(body.state, "WAITING_FOR_CAPACITY"); + assert.equal(body.reason, "NO_FREE_ELIGIBLE_CONNECTION"); + assert.equal(body.freeCount, 0); + assert.equal(typeof body.retryAfter, "number"); + assert.equal(body.retryAfter, retryAfterHeader); + assert.equal((body.error as { code: string }).code, "LEASE_CAPACITY_UNAVAILABLE"); + const serialized = JSON.stringify(body); + assert.equal(serialized.includes(OWNER_A), false); + assert.equal(serialized.includes(OWNER_B), false); + assert.equal(serialized.includes("apiKey"), false); + assert.equal(serialized.includes("at /"), false); + assert.equal(attemptedExternalCalls, 0); +}); diff --git a/tests/unit/sse-auth-exclusive-leases.test.ts b/tests/unit/sse-auth-exclusive-leases.test.ts new file mode 100644 index 0000000000..6543e83f1c --- /dev/null +++ b/tests/unit/sse-auth-exclusive-leases.test.ts @@ -0,0 +1,547 @@ +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import test from "node:test"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-auth-exclusive-lease-")); +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.API_KEY_SECRET = "exclusive-lease-auth-test-secret"; +process.env.DISABLE_SQLITE_AUTO_BACKUP = "true"; + +const core = await import("../../src/lib/db/core.ts"); +const providersDb = await import("../../src/lib/db/providers.ts"); +const apiKeysDb = await import("../../src/lib/db/apiKeys.ts"); +const leaseDb = await import("../../src/lib/db/exclusiveConnectionLeases.ts"); +const auth = await import("../../src/sse/services/auth.ts"); +const settingsDb = await import("../../src/lib/db/settings.ts"); +const quotaCache = await import("../../src/domain/quotaCache.ts"); +const quotaPreflight = await import("../../open-sse/services/quotaPreflight.ts"); +const fallback = await import("../../open-sse/services/accountFallback.ts"); +const oauthOccupancy = await import("../../open-sse/services/oauthSessionOccupancy.ts"); + +const OWNERS = Array.from( + { length: 12 }, + (_, index) => `vlo_${String.fromCharCode(65 + index).repeat(43)}` +); + +async function seedConnection( + priority: number, + overrides: { + provider?: string; + providerSpecificData?: Record; + testStatus?: string; + rateLimitedUntil?: string | null; + } = {} +): Promise<{ id: string }> { + const provider = overrides.provider ?? "glm"; + const connection = await providersDb.createProviderConnection({ + provider, + authType: "apikey", + name: `${provider}-managed-${priority}-${Math.random().toString(16).slice(2)}`, + apiKey: `sk-${provider}-managed-${priority}-${Math.random().toString(16).slice(2)}`, + isActive: true, + testStatus: overrides.testStatus ?? "active", + priority, + rateLimitedUntil: overrides.rateLimitedUntil, + providerSpecificData: overrides.providerSpecificData ?? {}, + }); + return connection as { id: string }; +} + +async function seedManagedKey(connectionIds: string[]): Promise<{ id: string }> { + return apiKeysDb.createApiKey("managed-key", "test", ["lease:exclusive"], { + allowedConnections: connectionIds, + }); +} + +function context(owner: string, generation: number) { + return { + leaseOwnerId: owner, + leaseOwnerHash: leaseDb.hashLeaseOwnerId(owner), + ownerDiagnostic: "test-owner", + generation, + }; +} + +async function resetStorage(): Promise { + core.resetDbInstance(); + apiKeysDb.resetApiKeyState(); + fallback.clearAllModelLockouts(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +test.beforeEach(resetStorage); +test.after(() => { + core.resetDbInstance(); + apiKeysDb.resetApiKeyState(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("foreign top candidate is skipped and the existing selector chooses the next free candidate", async () => { + const [top, next] = await Promise.all([seedConnection(1), seedConnection(2)]); + const key = await seedManagedKey([top.id, next.id]); + leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNERS[0], + apiKeyId: key.id, + provider: "glm", + connectionId: top.id, + }); + + const selected = await auth.getProviderCredentials("glm", null, [top.id, next.id], "glm-4.6", { + lease: { apiKeyId: key.id, context: context(OWNERS[1], 1), mode: "acquire" }, + materializeCredentials: false, + }); + assert.equal(selected?.connectionId, next.id); + assert.equal((selected as auth.ExclusiveLeaseSelectionResult).exclusiveLease.generation, 1); +}); + +test("all eligible candidates foreign returns WAITING without credentials", async () => { + const [first, second] = await Promise.all([seedConnection(1), seedConnection(2)]); + const key = await seedManagedKey([first.id, second.id]); + leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNERS[0], + apiKeyId: key.id, + provider: "glm", + connectionId: first.id, + }); + leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNERS[1], + apiKeyId: key.id, + provider: "glm", + connectionId: second.id, + }); + + const selected = await auth.getProviderCredentials( + "glm", + null, + [first.id, second.id], + "glm-4.6", + { + lease: { apiKeyId: key.id, context: context(OWNERS[2], 1), mode: "acquire" }, + materializeCredentials: false, + } + ); + assert.equal(selected?.waitingForCapacity, true); + assert.equal(selected?.freeCount, 0); + assert.equal("apiKey" in selected, false); +}); + +test("an eligible live owner binding is reused despite softer priority scoring", async () => { + const [preferred, bound] = await Promise.all([seedConnection(1), seedConnection(2)]); + const key = await seedManagedKey([preferred.id, bound.id]); + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNERS[0], + apiKeyId: key.id, + provider: "glm", + connectionId: bound.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + + const selected = await auth.getProviderCredentials( + "glm", + null, + [preferred.id, bound.id], + "glm-4.6", + { + lease: { + apiKeyId: key.id, + context: context(OWNERS[0], acquired.lease.generation), + mode: "request", + }, + } + ); + assert.equal(selected?.connectionId, bound.id); + assert.match(selected?.apiKey ?? "", /^sk-glm-managed-2-/); +}); + +test("quota-preflight wrapper claims before returning even when live preflight is a no-op", async () => { + const connection = await seedConnection(1); + const key = await seedManagedKey([connection.id]); + const selected = await auth.getProviderCredentialsWithQuotaPreflight( + "glm", + null, + [connection.id], + "glm-4.6", + { + lease: { apiKeyId: key.id, context: context(OWNERS[0], 1), mode: "acquire" }, + materializeCredentials: false, + reserveOAuthSession: false, + } + ); + + assert.equal(selected?.connectionId, connection.id); + assert.equal((selected as auth.ExclusiveLeaseSelectionResult).exclusiveLease.generation, 1); + assert.equal(leaseDb.getActiveExclusiveConnectionLease(OWNERS[0])?.connectionId, connection.id); + assert.equal("apiKey" in selected, false); +}); + +test("lifecycle pre-acquire disables request-scoped OAuth occupancy reservation", async () => { + const connection = await seedConnection(1, { provider: "codex" }); + const key = await seedManagedKey([connection.id]); + const selected = await auth.getProviderCredentialsWithQuotaPreflight( + "codex", + null, + [connection.id], + "gpt-5.6-sol", + { + lease: { apiKeyId: key.id, context: context(OWNERS[0], 1), mode: "acquire" }, + materializeCredentials: false, + reserveOAuthSession: false, + sessionKey: "routing-session-not-owner", + } + ); + + assert.equal(selected?.connectionId, connection.id); + assert.equal( + oauthOccupancy.getForeignOAuthSessionCount(connection.id, "some-other-routing-session"), + 0 + ); +}); + +test("acquire is idempotent without adopting a caller-supplied placeholder generation", async () => { + const connection = await seedConnection(1); + const key = await seedManagedKey([connection.id]); + const first = await auth.getProviderCredentialsWithQuotaPreflight( + "glm", + null, + [connection.id], + "glm-4.6", + { + lease: { apiKeyId: key.id, context: context(OWNERS[0], 1), mode: "acquire" }, + materializeCredentials: false, + reserveOAuthSession: false, + } + ); + const second = await auth.getProviderCredentialsWithQuotaPreflight( + "glm", + null, + [connection.id], + "glm-4.6", + { + lease: { apiKeyId: key.id, context: context(OWNERS[0], 999), mode: "acquire" }, + materializeCredentials: false, + reserveOAuthSession: false, + } + ); + + assert.equal((first as auth.ExclusiveLeaseSelectionResult).exclusiveLease.generation, 1); + assert.equal((second as auth.ExclusiveLeaseSelectionResult).exclusiveLease.generation, 1); +}); + +test("managed request distinguishes missing lease from stale generation", async () => { + const connection = await seedConnection(1); + const key = await seedManagedKey([connection.id]); + const missing = await auth.getProviderCredentials("glm", null, [connection.id], "glm-4.6", { + lease: { apiKeyId: key.id, context: context(OWNERS[0], 1), mode: "request" }, + }); + assert.equal(missing?.leaseRequired, true); + + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNERS[0], + apiKeyId: key.id, + provider: "glm", + connectionId: connection.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + const stale = await auth.getProviderCredentials("glm", null, [connection.id], "glm-4.6", { + lease: { + apiKeyId: key.id, + context: context(OWNERS[0], acquired.lease.generation + 1), + mode: "request", + }, + }); + assert.equal(stale?.leaseFenceStale, true); +}); + +test("unmanaged selection cannot receive lease-only connections even while free", async () => { + const [managed, ordinary] = await Promise.all([seedConnection(1), seedConnection(2)]); + await seedManagedKey([managed.id]); + assert.equal((await apiKeysDb.getExclusiveLeaseConnectionIds()).has(managed.id), true); + + const selected = await auth.getProviderCredentials("glm", null, null, "glm-4.6"); + assert.equal(selected?.connectionId, ordinary.id); +}); + +test("generic lease selection is provider-neutral across GLM and OpenAI fixtures", async () => { + for (const [provider, model] of [ + ["glm", "glm-4.6"], + ["openai", "gpt-4.1"], + ] as const) { + const connection = await seedConnection(1, { provider }); + const key = await seedManagedKey([connection.id]); + const selected = await auth.getProviderCredentials( + provider, + null, + [connection.id], + model, + { + lease: { apiKeyId: key.id, context: context(OWNERS[0], 1), mode: "acquire" }, + materializeCredentials: false, + } + ); + assert.equal(selected?.connectionId, connection.id, provider); + leaseDb.releaseExclusiveConnectionLease({ + leaseOwnerId: OWNERS[0], + generation: selected!.exclusiveLease.generation, + apiKeyId: key.id, + }); + } +}); + +test("managed capacity scales one owner per connection and the next owner waits", async () => { + const connections = await Promise.all( + Array.from({ length: 9 }, (_, index) => seedConnection(index + 1)) + ); + const key = await seedManagedKey(connections.map((connection) => connection.id)); + const selectedIds = new Set(); + for (let index = 0; index < connections.length; index += 1) { + const selected = await auth.getProviderCredentials( + "glm", + null, + connections.map((connection) => connection.id), + "glm-4.6", + { + lease: { apiKeyId: key.id, context: context(OWNERS[index], 1), mode: "acquire" }, + materializeCredentials: false, + } + ); + selectedIds.add(selected?.connectionId ?? ""); + if ([1, 2, 5, 9].includes(index + 1)) { + assert.equal( + selectedIds.size, + index + 1, + `${index + 1} owners must hold distinct connections` + ); + } + } + assert.equal(selectedIds.size, 9); + + const waiting = await auth.getProviderCredentials( + "glm", + null, + connections.map((connection) => connection.id), + "glm-4.6", + { + lease: { apiKeyId: key.id, context: context(OWNERS[9], 1), mode: "acquire" }, + materializeCredentials: false, + } + ); + assert.equal(waiting?.waitingForCapacity, true); +}); + +for (const strategy of ["fill-first", "round-robin", "random", "p2c", "strict-random"]) { + test(`managed filtering preserves the existing ${strategy} selector among FREE candidates`, async () => { + await settingsDb.updateSettings({ fallbackStrategy: strategy }); + const [foreignTop, freeA, freeB] = await Promise.all([ + seedConnection(1), + seedConnection(2), + seedConnection(3), + ]); + const ids = [foreignTop.id, freeA.id, freeB.id]; + const key = await seedManagedKey(ids); + const foreign = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNERS[0], + apiKeyId: key.id, + provider: "glm", + connectionId: foreignTop.id, + }); + assert.equal(foreign.kind, "ACQUIRED"); + + const selected = await auth.getProviderCredentials("glm", null, ids, "glm-4.6", { + lease: { apiKeyId: key.id, context: context(OWNERS[1], 1), mode: "acquire" }, + materializeCredentials: false, + }); + + assert.ok([freeA.id, freeB.id].includes(selected?.connectionId)); + assert.notEqual(selected?.connectionId, foreignTop.id); + }); +} + +test("managed live quota preflight rejects one candidate and claims the next without model execution", async () => { + const [blocked, healthy] = await Promise.all([ + seedConnection(1, { providerSpecificData: { quotaPreflightEnabled: true } }), + seedConnection(2, { providerSpecificData: { quotaPreflightEnabled: true } }), + ]); + const key = await seedManagedKey([blocked.id, healthy.id]); + const calls: string[] = []; + quotaPreflight.registerQuotaFetcher("glm", async (connectionId) => { + calls.push(connectionId); + return { + used: connectionId === blocked.id ? 100 : 20, + total: 100, + percentUsed: connectionId === blocked.id ? 1 : 0.2, + resetAt: new Date(Date.now() + 60_000).toISOString(), + }; + }); + + const selected = await auth.getProviderCredentialsWithQuotaPreflight( + "glm", + null, + [blocked.id, healthy.id], + "glm-4.6", + { + lease: { apiKeyId: key.id, context: context(OWNERS[0], 1), mode: "acquire" }, + materializeCredentials: false, + reserveOAuthSession: false, + } + ); + + assert.deepEqual(calls, [blocked.id, healthy.id]); + assert.equal(selected?.connectionId, healthy.id); + assert.equal(leaseDb.getActiveExclusiveConnectionLease(OWNERS[0])?.connectionId, healthy.id); +}); + +test("managed cached quota ineligibility transitions the live owner to a FREE connection", async () => { + const [bound, free] = await Promise.all([ + seedConnection(1, { + providerSpecificData: { + limitPolicy: { enabled: true, thresholdPercent: 75, windows: ["daily"] }, + }, + }), + seedConnection(2), + ]); + const key = await seedManagedKey([bound.id, free.id]); + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNERS[0], + apiKeyId: key.id, + provider: "glm", + connectionId: bound.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + quotaCache.setQuotaCache(bound.id, "glm", { + daily: { remainingPercentage: 1, resetAt: new Date(Date.now() + 60_000).toISOString() }, + }); + + const selected = await auth.getProviderCredentials("glm", null, [bound.id, free.id], "glm-4.6", { + lease: { + apiKeyId: key.id, + context: context(OWNERS[0], acquired.lease.generation), + mode: "request", + }, + materializeCredentials: false, + }); + + assert.equal(selected?.connectionId, free.id); + assert.equal( + leaseDb.getActiveExclusiveConnectionLease(OWNERS[0])?.generation, + acquired.lease.generation + ); +}); + +test("managed cooldown and terminal-auth ineligibility transition only to a FREE connection", async () => { + for (const kind of ["cooldown", "terminal"] as const) { + await resetStorage(); + const [bound, free] = await Promise.all([seedConnection(1), seedConnection(2)]); + const key = await seedManagedKey([bound.id, free.id]); + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNERS[0], + apiKeyId: key.id, + provider: "glm", + connectionId: bound.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + await providersDb.updateProviderConnection( + bound.id, + kind === "cooldown" + ? { + testStatus: "unavailable", + rateLimitedUntil: new Date(Date.now() + 60_000).toISOString(), + } + : { testStatus: "banned" } + ); + + const selected = await auth.getProviderCredentials( + "glm", + null, + [bound.id, free.id], + "glm-4.6", + { + lease: { + apiKeyId: key.id, + context: context(OWNERS[0], acquired.lease.generation), + mode: "request", + }, + materializeCredentials: false, + } + ); + assert.equal(selected?.connectionId, free.id, kind); + } +}); + +test("managed model lockout transitions the same generation to a FREE connection", async () => { + const [bound, free] = await Promise.all([ + seedConnection(1, { provider: "gemini" }), + seedConnection(2, { provider: "gemini" }), + ]); + const key = await seedManagedKey([bound.id, free.id]); + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNERS[0], + apiKeyId: key.id, + provider: "gemini", + connectionId: bound.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + await auth.markAccountUnavailable( + bound.id, + 429, + "synthetic model lockout", + "gemini", + "gemini-2.5-pro" + ); + + const selected = await auth.getProviderCredentials( + "gemini", + null, + [bound.id, free.id], + "gemini-2.5-pro", + { + lease: { + apiKeyId: key.id, + context: context(OWNERS[0], acquired.lease.generation), + mode: "request", + }, + materializeCredentials: false, + } + ); + assert.equal(selected?.connectionId, free.id); + assert.equal( + leaseDb.getActiveExclusiveConnectionLease(OWNERS[0])?.generation, + acquired.lease.generation + ); +}); + +test("managed request invalidates an unsafe binding when no FREE failover target exists", async () => { + const bound = await seedConnection(1); + const key = await seedManagedKey([bound.id]); + const acquired = leaseDb.acquireExclusiveConnectionLease({ + leaseOwnerId: OWNERS[0], + apiKeyId: key.id, + provider: "glm", + connectionId: bound.id, + }); + assert.equal(acquired.kind, "ACQUIRED"); + if (acquired.kind !== "ACQUIRED") return; + await providersDb.updateProviderConnection(bound.id, { + testStatus: "unavailable", + rateLimitedUntil: new Date(Date.now() + 60_000).toISOString(), + }); + + const selected = await auth.getProviderCredentials("glm", null, [bound.id], "glm-4.6", { + lease: { + apiKeyId: key.id, + context: context(OWNERS[0], acquired.lease.generation), + mode: "request", + }, + materializeCredentials: false, + }); + + assert.equal(selected?.allRateLimited, true); + assert.equal(leaseDb.getActiveExclusiveConnectionLease(OWNERS[0]), null); +}); diff --git a/tests/unit/warmupScheduler.test.ts b/tests/unit/warmupScheduler.test.ts index 97202365c8..7b77d00a85 100644 --- a/tests/unit/warmupScheduler.test.ts +++ b/tests/unit/warmupScheduler.test.ts @@ -19,6 +19,7 @@ const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-warmup-or process.env.DATA_DIR = TEST_DATA_DIR; process.env.NODE_ENV = "test"; process.env.DISABLE_SQLITE_AUTO_BACKUP = "true"; +process.env.API_KEY_SECRET = "warmup-exclusive-lease-test-secret"; const core = await import("../../src/lib/db/core.ts"); const providersDb = await import("../../src/lib/db/providers.ts"); @@ -193,6 +194,41 @@ test("integration: opted-in claude_pro connection → fetch fires with Bearer to delete process.env.OMNIROUTE_WARMUP_CRON; }); +test("hard lease isolation skips an opted-in lease-only connection with zero model calls", async () => { + const { startWarmupScheduler, stopWarmupScheduler } = + await import("../../src/lib/warmupScheduler.ts"); + const settingsDb = await import("../../src/lib/db/settings.ts"); + const apiKeysDb = await import("../../src/lib/db/apiKeys.ts"); + + const conn = await providersDb.createProviderConnection({ + provider: "claude", + authType: "oauth", + name: "Managed Pro User", + accessToken: "synthetic-token", + refreshToken: "synthetic-refresh", + isActive: true, + providerSpecificData: { organizationType: "claude_pro" }, + }); + await apiKeysDb.createApiKey("managed warmup key", "test", ["lease:exclusive"], { + allowedConnections: [conn.id], + }); + await settingsDb.updateSettings({ claudeWarmup: { connections: { [conn.id]: true } } }); + + const mock = installMockFetch(() => { + throw new Error("unexpected model warmup call"); + }); + process.env.OMNIROUTE_WARMUP_ENABLED = "1"; + process.env.OMNIROUTE_WARMUP_CRON = "*/1 * * * *"; + startWarmupScheduler(); + await new Promise((resolve) => setTimeout(resolve, 50)); + stopWarmupScheduler(); + mock.restore(); + + assert.equal(mock.calls.length, 0); + delete process.env.OMNIROUTE_WARMUP_ENABLED; + delete process.env.OMNIROUTE_WARMUP_CRON; +}); + test("integration: message rotation — different content across sequential pings", async () => { const { startWarmupScheduler, stopWarmupScheduler } = await import("../../src/lib/warmupScheduler.ts");