Skip to content

refactor(ai-gateway): split the LLM proxy handler into phases - #9717

Open
chrarnoldus wants to merge 11 commits into
mainfrom
ai-gateway/llm-proxy-phases
Open

chrarnoldus wants to merge 11 commits into
mainfrom
ai-gateway/llm-proxy-phases

Conversation

@chrarnoldus

Copy link
Copy Markdown
Contributor

Summary

handleLlmProxyRequest was a single ~1,090-line function and the gateway's top churn hotspot. This splits it into phases. Each extracted phase returns { ok: true, ... } | { ok: false, response }, except model resolution, which returns the result type described below. No generic pipeline abstraction is introduced.

The auth, balance, group-policy, guardrail and bouncer promises still start at the same points and overlap as before. Auth still starts before the body is parsed, so it overlaps JSON.parse of large bodies. Several invalid requests now return before anything awaits auth, so a no-op catch keeps an auth failure on those paths from becoming an unhandled rejection. Malformed JSON already had this exposure.

Phases

Phase Where Notes
Ingress inline in llm-proxy.ts path, client IP, country, account rate limit
Parse and validate handlers/llm-proxy-request.ts → parseLlmProxyRequest pure, no I/O: JSON, providerOptions, model, messages/input, provider privacy, missing client IP, Responses store/previous_response_id. Also readLlmProxyHeaders, which reads each header once.
Context inline auth, balance, group policy, guardrail start, disconnect logging, effective privacy, auto-routing provider hints
Model resolution handlers/llm-proxy-model-resolution.ts → resolveModel Kilo Auto ids, custom routers, signed x-kilo-auto-decision, catalog-mode mapping. Also autoRoutingFailureResponse and billLegacyClassifierCost.
Admission inline custom LLM lookup, free-model rate limit, auth outcome and anonymous handling, bouncer start, legacy classifier billing
Provider and policy handlers/llm-proxy-provider-policy.ts → resolveProviderAndPolicy non-trial enterprise lookup, getProvider, Organization Auto target checks, API kind, unavailable model, balance, org access policy, data-collection check, custom LLM provider intersection
Upstream inline usage context, bouncer and guardrail verdicts, sendUpstreamAttempt
Respond handlers/llm-proxy-respond.ts → respondToUpstream metrics, upstream error policy, accounting, briefing receipt, makeErrorReadable, rewriteModelResponse

The orchestrator is about 430 lines, more than the ~150 originally estimated. Ingress, context, admission and upstream stay inline because extracting them would only move code around. Most of the remaining length is the auth outcome, balance and usage-context code. A parallel PR that adds shared admission helpers (resolveGatewayCaller) replaces that code, so this PR leaves it as is to keep the two PRs from conflicting.

Auto-routing result type

The five mutable auto-routing variables (autoModel, routingTarget, legacyClassifierCostUsd, autoRoutingUnavailable, isAutoRoutedRequest) and the closures that wrote them are gone. resolveModel returns one discriminated result:

type ModelResolution =
  | { kind: 'direct' }
  | { kind: 'routed'; autoModel: string; routingTarget: string | null; poolRouted: boolean; classifierCostUsd: number }
  | { kind: 'rejected'; failure: AutoRoutingFailure; autoModel: string | null; classifierCostUsd: number };

AutoRoutingFailure names each failure: invalid_token, auth_required, decision_secret_missing, invalid_auto_decision, no_free_models_available, organization_auto_configuration, custom_router or unavailable. The signed x-kilo-auto-decision branch returns the same type: routed with poolRouted: false and no classifier cost, or rejected.

unavailable is answered later than the other failures, as before. It yields to the caller's auth failure, and the legacy classifier cost it carries is billed first. The other failures are answered right after resolution. Legacy positive-cost classifier billing still works and reads its cost from the result. Deleting it is a separate change that needs the deployed worker versions confirmed first. The provider and policy phase reads routingTarget and poolRouted from the result.

applyResolvedAutoModel is unchanged. Its signature already fits the phase, and keeping it avoids rewriting the many handler tests that mock it.

Other changes (one commit each)

  • Removed the settings deny-list branch from autoRoutingDenials along with its candidateIds parameter. It could never run: group policy is null only when the caller has no organization, and then plan is undefined. Removed the test that reached it only through impossible mocks.
  • Removed the max-tokens guard (MAX_TOKENS_LIMIT = 99999999999). It logged the full prompt through console.warn, which Sentry Logs collects. It also returned a retryable 503 for a client error, and with that limit it protected nothing. getMaxTokens is now called once.
  • The session id (taskId ?? sessionHeader) and the editor-name header are each read once.
  • Removed the effectiveProviderContext alias and the single-use createAccessCheckResolver factory, along with stale comments.
  • validatePath accepts only /api/v1/.... apps/web/ai-gateway-rewrites.mjs always rewrites /api/gateway and /api/openrouter to /api/v1, and apps/ai-gateway has only /api/v1 routes. invalidPathResponse now lists all three paths in its message; the status and error_type are unchanged. The tests' default prefix is /api/v1, a new test covers each rejected legacy prefix, and apps/ai-gateway/AGENTS.md is updated.

Error precedence changes

These matter only when several errors apply to one request:

  • Responses store / previous_response_id now returns 400 before any auth, auto-routing or rate-limit error. Before, it ran after auto-routing, so a rejected request could still call the paid classifier. The free-model rate limit 429, invalid-token and paid-model-auth 401s, auto-routing errors and the missing-IP 400 also came before it. A new test asserts that the classifier is not called.
  • Missing client IP now returns 400 before auth failures, auto-routing errors and the classifier call (new test). Before, it ran after model resolution.
  • Invalid provider privacy fields already returned 400 before any other error could occur, so the only change is that the check no longer waits for the balance read. A failed balance read no longer comes first.
  • Data-collection and custom LLM provider checks now run before the usage context is built and before console.debug and setTag. The response is unchanged.
  • isUnavailableModel stays where it was, after the free-model rate limit, because it does I/O. The test that expects the Gemma free-model 429 before the unavailable-model error still passes unchanged.

Owner decision, deferred

body.provider is still replaced wholesale by the organization's effective provider config when organization policy applies (request.body.provider = effectiveProviderConfig in resolveProviderAndPolicy). Tests lock this behavior in, and whether to merge the caller's provider object instead is an owner decision, so this PR does not change it.

Not touched

  • providers/get-provider.ts, which a parallel PR is refactoring.
  • The auth outcome normalization, beyond what the phase split needed.

Verification

  • pnpm --filter @kilocode/web-shared typecheck: passes
  • llm-proxy.test.ts: 219 tests pass. Changes are the removed dead-branch test, the /api/v1 default prefix and legacy-prefix test, and new precedence tests.
  • Every gateway suite under the apps/web Jest config (packages/web-shared/src/lib/ai-gateway/**, including auto-model/*, auto-routing-*, and the apps/ai-gateway route tests): 164 suites pass. One DB-backed suite (auto-free-config.test.ts) timed out during the parallel run, which shares a local database with other sessions, and passed when rerun on its own.
  • pnpm exec oxfmt --check and oxlint --config .oxlintrc.json on changed files: clean. git diff --check: clean.

Models

Written with kilo/anthropic/claude-opus-5.5. A review subagent ran on the same model with no override.

@chrarnoldus chrarnoldus self-assigned this Nov 3, 2026
@kilo-code-bot

kilo-code-bot Bot commented Nov 3, 2026 •

Copy link
Copy Markdown
Contributor

Code Review Summary

Status: No Issues Found | Recommendation: Merge

Executive Summary

The phase split is behaviorally faithful to the previous handler: I traced every extracted statement (parse/validate, model resolution, provider and policy, respond) against the base revision and found no dropped checks, reordered side effects that change responses, or new unhandled promises.

Assumptions: [INFERENCE] The tightened path validation is safe because apps/ai-gateway has no /api/gateway or /api/openrouter routes and apps/web/ai-gateway-rewrites.mjs rewrites every legacy prefix to /api/v1; no in-repo caller reaches the handler with a legacy pathname. [INFERENCE] Error-precedence changes (Responses store/previous_response_id, missing client IP, provider privacy) and the removed max-tokens guard intentionally change 4xx/503 ordering, per the PR description; tests cover the classifier-not-called cases. [INFERENCE] Removed settings deny-list branch is unreachable: getBalanceAndOrgSettings returns no plan/settings without an organization, and an enterprise org with authFailedResponse is rejected before classifier denials are consumed. I did not execute the Jest suites (read-only review), so test results are taken from the PR description.

Files Reviewed (8 files)
  • apps/ai-gateway/AGENTS.md
  • packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy-request.ts (new)
  • packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy-model-resolution.ts (new)
  • packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy-provider-policy.ts (new)
  • packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy-respond.ts (new)
  • packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.ts
  • packages/web-shared/src/lib/ai-gateway/handlers/llm-proxy.test.ts
  • packages/web-shared/src/lib/ai-gateway/llm-proxy-helpers.ts

Reviewed by deepseek-flash-latest · Input: 134.5K · Output: 31.2K · Cached: 2.9M

Review guidance: REVIEW.md from base branch main

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant