Repository navigation
feat(discovery): per-provider admin API keys for external worker model discovery - #578
Conversation
…model discovery When SMG registers multiple external backends (OpenAI, Anthropic, xAI, Gemini), model discovery via /v1/models requires provider-specific API keys. Previously, --api-key was used for both discovery and inference across all workers, which doesn't work with multiple providers needing different credentials. Add per-provider "admin keys" resolved from environment variables, used only for model discovery (not inference). Resolution priority: 1. Per-provider env var (OPENAI_ADMIN_KEY, XAI_ADMIN_KEY, etc.) 2. --api-key flag (backward-compatible fallback) 3. No key → wildcard mode (unchanged) Also add provider-aware auth header handling: Anthropic uses x-api-key header instead of Authorization: Bearer. What changed: - protocols/src/worker.rs: Add from_url(), admin_key_env_var(), and uses_x_api_key() methods to ProviderType enum for URL-based provider detection and per-provider admin key env var mapping - model_gateway/src/core/steps/worker/external/discover_models.rs: Add resolve_discovery_api_key() function, update fetch_models() to accept provider for correct auth headers, update DiscoverModelsStep::execute() to use the new key resolution logic Fully backward compatible: --api-key without env vars works as before. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the model gateway's external worker model discovery by introducing a more robust and flexible API key management system. It addresses the challenge of integrating multiple external providers, each potentially requiring unique credentials, by allowing per-provider admin API keys to be configured via environment variables. This change also incorporates provider-specific authentication header handling, such as for Anthropic, ensuring broader compatibility while maintaining full backward compatibility with existing Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
📝 WalkthroughWalkthroughAdds provider-aware model discovery: provider is inferred from config URL, discovery API key is resolved with per-provider env var → config.api_key → none, and fetch_models uses provider to choose x-api-key vs. bearer authentication. Changes
Sequence Diagram(s)sequenceDiagram
participant WorkerStep as DiscoverModelsStep
participant ProviderUtil as ProviderType::from_url
participant Env as Environment
participant Fetch as fetch_models
participant Remote as Discovery API
Note over WorkerStep,ProviderUtil: DiscoverModelsStep.execute flow
WorkerStep->>ProviderUtil: parse config.url -> provider?
ProviderUtil-->>WorkerStep: provider or None
WorkerStep->>Env: check provider-specific env var (if provider)
Env-->>WorkerStep: env value or None
WorkerStep->>WorkerStep: resolve_discovery_api_key(env_var, config.api_key)
WorkerStep->>Fetch: fetch_models(url, discovery_key, provider)
Fetch->>Remote: HTTP request with header:
Note right of Fetch: if provider.uses_x_api_key -> "x-api-key: <key>"\nelse if key -> "Authorization: Bearer <key>"\nelse -> no auth (wildcard)
Remote-->>Fetch: models / error
Fetch-->>WorkerStep: models or error
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Comment |
github-actions
Bot
added
protocols
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a valuable feature for using per-provider admin API keys for model discovery. However, the current implementation of provider detection from URLs is flawed, using a simple string containment check (url.contains(...)) that can be bypassed to leak sensitive administrative API keys. Suggestions have been provided to enhance the robustness of provider detection and optimize the code by removing a redundant function call.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@protocols/src/worker.rs`:
- Around line 202-213: The provider detection in from_url is unsafe because it
uses url.contains; parse the input with url::Url::parse(url), get host_str(),
and match against canonical hostnames using exact-equals or suffix-with-dot
checks (e.g., host == "openai.com" || host.ends_with(".openai.com")) to prevent
matches like "openai.com.evil.example"; update the matching branches for OpenAI,
XAI, Anthropic, and Gemini in from_url to use the parsed host checks and return
None on parse failure.
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (2)
model_gateway/src/core/steps/worker/external/discover_models.rsprotocols/src/worker.rs
…tion Address code review feedback on PR #578: - Security: Parse URL host with url::Url instead of string contains() to prevent credential leakage via crafted URLs like http://attacker.com/openai.com/. Now checks host suffix only (ends_with on parsed host). Also tighten "anthropic" to "anthropic.com". - Simplify nested env var check using .ok().filter() idiom - Remove redundant ProviderType::from_url() call by passing the already-resolved provider into resolve_discovery_api_key() - Add url crate dependency to openai-protocol Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
There was a problem hiding this comment.
♻️ Duplicate comments (1)
protocols/src/worker.rs (1)
202-216:⚠️ Potential issue | 🔴 CriticalSecurity:
ends_with()hostname check still allows credential leakage to lookalike domains.The current
ends_with("openai.com")check matches malicious domains likeevilopenai.com, which would cause admin credentials to be sent to attacker-controlled servers. The same vulnerability exists for all provider checks.The fix requires checking for exact match OR dot-prefixed subdomain:
🔐 Proposed fix for secure hostname matching
pub fn from_url(url: &str) -> Option<Self> { let host = url::Url::parse(url).ok()?.host_str()?.to_lowercase(); - if host.ends_with("openai.com") { + if host == "openai.com" || host.ends_with(".openai.com") { Some(Self::OpenAI) - } else if host.ends_with("x.ai") { + } else if host == "x.ai" || host.ends_with(".x.ai") { Some(Self::XAI) - } else if host.ends_with("anthropic.com") { + } else if host == "anthropic.com" || host.ends_with(".anthropic.com") { Some(Self::Anthropic) - } else if host.ends_with("googleapis.com") { + } else if host == "googleapis.com" || host.ends_with(".googleapis.com") { Some(Self::Gemini) } else { None } }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@protocols/src/worker.rs` around lines 202 - 216, The hostname matching in from_url currently uses ends_with("openai.com") (and similar checks) which allows lookalike domains; change the checks in the from_url function to accept either an exact match OR a dot-prefixed subdomain (e.g., host == "openai.com" || host.ends_with(".openai.com")) for each provider (OpenAI, XAI, Anthropic, Gemini) so credentials are only sent to the real domains or their subdomains; keep the existing host normalization (to_lowercase) and apply the same pattern for "x.ai", "anthropic.com", and "googleapis.com".
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@protocols/src/worker.rs`:
- Around line 202-216: The hostname matching in from_url currently uses
ends_with("openai.com") (and similar checks) which allows lookalike domains;
change the checks in the from_url function to accept either an exact match OR a
dot-prefixed subdomain (e.g., host == "openai.com" ||
host.ends_with(".openai.com")) for each provider (OpenAI, XAI, Anthropic,
Gemini) so credentials are only sent to the real domains or their subdomains;
keep the existing host normalization (to_lowercase) and apply the same pattern
for "x.ai", "anthropic.com", and "googleapis.com".
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (3)
model_gateway/src/core/steps/worker/external/discover_models.rsprotocols/Cargo.tomlprotocols/src/worker.rs
…ixes The previous pin (fd080fc7) is the commit immediately before lightseekorg/tokenspeed#578, which adds defensive Finished-state handlers to the scheduler FSM. Without #578 the engine crashes under retract pressure with: RuntimeError: FSM transition invalid: event=tokenspeed::fsm::ExtendResultEvent; state=tokenspeed::fsm::Finished Reproduced on the nightly Qwen3-30B-A3B bench: when the host KV cache fills up and a retract fails, AbortEvent terminalizes the request → Finished, but overlap scheduling has already dispatched a forward batch including it, and the late ExtendResultEvent commit hits a strict FSM handler that throws and kills the scheduler event loop. Bump to current lightseekorg/tokenspeed main (eabeb106) so we also pick up #602 (release scheduler slot + cancel non-stream handlers on client disconnect), which removes the long pre-crash stream of ``Received output for rid=... but the state was deleted in AsyncLLM`` warnings caused by aborted requests still occupying engine slots. Signed-off-by: Chang Su <chang.s.su@oracle.com>
…ixes The previous pin (fd080fc7) is the commit immediately before lightseekorg/tokenspeed#578, which adds defensive Finished-state handlers to the scheduler FSM. Without #578 the engine crashes under retract pressure with: RuntimeError: FSM transition invalid: event=tokenspeed::fsm::ExtendResultEvent; state=tokenspeed::fsm::Finished Reproduced on the nightly Qwen3-30B-A3B bench: when the host KV cache fills up and a retract fails, AbortEvent terminalizes the request → Finished, but overlap scheduling has already dispatched a forward batch including it, and the late ExtendResultEvent commit hits a strict FSM handler that throws and kills the scheduler event loop. Bump to current lightseekorg/tokenspeed main (eabeb106) so we also pick up #602 (release scheduler slot + cancel non-stream handlers on client disconnect), which removes the long pre-crash stream of ``Received output for rid=... but the state was deleted in AsyncLLM`` warnings caused by aborted requests still occupying engine slots. Signed-off-by: key4ng <rukeyang@gmail.com>
1 participant
Summary
x-api-keyinstead of Bearer)--api-keywithout env vars works identically to beforeRefs: external worker registration support
What changed
protocols/src/worker.rs— 3 new methods onProviderType:from_url(url)→Option<ProviderType>— detect provider from URL domain (openai.com,x.ai,anthropic,googleapis.com)admin_key_env_var()→Option<&str>— maps provider to env var name (OPENAI_ADMIN_KEY,XAI_ADMIN_KEY,ANTHROPIC_ADMIN_KEY,GEMINI_ADMIN_KEY)uses_x_api_key()→bool— Anthropic requiresx-api-keyheadermodel_gateway/src/core/steps/worker/external/discover_models.rs:resolve_discovery_api_key()— resolves API key with priority: per-provider env var →--api-key→ None (wildcard mode)fetch_models()— acceptsOption<&ProviderType>, sendsx-api-keyfor Anthropic, Bearer for all othersDiscoverModelsStep::execute()— uses new resolution logic instead of checkingconfig.api_keydirectlyWhy
When registering multiple external backends (OpenAI, Anthropic, xAI, Gemini), model discovery requires calling
/v1/modelswith a valid API key. The existing--api-keyflag serves as a single key for all providers, which doesn't work when different providers need different credentials. Per-provider admin keys solve this by allowing operators to setOPENAI_ADMIN_KEY,XAI_ADMIN_KEY, etc. as environment variables, each used only for discovery on the matching provider.How
API key resolution follows a clear priority chain: first check for a per-provider env var based on URL domain detection (
ProviderType::from_url), then fall back to--api-key, then enter wildcard mode. TheProviderTypeenum in the protocols crate was extended (rather than creating a new enum) to keep provider logic centralized and reusable by bothcoreandrouterslayers.Usage examples
Test plan
cargo check -p smg— compilescargo clippy -p smg -- -D warnings— no warningscargo clippy -p openai-protocol -- -D warnings— no warningscargo test -p smg— all 413 tests passOPENAI_ADMIN_KEY, start with--worker-url https://api.openai.com --enable-igw, verify models discovered--api-keywithout env vars still works as beforeSummary by CodeRabbit