[model-gateway] move oai header util to router header util (#14441)

Co-authored-by: key4ng <rukeyang@gmail.com>
This commit is contained in:
Simo Lin
2025-12-04 11:02:50 -08:00
committed by GitHub
parent 7dfcc78155
commit d8faf2f361
4 changed files with 92 additions and 102 deletions

View File

@@ -1,4 +1,8 @@
use axum::{body::Body, extract::Request, http::HeaderMap};
use axum::{
body::Body,
extract::Request,
http::{HeaderMap, HeaderValue},
};
/// Copy request headers to a Vec of name-value string pairs
/// Used for forwarding headers to backend workers
@@ -92,3 +96,89 @@ pub fn apply_request_headers(
request_builder
}
/// API provider types for provider-specific header handling
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ApiProvider {
Anthropic,
Xai,
OpenAi,
Gemini,
Generic,
}
impl ApiProvider {
/// Detect provider type from URL
pub fn from_url(url: &str) -> Self {
if url.contains("anthropic") {
ApiProvider::Anthropic
} else if url.contains("x.ai") {
ApiProvider::Xai
} else if url.contains("openai.com") {
ApiProvider::OpenAi
} else if url.contains("googleapis.com") {
ApiProvider::Gemini
} else {
ApiProvider::Generic
}
}
}
/// Apply provider-specific headers to request
pub fn apply_provider_headers(
mut req: reqwest::RequestBuilder,
url: &str,
auth_header: Option<&HeaderValue>,
) -> reqwest::RequestBuilder {
let provider = ApiProvider::from_url(url);
match provider {
ApiProvider::Anthropic => {
// Anthropic requires x-api-key instead of Authorization
// Extract Bearer token and use as x-api-key
if let Some(auth) = auth_header {
if let Ok(auth_str) = auth.to_str() {
let api_key = auth_str.strip_prefix("Bearer ").unwrap_or(auth_str);
req = req
.header("x-api-key", api_key)
.header("anthropic-version", "2023-06-01");
}
}
}
ApiProvider::Gemini | ApiProvider::Xai | ApiProvider::OpenAi | ApiProvider::Generic => {
// Standard OpenAI-compatible: use Authorization header as-is
if let Some(auth) = auth_header {
req = req.header("Authorization", auth);
}
}
}
req
}
/// Extract auth header with passthrough semantics.
///
/// Passthrough mode: User's Authorization header takes priority.
/// Fallback: Worker's API key is used only if user didn't provide auth.
///
/// This enables use cases where:
/// 1. Users send their own API keys (multi-tenant, BYOK)
/// 2. Router has a default key for users who don't provide one
pub fn extract_auth_header(
headers: Option<&HeaderMap>,
worker_api_key: &Option<String>,
) -> Option<HeaderValue> {
// Passthrough: Try user's auth header first
let user_auth = headers.and_then(|h| {
h.get("authorization")
.or_else(|| h.get("Authorization"))
.cloned()
});
// Return user's auth if provided, otherwise use worker's API key
user_auth.or_else(|| {
worker_api_key
.as_ref()
.and_then(|k| HeaderValue::from_str(&format!("Bearer {}", k)).ok())
})
}

View File

@@ -22,10 +22,6 @@ use crate::{
protocols::responses::{generate_id, ResponseInput, ResponsesRequest},
};
// ============================================================================
// Persistence Operations (OpenAI-specific)
// ============================================================================
/// Persist conversation items to storage
///
/// This function:

View File

@@ -30,7 +30,6 @@ use super::{
provider::ProviderRegistry,
responses::{mask_tools_as_mcp, patch_streaming_response_json},
streaming::handle_streaming_response,
utils::{apply_provider_headers, extract_auth_header},
};
use crate::{
app_context::AppContext,
@@ -48,6 +47,7 @@ use crate::{
ResponsesGetParams, ResponsesRequest,
},
},
routers::header_utils::{apply_provider_headers, extract_auth_header},
};
pub struct OpenAIRouter {

View File

@@ -2,8 +2,6 @@
use std::collections::HashMap;
use axum::http::HeaderValue;
// ============================================================================
// SSE Event Type Constants
// ============================================================================
@@ -95,100 +93,6 @@ impl OutputIndexMapper {
}
}
// ============================================================================
// Provider Detection and Header Handling
// ============================================================================
/// API provider types
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ApiProvider {
Anthropic,
Xai,
OpenAi,
Gemini,
Generic,
}
impl ApiProvider {
/// Detect provider type from URL
pub fn from_url(url: &str) -> Self {
if url.contains("anthropic") {
ApiProvider::Anthropic
} else if url.contains("x.ai") {
ApiProvider::Xai
} else if url.contains("openai.com") {
ApiProvider::OpenAi
} else if url.contains("googleapis.com") {
ApiProvider::Gemini
} else {
ApiProvider::Generic
}
}
}
/// Apply provider-specific headers to request
pub fn apply_provider_headers(
mut req: reqwest::RequestBuilder,
url: &str,
auth_header: Option<&HeaderValue>,
) -> reqwest::RequestBuilder {
let provider = ApiProvider::from_url(url);
match provider {
ApiProvider::Anthropic => {
// Anthropic requires x-api-key instead of Authorization
// Extract Bearer token and use as x-api-key
if let Some(auth) = auth_header {
if let Ok(auth_str) = auth.to_str() {
let api_key = auth_str.strip_prefix("Bearer ").unwrap_or(auth_str);
req = req
.header("x-api-key", api_key)
.header("anthropic-version", "2023-06-01");
}
}
}
ApiProvider::Gemini | ApiProvider::Xai | ApiProvider::OpenAi | ApiProvider::Generic => {
// Standard OpenAI-compatible: use Authorization header as-is
if let Some(auth) = auth_header {
req = req.header("Authorization", auth);
}
}
}
req
}
// ============================================================================
// Auth Header Resolution
// ============================================================================
/// Extract auth header with passthrough semantics.
///
/// Passthrough mode: User's Authorization header takes priority.
/// Fallback: Worker's API key is used only if user didn't provide auth.
///
/// This enables use cases where:
/// 1. Users send their own API keys (multi-tenant, BYOK)
/// 2. Router has a default key for users who don't provide one
pub fn extract_auth_header(
headers: Option<&http::HeaderMap>,
worker_api_key: &Option<String>,
) -> Option<HeaderValue> {
// Passthrough: Try user's auth header first
let user_auth = headers.and_then(|h| {
h.get("authorization")
.or_else(|| h.get("Authorization"))
.cloned()
});
// Return user's auth if provided, otherwise use worker's API key
user_auth.or_else(|| {
worker_api_key
.as_ref()
.and_then(|k| HeaderValue::from_str(&format!("Bearer {}", k)).ok())
})
}
// ============================================================================
// Re-export FunctionCallInProgress from mcp module
// ============================================================================