Repository navigation
feat(gateway): gate service tiers to upstream provider keys - #2821
Conversation
|
Warning Review limit reached
More reviews will be available in 21 minutes and 21 seconds. Learn how PR review limits work. Your organization has used up its prepaid credits, and credit purchases are no longer available. Enable the review add-on in the billing tab to keep reviews running — you're only billed for reviews past your plan's rate limits ($0.25/file). ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based credits. 🚦 How do rate limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (7)
WalkthroughThe PR adds base-URL eligibility checks for Flex/Priority service tiers, applies them in gateway routing and provider-context selection, and updates the service-tiers documentation with Google provider rules. ChangesService-tier provider-key eligibility
Sequence Diagram(s)sequenceDiagram
participant chat.ts
participant resolve-provider-context.ts
participant apply-google-service-tier.ts
participant get-provider-env.ts
chat.ts->>apply-google-service-tier.ts: check provider key base URL eligibility
chat.ts->>get-provider-env.ts: exclude ineligible env credential indices
resolve-provider-context.ts->>apply-google-service-tier.ts: filter api-keys and hybrid selections
resolve-provider-context.ts->>get-provider-env.ts: exclude ineligible env credential indices
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 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: cfb283ca11
ℹ️ 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".
| const envEligible = | ||
| (project.mode === "credits" || project.mode === "hybrid") && | ||
| hasProviderEnvironmentToken(providerId as Provider); |
There was a problem hiding this comment.
Check env base URLs before tier routing
When a credits or hybrid project uses LLM_GOOGLE_AI_STUDIO_BASE_URL or LLM_GOOGLE_VERTEX_BASE_URL to point Google env credentials at a proxy, this new branch still treats the provider as service-tier eligible solely because the API-key env var exists. Those base URL env vars are supported in packages/models/src/providers.ts and are used by getProviderEndpoint, so a Flex/Priority request can still be sent through the proxy and silently downgraded—the case this change is trying to prevent. Gate envEligible with the same canonical-base-URL check against the env base URL, or reject env fallback when that override is non-canonical.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
apps/gateway/src/chat/tools/resolve-provider-context.ts (1)
313-320: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: reuse the shared tier predicate for consistency.
chat.tsgates the identical filter viaisRequestedServiceTier(service_tier)while here the"flex" | "priority"check is inlined. Sharing one predicate keeps the two service-tier paths from drifting if accepted tier values change. Only worth it ifisRequestedServiceTieris already exported/importable.🤖 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/chat/tools/resolve-provider-context.ts` around lines 313 - 320, The service-tier gating in resolve-provider-context should use the shared isRequestedServiceTier predicate instead of inlining the "flex" | "priority" check so it stays consistent with chat.ts. Update the serviceTierKeyFilter logic in resolveProviderContext to call the shared predicate if it is exported/importable, and keep providerKeyBaseUrlSupportsServiceTier as the actual filter body.
🤖 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 `@packages/actions/src/apply-google-service-tier.ts`:
- Around line 47-49: The trailing-slash cleanup in normalizeServiceTierBaseUrl
uses a superlinear regex that CodeQL flags, so replace the /\/+$/ logic with a
non-regex approach that removes trailing slashes in linear time while preserving
the existing trim and lowercase normalization. Keep the fix localized to
normalizeServiceTierBaseUrl in apply-google-service-tier so the stored provider
key normalization still produces the same canonical result without backtracking
risk.
---
Nitpick comments:
In `@apps/gateway/src/chat/tools/resolve-provider-context.ts`:
- Around line 313-320: The service-tier gating in resolve-provider-context
should use the shared isRequestedServiceTier predicate instead of inlining the
"flex" | "priority" check so it stays consistent with chat.ts. Update the
serviceTierKeyFilter logic in resolveProviderContext to call the shared
predicate if it is exported/importable, and keep
providerKeyBaseUrlSupportsServiceTier as the actual filter body.
🪄 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: 36593ffd-b971-43a7-a303-850c95c887e3
📒 Files selected for processing (5)
apps/docs/content/features/service-tiers.mdxapps/gateway/src/chat/chat.tsapps/gateway/src/chat/tools/resolve-provider-context.tspackages/actions/src/apply-google-service-tier.spec.tspackages/actions/src/apply-google-service-tier.ts
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9a98c7ab6f
ℹ️ 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".
| if (project.mode === "api-keys") { | ||
| return hasCompliantDbKey; | ||
| } | ||
| return hasCompliantDbKey || envEligible; |
There was a problem hiding this comment.
Ignore DB keys in credits-mode tier eligibility
In credits mode this lets an unrelated compliant DB key make the provider look service-tier eligible, even though credits-mode resolution never uses DB provider keys and later only calls getProviderEnv(...). With a credits project where the Google env credential is missing or every env base URL is a proxy, but the org also has a canonical DB key, a pinned/selected Priority request passes this guard and then fails during env selection instead of being filtered or returning the intended 400. Return only envEligible for project.mode === "credits".
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
🧹 Nitpick comments (1)
apps/gateway/src/chat/tools/resolve-provider-context.ts (1)
316-323: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider reusing
isRequestedServiceTierfor the tier check.The chat.ts variant of this filter uses
isRequestedServiceTier(service_tier), whereas here theflex/priorityset is inlined. Centralizing on the helper avoids drift if the set of "requested" tiers changes later.🤖 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/chat/tools/resolve-provider-context.ts` around lines 316 - 323, The service tier filter in resolve-provider-context is inlining the requested tier check instead of using the shared helper. Update the serviceTierKeyFilter logic to reuse isRequestedServiceTier for the service_tier condition, matching the chat.ts implementation and keeping the requested-tier set centralized in one place.
🤖 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.
Nitpick comments:
In `@apps/gateway/src/chat/tools/resolve-provider-context.ts`:
- Around line 316-323: The service tier filter in resolve-provider-context is
inlining the requested tier check instead of using the shared helper. Update the
serviceTierKeyFilter logic to reuse isRequestedServiceTier for the service_tier
condition, matching the chat.ts implementation and keeping the requested-tier
set centralized in one place.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: be2f7b4e-4bde-4d57-9943-2b3595a4b7af
📒 Files selected for processing (4)
apps/gateway/src/chat/chat.tsapps/gateway/src/chat/tools/get-provider-env.spec.tsapps/gateway/src/chat/tools/get-provider-env.tsapps/gateway/src/chat/tools/resolve-provider-context.ts
🚧 Files skipped from review as they are similar to previous changes (1)
- apps/gateway/src/chat/chat.ts
9a98c7a to
25d006c
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 25d006cf79
ℹ️ 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".
| const baseUrl = getProviderEnvValue(provider, "baseUrl", index); | ||
| if (!providerKeyBaseUrlSupportsServiceTier(provider, baseUrl)) { | ||
| ineligible.add(index); |
There was a problem hiding this comment.
Exclude non-global Vertex env indices for tiers
When google-vertex env credentials are comma-indexed by both base URL and region, filtering only on baseUrl can move a Flex/Priority request onto a non-global Vertex index. For example, if index 0 is baseUrl=proxy with region=global and index 1 is the canonical/default base URL with region=us-central1, the earlier service-tier support check can pass using config index 0, but this helper excludes index 0 so getProviderEnv selects index 1; later getForwardedServiceTier(..., configIndex) sees us-central1 and omits the Vertex tier header, sending the request as standard instead of rejecting or routing elsewhere. Treat non-global Vertex env indices as ineligible for service-tier selection as well.
Useful? React with 👍 / 👎.
25d006c to
dc76ac5
Compare
Flex/Priority is only honored when the request reaches Google directly. A google-vertex/google-ai-studio credential pointing at a custom base URL (proxy) may silently drop the tier header (Vertex) or body field (AI Studio) and be billed/reported as standard. Route service-tier requests only to credentials targeting the canonical upstream, covering both DB provider keys and env credentials: - providerKeyBaseUrlSupportsServiceTier(): a credential is eligible when it has no custom base URL (managed default) or its base URL matches the provider's canonical upstream. Empty/unset base URL = default = eligible. Non-google providers are unaffected. - enforceServiceTierKeyEligibility() excludes ineligible providers from the routing candidate set (initial selection + selectNextProvider fallback), mirroring the compliance policy: auto/model-id routing drops them and only fails when none remain; a pinned provider with no eligible credential returns a clear 400 instead of downgrading. - findProviderKey filter skips non-compliant DB keys (hybrid falls back to env); getProviderEnv excludes env indices whose per-index base URL is a custom proxy. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
dc76ac5 to
b87aeb4
Compare
Follow-up to #2816, which populated requestedServiceTier / usedServiceTier only on the chat success paths and the two direct _insertLog rejections. Other paths (guardrail/validation rejections, cache hits, streaming timeouts/upstream errors, fetch errors) go through the shared insertLogEntry wrapper and omitted both fields. - Default requestedServiceTier + usedServiceTier (from servedServiceTier) in the insertLogEntry wrapper so every chat log path records them; explicit values still win. Declare servedServiceTier ahead of the wrapper and drop the now-redundant explicit fields from the two success inserts. - Reset servedServiceTier at the start of each retry attempt (both the streaming and non-streaming loops) so a fallback that fails before fetch returns logs no served tier instead of inheriting the previous provider's. - Update the upstream-error test to openai/gpt-5.5 (priority, not subject to the upstream base-URL restriction added in #2821) so it still reaches the error path, and make the mock /v1/responses handler detect TRIGGER_ERROR / status triggers in array-form input. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Surface the requested vs served processing tier (Flex/Priority) on both the response and every chat log path. Response: - Add requested_service_tier + used_service_tier to the response metadata for non-streaming completions and the streaming final usage chunk (via buildFinalResponseMetadata / toResponseMetadataExtras). Emitted only when a premium tier was requested; used_service_tier is null when the upstream downgraded to standard. Logs (follow-up to #2816, which only covered success + the two direct _insertLog rejections): - Default requestedServiceTier + usedServiceTier in the insertLogEntry wrapper so guardrail/validation rejections, cache hits, streaming timeouts/upstream errors and fetch errors record them too. Declare servedServiceTier ahead of the wrapper and drop the now-redundant explicit fields from the two success inserts. - Reset servedServiceTier at the start of each retry attempt (both loops) so a fallback that fails before fetch returns logs no served tier instead of inheriting the previous provider's. Tests/infra: - Update the upstream-error test to openai/gpt-5.5 (priority, not subject to the upstream base-URL restriction from #2821) so it still reaches the error path; make the mock /v1/responses handler array-safe for TRIGGER_ERROR / status triggers and echo service_tier in its streamed response object. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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>
## 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 `priority` was requested. This is why production showed `requested=priority` / `used=null` on `aiplatform.googleapis.com` projects 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](https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/priority-paygo) and [liteLLM](https://docs.litellm.ai/docs/tutorials/vertex_ai_pay_go): *"To bypass Provisioned Throughput entirely, also set requestType: 'shared'."* ## Also in this PR: surface & log requested vs served tier - **Response**: `requested_service_tier` + `used_service_tier` added 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 so callers can detect it. - **Logs** (follow-up to theopenco#2816): 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 instead of inheriting the previous provider's. ## Tests / infra - `get-provider-headers` test asserts **both** Vertex headers on flex/priority and their absence otherwise. - Upstream-error test uses `openai/gpt-5.5` (not subject to theopenco#2821's upstream base-URL rule) so it still reaches the error path; mock `/v1/responses` handler is array-safe for `TRIGGER_ERROR` / `TRIGGER_STATUS_*` and echoes `service_tier` when streaming. - New tests: non-streaming metadata present on a priority request / absent without one; streaming final chunk carries both fields. ## Verification `pnpm test:unit` (132 files), full `apps/gateway/src/api.spec.ts` (101 passed / 2 skipped), `pnpm format`, full `pnpm build` all pass. The two-header behavior was confirmed against Google's Priority PayGo docs. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Summary
Flex/Priority service tiers are only honored when the request reaches Google directly. A
google-vertex/google-ai-studiocredential pointing at a custom base URL (a proxy) may silently drop the tier header (Vertex) or body field (AI Studio) — the upstream then serves standard, and the gateway bills/reports standard, even though the caller asked for Priority. This PR routes service-tier requests only to credentials that target the canonical upstream endpoint, covering both DB provider keys and env credentials.Behavior
iamFilteredModelProviders, mirroring the existing compliance-policy pattern). The gateway routes around the ineligible credential to a compliant provider; it does not 4xx when a compliant alternative exists.findProviderKey'sfilter(in hybrid mode it falls back to the env credential).LLM_..._BASE_URL). Indices whose base URL is a custom proxy are excluded from selection for service-tier requests, so the request lands on an upstream index.400instead of silently downgrading.Eligibility: a credential is eligible if it has no custom base URL (managed default) or its base URL exactly matches the provider's canonical upstream (
https://aiplatform.googleapis.comfor Vertex,https://generativelanguage.googleapis.comfor AI Studio). Providers without an upstream-only rule (e.g. OpenAI) are always eligible.Implementation
packages/actions:providerKeyBaseUrlSupportsServiceTier(provider, baseUrl)helper + unit tests.apps/gatewaychat.ts:enforceServiceTierKeyEligibility()filters the candidate provider list (initial selection +selectNextProviderfallback) and 400s only when pinned/none-remain;serviceTierKeyFilterpassed to inlinefindProviderKey; envgetProviderEnvcalls exclude non-compliant indices.get-provider-env.ts:getServiceTierIneligibleEnvIndices/hasServiceTierEligibleEnvCredential+ unit tests.resolve-provider-context.ts(retry path): same DB-key and env-index filtering.service-tiers.mdx.Verification (live, real Vertex/AI Studio)
[proxy, upstream], priorityconfigIndex=1->aiplatform.googleapis.com(proxy excluded)configIndex=0->vertex-proxy-test.invalid(proxy used; unaffected)requested_service_tier=priority,used_service_tier=priorityThe env-credential case was confirmed deterministically via the resolved endpoint URL: a
priorityrequest selected the upstream index while a non-tier request selected the proxy index. The persisted log row shows both requested and used tier =priority.pnpm format, fullpnpm build, and unit tests (21:apply-google-service-tier+get-provider-env) all pass.🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
Documentation
Tests