Repository navigation
fix(gateway): serve and expose Vertex service tiers - #2819
Conversation
WalkthroughThe PR centralizes service-tier logging in ChangesService-tier logging and error test
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e3f078a771
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // rejections, cache hits, streaming/upstream errors, fetch errors) | ||
| // records them. Explicit values in logData still win. | ||
| requestedServiceTier, | ||
| usedServiceTier: servedServiceTier, |
There was a problem hiding this comment.
Reset served tier between retry attempts
When a fallback retry follows an attempt that already populated servedServiceTier, this default can be stale for any later attempt that fails before fetch returns (timeout/connection error), because the retry-loop reset clears res/fetchError but not servedServiceTier. In that scenario the timeout/error log for the later provider records the previous provider's usedServiceTier even though no tier was served for that attempt; reset the served tier at the start of each attempt or only default it when it belongs to the current response.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@apps/gateway/src/api.spec.ts`:
- Around line 914-920: The test is assuming waitForLogs(1) returns only the
current request’s row, but it actually returns all log rows once the threshold
is reached. Update the spec in api.spec.ts to locate the request-specific entry
before asserting hasError/requestedServiceTier/usedServiceTier, using a
request-correlated field or explicit filter instead of logs[0]. Keep the
assertions tied to the log created by this request rather than relying on array
length or ordering.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 84eda595-2fdc-4a1c-aa2e-a2a55b0e7cc2
📒 Files selected for processing (2)
apps/gateway/src/api.spec.tsapps/gateway/src/chat/chat.ts
| const logs = await waitForLogs(1); | ||
| expect(logs.length).toBe(1); | ||
| expect(logs[0].hasError).toBe(true); | ||
| // The upstream errored before serving a tier, but the requested tier is | ||
| // still threaded onto the error log via the insertLogEntry wrapper. | ||
| expect(logs[0].requestedServiceTier).toBe("priority"); | ||
| expect(logs[0].usedServiceTier).toBeNull(); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Target the log row for this request instead of assuming waitForLogs(1) returns only it.
waitForLogs returns every row from db.query.log.findMany({}) once the count threshold is met, so logs.length === 1 and logs[0] are brittle if another log already exists or is written during this test. Query or filter for the request-specific log before asserting tier fields. Based on apps/gateway/src/test-utils/test-helpers.ts:17-47, waitForLogs does not filter or order log rows.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@apps/gateway/src/api.spec.ts` around lines 914 - 920, The test is assuming
waitForLogs(1) returns only the current request’s row, but it actually returns
all log rows once the threshold is reached. Update the spec in api.spec.ts to
locate the request-specific entry before asserting
hasError/requestedServiceTier/usedServiceTier, using a request-correlated field
or explicit filter instead of logs[0]. Keep the assertions tied to the log
created by this request rather than relying on array length or ordering.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 808bd791be
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| token: "google-test-key", | ||
| provider: "google-vertex", | ||
| organizationId: "org-id", | ||
| baseUrl: mockServerUrl, |
There was a problem hiding this comment.
Use a tier-eligible key in the new test
With the harness seeding the project in api-keys mode (apps/gateway/src/test-utils/gateway-api-test-harness.ts:65-67), this Google Vertex key is treated as a custom/proxy base URL. Premium service_tier requests then run through the eligibility check (apps/gateway/src/chat/chat.ts:2695-2715), and providerKeyBaseUrlSupportsServiceTier rejects Vertex keys unless the base URL is null or the canonical Google upstream, so the request is rejected before hitting TRIGGER_ERROR and no log row is inserted; waitForLogs(1) will fail instead of validating the upstream-error path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 18a98fb5ed
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| requestedServiceTier, | ||
| usedServiceTier: servedServiceTier, |
There was a problem hiding this comment.
Thread tiers through content-filter logs
When contentFilterMode === "enabled" blocks a request that included service_tier, the log still bypasses these new defaults because the block path later calls the local insertLog directly at apps/gateway/src/chat/chat.ts:5029 instead of insertLogEntry. In that scenario no upstream tier is served, but the activity row still loses the requested tier, so the patch does not cover this gateway-rejection log path despite the wrapper comment saying every log path records it; use insertLogEntry there or pass the two tier fields explicitly.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@apps/gateway/src/api.spec.ts`:
- Around line 912-920: The test in api.spec.ts is only asserting a non-200
response, which can still pass for local rejections and no longer distinguishes
upstream failures because insertLogEntry now always adds requestedServiceTier
and usedServiceTier. Update the assertion around the request/response flow in
this test to check the concrete upstream-failure signal first, such as the exact
500 status, the response error type, or the log finishReason, and only then
assert the tier fields on the log entry.
In `@apps/gateway/src/test-utils/mock-openai-server.ts`:
- Around line 589-592: The mock OpenAI server is using unsafe `as any` casts
when setting HTTP status codes; update the `statusTrigger` handling in
`mock-openai-server` to use a proper `StatusCode` type from `hono/context`
instead. Replace each `c.status(statusTrigger.statusCode as any)` occurrence in
the relevant response paths with a type-safe `StatusCode`-compatible cast or
numeric handling, and make sure the `extractStatusCodeTrigger` flow still
accepts any three-digit status code. Also extend the TypeError-related response
typing in the same module so the new status-code typing is reflected
consistently across the mock server helpers.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 71ff89c7-b962-43a5-b717-7ed0b31f12a9
📒 Files selected for processing (3)
apps/gateway/src/api.spec.tsapps/gateway/src/chat/chat.tsapps/gateway/src/test-utils/mock-openai-server.ts
| const statusTrigger = extractStatusCodeTrigger(userMessage); | ||
| if (statusTrigger) { | ||
| c.status(statusTrigger.statusCode as any); | ||
| return c.json(statusTrigger.errorResponse); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor
🧩 Analysis chain
🏁 Script executed:
cat -n apps/gateway/src/test-utils/mock-openai-server.ts | sed -n '580,600p'Repository: theopenco/llmgateway
Length of output: 1146
🏁 Script executed:
cat -n apps/gateway/src/test-utils/mock-openai-server.ts | sed -n '1,100p'Repository: theopenco/llmgateway
Length of output: 3319
🏁 Script executed:
cat -n apps/gateway/src/test-utils/mock-openai-server.ts | sed -n '100,200p'Repository: theopenco/llmgateway
Length of output: 3134
🏁 Script executed:
cat -n apps/gateway/src/test-utils/mock-openai-server.ts | sed -n '585,600p'Repository: theopenco/llmgateway
Length of output: 850
🏁 Script executed:
cat -n apps/gateway/src/test-utils/mock-openai-server.ts | sed -n '570,595p'Repository: theopenco/llmgateway
Length of output: 1314
🏁 Script executed:
cd apps/gateway && grep -n "status(" src/test-utils/mock-openai-server.ts | head -10Repository: theopenco/llmgateway
Length of output: 540
🏁 Script executed:
head -30 apps/gateway/src/test-utils/mock-openai-server.tsRepository: theopenco/llmgateway
Length of output: 864
🏁 Script executed:
cd apps/gateway && cat package.json | jq '.dependencies.hono'Repository: theopenco/llmgateway
Length of output: 169
🏁 Script executed:
cd apps/gateway && grep -n "import.*Context" src/test-utils/mock-openai-server.tsRepository: theopenco/llmgateway
Length of output: 197
🏁 Script executed:
cd apps/gateway && grep -A 2 -B 2 "Status" node_modules/hono/dist/types/context.d.ts 2>/dev/null | head -20 || grep -A 2 -B 2 "status" node_modules/hono/dist/hono-base.js 2>/dev/null || echo "Checking types..."Repository: theopenco/llmgateway
Length of output: 158
🏁 Script executed:
cd apps/gateway && grep -n "StatusCode" src/test-utils/mock-openai-server.tsRepository: theopenco/llmgateway
Length of output: 1449
Eliminate as any cast for HTTP status codes. Extend TypeError error types.
File: apps/gateway/src/test-utils/mock-openai-server.ts uses a fallback for any three-digit status code, so restricting the TRIGGER_STATUS_NNN regex is unnecessary. However, the project's TypeScript guidelines forbid the use of any. Replace c.status(statusTrigger.statusCode as any) in three locations (lines 591, 767, 1141, etc.) with c.status(statusTrigger.statusCode as unknown as StatusCode) or c.status(statusTrigger.statusCode as number), and import StatusCode from hono/context.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@apps/gateway/src/test-utils/mock-openai-server.ts` around lines 589 - 592,
The mock OpenAI server is using unsafe `as any` casts when setting HTTP status
codes; update the `statusTrigger` handling in `mock-openai-server` to use a
proper `StatusCode` type from `hono/context` instead. Replace each
`c.status(statusTrigger.statusCode as any)` occurrence in the relevant response
paths with a type-safe `StatusCode`-compatible cast or numeric handling, and
make sure the `extractStatusCodeTrigger` flow still accepts any three-digit
status code. Also extend the TypeError-related response typing in the same
module so the new status-code typing is reflected consistently across the mock
server helpers.
Source: Coding guidelines
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7b57be9e55
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // Read at call time, so the streaming final usage chunk reflects the tier | ||
| // resolved from the upstream response. | ||
| requestedServiceTier, | ||
| usedServiceTier: servedServiceTier, |
There was a problem hiding this comment.
Preserve cached stream tier metadata
When response caching is enabled, replaying a cached streamed request with service_tier: "flex"/"priority" never contacts upstream, so servedServiceTier is still null here. The cached-stream replay builds metadata with this value (see buildFinalResponseMetadata at the cache replay) and merges it after the cached chunk metadata, which overwrites the original used_service_tier from the cached final usage chunk with null; callers then see a false downgrade on cache hits. Preserve the cached tier metadata or avoid overwriting used_service_tier when no new upstream response was served.
Useful? React with 👍 / 👎.
Vertex Flex/Priority PayGo needs two headers, but the gateway only sent one. `X-Vertex-AI-LLM-Shared-Request-Type` selects the tier; `X-Vertex-AI-LLM-Request-Type: shared` forces the shared PayGo path. Without the latter, a project with Provisioned Throughput allocated consumes PT first and never reaches the shared tier, so the request is silently served (and billed) as standard even though Priority was requested. Send both headers for flex/priority (no-op on projects without PT). Also surface and log the requested vs served tier: - Add requested_service_tier + used_service_tier to the response metadata (non-streaming) and the streaming final usage chunk, via buildFinalResponseMetadata / toResponseMetadataExtras. Emitted only when a premium tier was requested; used is null on a downgrade. - Default both fields in the insertLogEntry wrapper so every chat log path (guardrail/validation rejections, cache hits, streaming/upstream errors, fetch errors) records them; reset servedServiceTier per retry attempt so a fallback that fails before fetch returns logs no served tier. Tests/infra: - get-provider-headers test asserts both Vertex headers on flex/priority. - Upstream-error test uses openai/gpt-5.5 (not subject to #2821's upstream base-URL rule); mock /v1/responses handler is array-safe for TRIGGER_ERROR / status triggers and echoes service_tier when streaming. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0e80b1e7bb
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| ...(extras.requestedServiceTier | ||
| ? { | ||
| requested_service_tier: extras.requestedServiceTier, | ||
| used_service_tier: extras.usedServiceTier ?? null, |
There was a problem hiding this comment.
Declare tier metadata in the response schema
When this starts emitting requested_service_tier and used_service_tier, the /v1/chat/completions OpenAPI response schema still declares only the existing metadata fields in apps/gateway/src/chat/chat.ts:1335-1347. Generated docs/clients therefore won't expose the new API fields even though runtime responses contain them; add both fields there, with used_service_tier nullable for downgrades.
Useful? React with 👍 / 👎.
The fix: Vertex Priority/Flex PayGo needs two headers
Vertex requires both:
X-Vertex-AI-LLM-Shared-Request-Type: priority|flex— selects the tier (the gateway already sent this)X-Vertex-AI-LLM-Request-Type: shared— forces the shared PayGo path (was missing)Without the second header, a project that has Provisioned Throughput allocated consumes PT first and never reaches the shared Flex/Priority tier — so the request is silently served (and billed) as standard even though
prioritywas requested. This is why production showedrequested=priority/used=nullonaiplatform.googleapis.comprojects with PT, while projects without PT worked. We now send both headers for flex/priority (no-op on projects without PT, so it's safe everywhere).Per Google's docs and liteLLM: "To bypass Provisioned Throughput entirely, also set requestType: 'shared'."
Also in this PR: surface & log requested vs served tier
requested_service_tier+used_service_tieradded to the responsemetadata(non-streaming) and the streaming final usage chunk (viabuildFinalResponseMetadata/toResponseMetadataExtras). Emitted only when a premium tier was requested;usedisnullon a downgrade so callers can detect it.insertLogEntrywrapper so every chat log path (guardrail/validation rejections, cache hits, streaming/upstream errors, fetch errors) records them; resetservedServiceTierper retry attempt so a fallback that fails beforefetchreturns logs no served tier instead of inheriting the previous provider's.Tests / infra
get-provider-headerstest asserts both Vertex headers on flex/priority and their absence otherwise.openai/gpt-5.5(not subject to feat(gateway): gate service tiers to upstream provider keys #2821's upstream base-URL rule) so it still reaches the error path; mock/v1/responseshandler is array-safe forTRIGGER_ERROR/TRIGGER_STATUS_*and echoesservice_tierwhen streaming.Verification
pnpm test:unit(132 files), fullapps/gateway/src/api.spec.ts(101 passed / 2 skipped),pnpm format, fullpnpm buildall pass. The two-header behavior was confirmed against Google's Priority PayGo docs.🤖 Generated with Claude Code