fix(xai): restore Codex compatibility and runtime semantics

Preserve transient credential failures as runtime cooldowns, restore account-aware retry and observability facts, and align Grok model/tool adaptation with Codex clients.

Keep complete SSE events on the live path while retaining only protocol-required tool and continuation state.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-25 23:40:15 +08:00
co-authored by Codex
parent 0879515ee8
commit 574f0d580e
48 changed files with 1756 additions and 394 deletions
+1
View File
@@ -1152,6 +1152,7 @@ dependencies = [
"rand_core 0.6.4",
"secrecy",
"serde",
"serde_json",
"subtle",
"thiserror",
"tokio",
+1
View File
@@ -27,6 +27,7 @@ thiserror = "2.0.18"
uuid = { version = "1.23.3", features = ["v7"] }
[dev-dependencies]
serde_json = { version = "1.0.150", features = ["preserve_order"] }
tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt"] }
[lints]
@@ -29,7 +29,6 @@ pub enum AccountStatus {
QuotaExhausted,
Disabled,
Banned,
Attention,
}
/// 账号列表排序字段。
@@ -184,7 +183,7 @@ pub struct AccountSummary {
pub total: u64,
pub active: u64,
pub quota_exhausted: u64,
pub attention: u64,
pub unavailable: u64,
}
/// 账号启停写入命令。
@@ -222,6 +221,8 @@ pub enum AccountConnectionTestEvent {
},
Failed {
message: String,
provider_error_code: Option<String>,
provider_error_type: Option<String>,
account_status: AccountStatus,
},
}
@@ -718,7 +718,8 @@ pub struct OpsError {
pub upstream_model_id: Option<String>,
pub upstream_transport: Option<String>,
pub failure_kind: String,
pub status_code: Option<u16>,
pub client_status_code: Option<u16>,
pub upstream_status_code: Option<u16>,
pub provider_error_code: Option<String>,
pub client_response_id: Option<String>,
pub upstream_request_id: Option<String>,
@@ -10,7 +10,6 @@ use gateway_core::{
credential::ProviderAccountId,
probe::{AccountProbe, AccountProbeRequest},
},
error::GatewayErrorKind,
routing::{ProviderKind, UpstreamModelId, snapshot::SnapshotControl},
};
@@ -454,12 +453,9 @@ impl AccountsService for DefaultAccountsService {
}))
.collect(),
Err(error) => vec![AccountConnectionTestEvent::Failed {
message: match error.kind() {
GatewayErrorKind::InvalidRequest
| GatewayErrorKind::Unsupported
| GatewayErrorKind::ModelNotFound => error.safe_message().to_owned(),
_ => "Provider connection test failed".to_owned(),
},
message: error.client_message().to_owned(),
provider_error_code: error.client_error_code().map(ToOwned::to_owned),
provider_error_type: error.client_error_type().map(ToOwned::to_owned),
account_status: status,
}],
}
@@ -485,8 +481,7 @@ fn account_status(account: &AccountRecord, now: chrono::DateTime<Utc>) -> Accoun
{
AccountStatus::Expired
}
crate::model::accounts::AccountAvailability::Ready => AccountStatus::Active,
_ => AccountStatus::Attention,
_ => AccountStatus::Active,
}
}
}
@@ -2,13 +2,14 @@ use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use chrono::{TimeDelta, Utc};
use futures::future::BoxFuture;
use futures::{StreamExt as _, future::BoxFuture};
use gateway_core::{
engine::{
credential::{OpaqueProviderData, ProviderAccountId},
probe::{AccountProbe, AccountProbeRequest, AccountProbeResult},
},
error::GatewayError,
error::{ClientVisibleUpstreamError, GatewayError, GatewayErrorKind},
operation::{GenerateRequest, Operation, ProtocolPayload},
routing::ProviderKind,
};
@@ -17,9 +18,9 @@ use gateway_admin::{
model::{
MutationContext, Revision,
accounts::{
AccountAvailability, AccountListQuery, AccountPage, AccountRecord, AccountSummary,
AccountUsage, AccountUsageWindowQuery, AccountUsageWindowResult, DeleteAccounts,
SetAccountEnabled,
AccountAvailability, AccountConnectionTestEvent, AccountListQuery, AccountPage,
AccountRecord, AccountSummary, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
},
observability::TimeRange,
provider_credentials::{
@@ -48,6 +49,7 @@ use gateway_admin::{
},
},
};
use serde_json::{Map, json};
pub(super) type EventLog = Arc<Mutex<Vec<&'static str>>>;
@@ -159,10 +161,22 @@ impl ProviderAdmin for FakeProviderAdmin {
fn connection_test_operation(
&self,
_: &gateway_core::routing::UpstreamModelId,
_: &str,
model: &gateway_core::routing::UpstreamModelId,
input: &str,
) -> Result<gateway_core::operation::Operation, ProviderAdminError> {
Err(ProviderAdminError::new(ProviderAdminErrorKind::Unsupported))
let payload = ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!(model.as_str())),
("input".to_owned(), json!(input)),
("stream".to_owned(), json!(true)),
("store".to_owned(), json!(false)),
]),
)
.map_err(|_| ProviderAdminError::new(ProviderAdminErrorKind::Invalid))?;
Ok(Operation::Generate(GenerateRequest::from_protocol_payload(
payload,
)))
}
fn dashboard_wire_profile(
@@ -395,7 +409,7 @@ impl AccountStore for FakeAccountStore {
total: 1,
active: 1,
quota_exhausted: 0,
attention: 0,
unavailable: 0,
},
})
}
@@ -686,6 +700,37 @@ fn provider_registry_should_reject_duplicate_kind() {
));
}
#[tokio::test]
async fn connection_test_should_preserve_safe_provider_error_details() {
let provider = FakeProviderAdmin::new("xai", events());
let store = FakeAccountStore::new("xai", events());
let services =
accounts_service_with_probe(provider, store, Arc::new(FailingAccountProbe)).await;
let events = services
.accounts()
.test_connection(
ProviderAccountId::new("acct_test").expect("account ID"),
gateway_core::routing::UpstreamModelId::new("grok-4.5").expect("model"),
)
.await
.expect("connection test stream")
.collect::<Vec<_>>()
.await;
assert!(matches!(
events.last(),
Some(AccountConnectionTestEvent::Failed {
message,
provider_error_code: Some(code),
provider_error_type: Some(error_type),
..
}) if message == "included usage exhausted"
&& code == "usage_exhausted"
&& error_type == "invalid_request_error"
));
}
#[tokio::test]
async fn accounts_export_should_pass_store_loaded_timestamps_and_material_to_provider() {
let provider = FakeProviderAdmin::new("openai", events());
@@ -1005,12 +1050,20 @@ fn rotation_result(command: CredentialRotationCommit) -> CredentialMutationResul
async fn accounts_service(
provider: Arc<FakeProviderAdmin>,
store: Arc<FakeAccountStore>,
) -> AdminServices {
accounts_service_with_probe(provider, store, Arc::new(SuccessfulAccountProbe)).await
}
async fn accounts_service_with_probe(
provider: Arc<FakeProviderAdmin>,
store: Arc<FakeAccountStore>,
probe: Arc<dyn AccountProbe>,
) -> AdminServices {
super::AdminHarness::new()
.accounts(store)
.settings(Arc::new(StaticSettingsStore))
.provider(provider)
.probe(Arc::new(SuccessfulAccountProbe))
.probe(probe)
.build()
.await
}
@@ -1030,6 +1083,29 @@ impl AccountProbe for SuccessfulAccountProbe {
}
}
struct FailingAccountProbe;
impl AccountProbe for FailingAccountProbe {
fn probe(
&self,
_: AccountProbeRequest,
) -> BoxFuture<'_, Result<AccountProbeResult, GatewayError>> {
Box::pin(async {
let detail = ClientVisibleUpstreamError::new(
"included usage exhausted",
Some("usage_exhausted".to_owned()),
Some("invalid_request_error".to_owned()),
)
.expect("safe provider error");
Err(GatewayError::new(
GatewayErrorKind::RateLimited,
"upstream capacity is temporarily unavailable",
)
.with_client_visible_upstream_error(detail))
})
}
}
fn store_unavailable() -> AdminStoreError {
AdminStoreError::new(
AdminStoreErrorKind::Unavailable,
@@ -174,7 +174,7 @@ pub struct AccountSummaryView {
pub total: u64,
pub active: u64,
pub quota_exhausted: u64,
pub attention: u64,
pub unavailable: u64,
}
/// 一条安全账号视图。
@@ -543,10 +543,14 @@ impl From<DomainConnectionTestEvent> for AccountConnectionTestEvent {
}),
DomainConnectionTestEvent::Failed {
message,
provider_error_code,
provider_error_type,
account_status,
} => serde_json::json!({
"type": "error",
"error": message,
"providerErrorCode": provider_error_code,
"providerErrorType": provider_error_type,
"accountStatus": account_status_name(account_status)
}),
};
@@ -1401,7 +1405,7 @@ fn account_page_data(
total: result.summary.total,
active: result.summary.active,
quota_exhausted: result.summary.quota_exhausted,
attention: result.summary.attention,
unavailable: result.summary.unavailable,
},
}
}
@@ -1798,7 +1802,6 @@ fn account_status_name(status: DomainAccountStatus) -> &'static str {
DomainAccountStatus::QuotaExhausted => "quota_exhausted",
DomainAccountStatus::Disabled => "disabled",
DomainAccountStatus::Banned => "banned",
DomainAccountStatus::Attention => "attention",
}
}
@@ -801,8 +801,8 @@ pub struct DashboardCardsView {
pub struct DashboardCredentialsCardView {
pub total: String,
pub total_value: u64,
pub enabled: String,
pub enabled_value: u64,
pub available: String,
pub available_value: u64,
pub unavailable: String,
pub unavailable_value: u64,
}
@@ -1235,7 +1235,6 @@ pub struct OpsErrorView {
pub account_id: Option<String>,
pub route: String,
pub model: Option<String>,
pub status_code: Option<i64>,
pub client_status_code: Option<i64>,
pub upstream_status_code: Option<i64>,
pub transport: Option<String>,
@@ -1954,7 +1953,6 @@ fn usage_page_view(
}
fn ops_error_view(error: domain::OpsError) -> OpsErrorView {
let status = error.status_code.map(i64::from);
OpsErrorView {
id: error.event_id,
request_id: error.request_id,
@@ -1964,9 +1962,8 @@ fn ops_error_view(error: domain::OpsError) -> OpsErrorView {
account_id: error.provider_account_ref,
route: error.operation,
model: error.upstream_model_id,
status_code: status,
client_status_code: None,
upstream_status_code: None,
client_status_code: error.client_status_code.map(i64::from),
upstream_status_code: error.upstream_status_code.map(i64::from),
transport: error.upstream_transport,
attempt_index: error.attempt_index,
failure_class: error.failure_kind,
@@ -2290,8 +2287,8 @@ fn dashboard_view(result: domain::DashboardResult, kind: TrendKind) -> Dashboard
credentials: DashboardCredentialsCardView {
total: format_compact_number(provider_accounts.total),
total_value: provider_accounts.total,
enabled: format_compact_number(provider_accounts.enabled),
enabled_value: provider_accounts.enabled,
available: format_compact_number(provider_accounts.active),
available_value: provider_accounts.active,
unavailable: format_compact_number(provider_accounts.unavailable),
unavailable_value: provider_accounts.unavailable,
},
@@ -6,7 +6,7 @@ use axum::{
http::{HeaderMap, StatusCode},
response::{IntoResponse, Response},
};
use gateway_core::routing::PublicModelId;
use gateway_core::routing::{PublicModelId, PublicModelProfile};
use serde_json::{Value, json};
use crate::ApiState;
@@ -17,6 +17,7 @@ use super::{
};
const MODEL_CREATED_TIMESTAMP: i64 = 1_700_000_000;
const CODEX_BASE_INSTRUCTIONS: &str = "You are Codex, a coding agent. Follow the user's instructions and use the available tools to complete software engineering tasks. Inspect relevant files before editing, preserve unrelated changes, and verify the result.";
/// `GET /v1/models`。
pub(crate) async fn models(State(state): State<ApiState>, headers: HeaderMap) -> Response {
@@ -30,6 +31,16 @@ pub(crate) async fn models(State(state): State<ApiState>, headers: HeaderMap) ->
.into_iter()
.map(|model| openai_model_json(&model))
.collect::<Vec<_>>();
let profiles = service.public_model_profiles(&client);
if !profiles.is_empty() {
let models = profiles
.iter()
.enumerate()
.map(|(index, profile)| codex_model_json(profile, index))
.collect::<Vec<_>>();
return (StatusCode::OK, Json(json!({ "models": models }))).into_response();
}
(
StatusCode::OK,
@@ -134,3 +145,75 @@ fn catalog_model_json(id: &str) -> Value {
"source": "gateway",
})
}
fn codex_model_json(profile: &PublicModelProfile, index: usize) -> Value {
let id = profile.model().as_str();
let presentation = profile.presentation();
let reasoning_levels = presentation
.supported_reasoning_efforts()
.iter()
.map(|effort| {
json!({
"effort": effort,
"description": reasoning_effort_description(effort),
})
})
.collect::<Vec<_>>();
let reasoning_supported = presentation
.supported_reasoning_efforts()
.iter()
.any(|effort| effort != "none");
let context_window = presentation.context_window_tokens();
let input_modalities = if presentation.image_input() {
vec!["text", "image"]
} else {
vec!["text"]
};
json!({
"slug": id,
"display_name": presentation.display_name().unwrap_or(id),
"description": presentation.description(),
"default_reasoning_level": presentation.default_reasoning_effort(),
"supported_reasoning_levels": reasoning_levels,
"shell_type": "shell_command",
"visibility": if presentation.hidden() { "hide" } else { "list" },
"supported_in_api": true,
"priority": index.saturating_add(1),
"additional_speed_tiers": [],
"service_tiers": [],
"availability_nux": Value::Null,
"upgrade": Value::Null,
"base_instructions": CODEX_BASE_INSTRUCTIONS,
"model_messages": Value::Null,
"include_skills_usage_instructions": false,
"supports_reasoning_summary_parameter": reasoning_supported,
"default_reasoning_summary": "auto",
"support_verbosity": presentation.verbosity(),
"default_verbosity": Value::Null,
"apply_patch_tool_type": presentation.agent_tools().then_some("freeform"),
"web_search_tool_type": "text",
"truncation_policy": { "mode": "tokens", "limit": 10_000 },
"supports_parallel_tool_calls": presentation.parallel_tool_calls(),
"supports_image_detail_original": presentation.image_detail_original(),
"context_window": context_window,
"max_context_window": context_window,
"effective_context_window_percent": 95,
"experimental_supported_tools": [],
"input_modalities": input_modalities,
"supports_search_tool": presentation.search_tool(),
"use_responses_lite": false,
})
}
fn reasoning_effort_description(effort: &str) -> &'static str {
match effort {
"none" => "No reasoning",
"low" => "Fast responses with lighter reasoning",
"medium" => "Balances speed and reasoning depth for everyday tasks",
"high" => "Greater reasoning depth for complex problems",
"xhigh" => "Extra high reasoning depth for complex problems",
"max" => "Maximum reasoning depth for the hardest problems",
_ => "Provider-supported reasoning level",
}
}
@@ -9,7 +9,7 @@ use gateway_core::engine::execution::{
};
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::lifecycle::{ConnectionDraining, ConnectionGuard, ConnectionLifecycle};
use gateway_core::routing::PublicModelId;
use gateway_core::routing::{PublicModelId, PublicModelProfile};
use uuid::Uuid;
use super::auth::ClientApiKeyAuthError;
@@ -51,6 +51,13 @@ impl OpenAiService {
.collect()
}
pub(crate) fn public_model_profiles(
&self,
client: &AuthenticatedClient,
) -> Vec<PublicModelProfile> {
self.execution.public_model_profiles(client)
}
pub(crate) fn contains_public_model(
&self,
client: &AuthenticatedClient,
@@ -276,6 +276,8 @@ mod actions {
},
DomainConnectionTestEvent::Failed {
message: "upstream unavailable".to_owned(),
provider_error_code: Some("usage_exhausted".to_owned()),
provider_error_type: Some("invalid_request_error".to_owned()),
account_status: AccountStatus::Active,
},
]
@@ -302,7 +304,13 @@ mod actions {
}),
json!({ "type": "content", "text": "OK" }),
json!({ "type": "test_complete", "success": true, "accountStatus": "active" }),
json!({ "type": "error", "error": "upstream unavailable", "accountStatus": "active" }),
json!({
"type": "error",
"error": "upstream unavailable",
"providerErrorCode": "usage_exhausted",
"providerErrorType": "invalid_request_error",
"accountStatus": "active"
}),
]
);
}
@@ -10,19 +10,28 @@ use gateway_core::engine::execution::{
StartedExecution,
};
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::routing::PublicModelId;
use gateway_core::routing::{ModelPresentation, PublicModelId, PublicModelProfile};
use tower::ServiceExt;
use super::{api_router, authenticated_client};
struct ModelsExecution {
client: AuthenticatedClient,
profiles: bool,
}
impl ModelsExecution {
fn new() -> Arc<Self> {
Arc::new(Self {
client: authenticated_client("sk_models_test"),
profiles: false,
})
}
fn with_profiles() -> Arc<Self> {
Arc::new(Self {
client: authenticated_client("sk_models_test"),
profiles: true,
})
}
}
@@ -46,6 +55,26 @@ impl ExecutionService for ModelsExecution {
.collect()
}
fn public_model_profiles(&self, _: &AuthenticatedClient) -> Vec<PublicModelProfile> {
if !self.profiles {
return Vec::new();
}
vec![PublicModelProfile::new(
PublicModelId::new("grok-4.5").expect("model"),
ModelPresentation::new(
Some("Grok 4.5".to_owned()),
Some("xAI Grok 4.5 frontier model.".to_owned()),
)
.with_reasoning(
Some("medium".to_owned()),
["low", "medium", "high"].map(str::to_owned).to_vec(),
)
.with_context_window_tokens(Some(500_000))
.with_image_input(true)
.with_agent_tools(true, true),
)]
}
fn contains_public_model(&self, _: &AuthenticatedClient, model: &PublicModelId) -> bool {
matches!(model.as_str(), "model-a" | "model-b")
}
@@ -91,6 +120,26 @@ async fn models_should_encode_the_service_visible_catalog() {
);
}
#[tokio::test]
async fn models_should_encode_provider_profiles_for_current_codex_clients() {
let response = api_router(ModelsExecution::with_profiles())
.await
.oneshot(authorized_request("/v1/models?client_version=0.145.0"))
.await
.expect("list Codex models response");
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("read Codex models body");
let value: serde_json::Value = serde_json::from_slice(&body).expect("models JSON");
let model = &value["models"][0];
assert_eq!(model["slug"], "grok-4.5");
assert_eq!(model["default_reasoning_level"], "medium");
assert_eq!(model["context_window"], 500_000);
assert_eq!(model["apply_patch_tool_type"], "freeform");
assert_eq!(model["service_tiers"], serde_json::json!([]));
}
#[tokio::test]
async fn model_detail_should_keep_the_official_path_id_contract() {
let response = api_router(ModelsExecution::new())
@@ -1,6 +1,6 @@
//! 唯一的账号重试、发送与下游 commit barrier owner。
use std::collections::{BTreeSet, VecDeque};
use std::collections::BTreeSet;
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime};
@@ -67,12 +67,6 @@ where
let continuation_attempt =
initial_continuation_attempt(&operation, &plan, continuation.as_ref());
let image_generation_requested = operation.image_generation_requested();
let commit_policy = if continuation_attempt == ContinuationAttempt::None {
StreamCommitPolicy::FirstCanonicalEvent
} else {
StreamCommitPolicy::UntilOutputOrTerminal
};
let mut session = ResponseExecutionSession {
engine: Arc::clone(&self.engine),
request_id,
@@ -86,7 +80,6 @@ where
required_account,
continuation,
continuation_attempt,
commit_policy,
account_state_owner,
cancellation,
attempts: 0,
@@ -106,7 +99,6 @@ where
upstream_response_id: None,
last_account_exhaustion: None,
provider_attempt_outcomes: Vec::new(),
before_commit: VecDeque::new(),
};
if session.cancellation.is_cancelled() {
@@ -141,12 +133,6 @@ struct FailureFinalization {
provider_response_id: Option<String>,
}
#[derive(Clone, Copy)]
enum StreamCommitPolicy {
FirstCanonicalEvent,
UntilOutputOrTerminal,
}
/// API 可逐事件消费的 Core 执行会话。
///
/// API 只能提交下游 delivery 边界;账号重试、断流终结与
@@ -164,7 +150,6 @@ pub struct ResponseExecutionSession<S: ?Sized> {
required_account: Option<crate::engine::credential::ProviderAccountId>,
continuation: Option<ContinuationBinding>,
continuation_attempt: ContinuationAttempt,
commit_policy: StreamCommitPolicy,
account_state_owner: Option<ProviderAccountStateOwner>,
cancellation: CancellationToken,
attempts: u32,
@@ -184,7 +169,6 @@ pub struct ResponseExecutionSession<S: ?Sized> {
upstream_response_id: Option<String>,
last_account_exhaustion: Option<ProviderError>,
provider_attempt_outcomes: Vec<ProviderAttemptOutcome>,
before_commit: VecDeque<ProviderEvent>,
}
impl<S: ?Sized> ResponseExecutionSession<S>
@@ -203,15 +187,6 @@ where
if self.finalized {
return Ok(None);
}
if self.downstream_committed_at.is_some()
&& let Some(event) = self.before_commit.pop_front()
{
return Ok(Some(CoordinatedEvent::single(
event,
CommitRequirement::AlreadyCommitted,
)));
}
loop {
match self.pull().await? {
PullOutcome::Event(event) => {
@@ -221,25 +196,13 @@ where
CommitRequirement::AlreadyCommitted,
)));
}
let commit_significant = match self.commit_policy {
StreamCommitPolicy::FirstCanonicalEvent => event.has_canonical_facts(),
StreamCommitPolicy::UntilOutputOrTerminal => event.is_commit_significant(),
};
self.before_commit.push_back(event);
if !commit_significant {
continue;
}
self.delivery_pending = true;
let events = self.before_commit.drain(..).collect();
return CoordinatedEvent::try_batch(
events,
return Ok(Some(CoordinatedEvent::single(
event,
CommitRequirement::CommitBeforeDelivery,
)
.map(Some);
}
PullOutcome::AttemptDiscarded => {
self.before_commit.clear();
)));
}
PullOutcome::AttemptDiscarded => {}
PullOutcome::End => {
if self.downstream_committed_at.is_some() {
self.finish_success().await?;
@@ -953,7 +916,6 @@ where
self.timings.first_text_ms = None;
self.timings.first_token_ms = None;
self.upstream_complete = false;
self.before_commit.clear();
}
async fn finish_success(&mut self) -> Result<(), EngineError> {
@@ -29,7 +29,9 @@ use crate::event::{GatewayEvent, ProviderEvent, ProviderResponseHeader};
use crate::operation::{Operation, ProviderSessionState};
use crate::policy::{ClientApiKeyId, ClientPolicy};
use crate::routing::snapshot::RuntimeSnapshotHandle;
use crate::routing::{ProviderKind, PublicModelId, RoutingContext, RuntimeSnapshot};
use crate::routing::{
ProviderKind, PublicModelId, PublicModelProfile, RoutingContext, RuntimeSnapshot,
};
const MODEL_REQUEST_DEADLINE: Duration = Duration::from_secs(10 * 60);
const CONTINUATION_AFFINITY_TIMEOUT: Duration = Duration::from_millis(100);
@@ -139,6 +141,9 @@ pub trait ExecutionService: Send + Sync {
plaintext: &str,
) -> Result<AuthenticatedClient, ClientAuthenticationError>;
fn public_models(&self, client: &AuthenticatedClient) -> Vec<PublicModelId>;
fn public_model_profiles(&self, _client: &AuthenticatedClient) -> Vec<PublicModelProfile> {
Vec::new()
}
fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool;
fn start(
&self,
@@ -568,6 +573,12 @@ impl ExecutionService for DefaultExecutionService {
.public_models_for_provider(client.policy.provider_kind())
}
fn public_model_profiles(&self, client: &AuthenticatedClient) -> Vec<PublicModelProfile> {
client
.snapshot
.public_model_profiles_for_provider(client.policy.provider_kind())
}
fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool {
client
.snapshot
@@ -20,7 +20,9 @@ use crate::error::{
};
use crate::event::{EventSequenceValidator, ProviderEvent};
use crate::operation::Operation;
use crate::routing::{ModelCapabilities, ProviderCandidate, ProviderKind, UpstreamModelId};
use crate::routing::{
ModelCapabilities, ModelPresentation, ProviderCandidate, ProviderKind, UpstreamModelId,
};
/// Box 只出现在 Provider Registry 的统一 event envelope 边界。
pub type EventStream =
@@ -433,6 +435,7 @@ pub struct ProviderRequestObservation {
pub struct ProviderModelCapabilities {
upstream_model: UpstreamModelId,
capabilities: ModelCapabilities,
presentation: Option<ModelPresentation>,
}
/// Provider 实时目录成功发布后的进程内单调代次。
@@ -459,9 +462,16 @@ impl ProviderModelCapabilities {
Self {
upstream_model,
capabilities,
presentation: None,
}
}
#[must_use]
pub fn with_presentation(mut self, presentation: ModelPresentation) -> Self {
self.presentation = Some(presentation);
self
}
#[must_use]
pub const fn upstream_model(&self) -> &UpstreamModelId {
&self.upstream_model
@@ -471,6 +481,11 @@ impl ProviderModelCapabilities {
pub const fn capabilities(&self) -> &ModelCapabilities {
&self.capabilities
}
#[must_use]
pub const fn presentation(&self) -> Option<&ModelPresentation> {
self.presentation.as_ref()
}
}
impl ProviderRequest {
-29
View File
@@ -756,18 +756,6 @@ impl ProviderEvent {
pub fn has_client_event(&self) -> bool {
!self.canonical.is_empty() || self.wire.is_some()
}
/// 返回是否包含足以冻结下游 commit barrier 的 canonical fact。
///
/// 单独的 wire event 不能冻结边界:协议层在看到 `Started`
/// 携带的网关响应身份前无法安全编码它。Core 会保留这些
/// 事件,直到首个 canonical 可见增量或终态再一并交付。
#[must_use]
pub fn is_commit_significant(&self) -> bool {
self.canonical
.iter()
.any(GatewayEvent::is_commit_significant)
}
}
impl From<GatewayEvent> for ProviderEvent {
@@ -776,23 +764,6 @@ impl From<GatewayEvent> for ProviderEvent {
}
}
impl GatewayEvent {
/// 返回该 fact 是否足以冻结下游 commit barrier。
///
/// 生命周期、结构与结算 facts 可以在换号前安全丢弃;首个可见增量或
/// Provider 正常终态才会使客户端输出不可撤回。
#[must_use]
pub const fn is_commit_significant(&self) -> bool {
matches!(
self,
Self::TextDelta(_)
| Self::ReasoningDelta(_)
| Self::ToolCallDelta(_)
| Self::Completed(_)
)
}
}
/// Canonical event 顺序错误。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum EventSequenceError {
@@ -476,7 +476,7 @@ pub trait ProviderCredentialStatePort: Send + Sync {
) -> BoxFuture<'a, Result<bool, ProviderStoreError>>;
}
/// 临时 cooldown 只保存调度截止时间;原因事实仍由 PostgreSQL 账号状态持有。
/// 临时 cooldown 只保存可丢失的调度截止时间,不进入账号持久状态。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderCooldown {
account_id: ProviderAccountId,
@@ -136,6 +136,177 @@ pub struct ModelCapabilities {
upstream_validates_features: bool,
}
/// Provider 为客户端模型目录提供的展示与交互能力。
///
/// 该值不参与路由,也不把任一 Provider 的 wire 类型带入 Core。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ModelPresentation {
display_name: Option<String>,
description: Option<String>,
default_reasoning_effort: Option<String>,
supported_reasoning_efforts: Vec<String>,
context_window_tokens: Option<u64>,
image_input: bool,
agent_tools: bool,
parallel_tool_calls: bool,
search_tool: bool,
image_detail_original: bool,
verbosity: bool,
hidden: bool,
}
impl ModelPresentation {
#[must_use]
pub fn new(display_name: Option<String>, description: Option<String>) -> Self {
Self {
display_name,
description,
..Self::default()
}
}
#[must_use]
pub fn with_reasoning(
mut self,
default_effort: Option<String>,
supported_efforts: Vec<String>,
) -> Self {
self.default_reasoning_effort = default_effort;
self.supported_reasoning_efforts = supported_efforts;
self
}
#[must_use]
pub const fn with_context_window_tokens(mut self, context_window_tokens: Option<u64>) -> Self {
self.context_window_tokens = context_window_tokens;
self
}
#[must_use]
pub const fn with_image_input(mut self, image_input: bool) -> Self {
self.image_input = image_input;
self
}
#[must_use]
pub const fn with_agent_tools(mut self, agent_tools: bool, parallel_tool_calls: bool) -> Self {
self.agent_tools = agent_tools;
self.parallel_tool_calls = parallel_tool_calls;
self
}
#[must_use]
pub const fn with_search_tool(mut self, search_tool: bool) -> Self {
self.search_tool = search_tool;
self
}
#[must_use]
pub const fn with_image_detail_original(mut self, image_detail_original: bool) -> Self {
self.image_detail_original = image_detail_original;
self
}
#[must_use]
pub const fn with_verbosity(mut self, verbosity: bool) -> Self {
self.verbosity = verbosity;
self
}
#[must_use]
pub const fn with_hidden(mut self, hidden: bool) -> Self {
self.hidden = hidden;
self
}
#[must_use]
pub fn display_name(&self) -> Option<&str> {
self.display_name.as_deref()
}
#[must_use]
pub fn description(&self) -> Option<&str> {
self.description.as_deref()
}
#[must_use]
pub fn default_reasoning_effort(&self) -> Option<&str> {
self.default_reasoning_effort.as_deref()
}
#[must_use]
pub fn supported_reasoning_efforts(&self) -> &[String] {
&self.supported_reasoning_efforts
}
#[must_use]
pub const fn context_window_tokens(&self) -> Option<u64> {
self.context_window_tokens
}
#[must_use]
pub const fn image_input(&self) -> bool {
self.image_input
}
#[must_use]
pub const fn agent_tools(&self) -> bool {
self.agent_tools
}
#[must_use]
pub const fn parallel_tool_calls(&self) -> bool {
self.parallel_tool_calls
}
#[must_use]
pub const fn search_tool(&self) -> bool {
self.search_tool
}
#[must_use]
pub const fn image_detail_original(&self) -> bool {
self.image_detail_original
}
#[must_use]
pub const fn verbosity(&self) -> bool {
self.verbosity
}
#[must_use]
pub const fn hidden(&self) -> bool {
self.hidden
}
}
/// 一个公开模型及其 Provider 编译后的客户端画像。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PublicModelProfile {
model: PublicModelId,
presentation: ModelPresentation,
}
impl PublicModelProfile {
#[must_use]
pub const fn new(model: PublicModelId, presentation: ModelPresentation) -> Self {
Self {
model,
presentation,
}
}
#[must_use]
pub const fn model(&self) -> &PublicModelId {
&self.model
}
#[must_use]
pub const fn presentation(&self) -> &ModelPresentation {
&self.presentation
}
}
impl ModelCapabilities {
#[must_use]
pub fn new(operations: BTreeSet<OperationKind>, max_output_tokens: Option<u64>) -> Self {
@@ -205,6 +376,7 @@ pub struct ProviderModel {
provider: ProviderKind,
upstream_model: UpstreamModelId,
capabilities: ModelCapabilities,
presentation: Option<ModelPresentation>,
}
impl ProviderModel {
@@ -218,9 +390,16 @@ impl ProviderModel {
provider,
upstream_model,
capabilities,
presentation: None,
}
}
#[must_use]
pub fn with_presentation(mut self, presentation: ModelPresentation) -> Self {
self.presentation = Some(presentation);
self
}
#[must_use]
pub const fn provider(&self) -> &ProviderKind {
&self.provider
@@ -230,6 +409,11 @@ impl ProviderModel {
pub const fn upstream_model(&self) -> &UpstreamModelId {
&self.upstream_model
}
#[must_use]
pub const fn presentation(&self) -> Option<&ModelPresentation> {
self.presentation.as_ref()
}
}
/// 本次请求选择 Provider 时使用的动态过滤事实。
@@ -231,11 +231,15 @@ async fn compile_runtime_snapshot(
continue;
};
provider_models.extend(models.into_iter().map(|model| {
ProviderModel::new(
let compiled = ProviderModel::new(
provider.clone(),
model.upstream_model().clone(),
model.capabilities().clone(),
)
);
match model.presentation().cloned() {
Some(presentation) => compiled.with_presentation(presentation),
None => compiled,
}
}));
}
@@ -287,6 +291,8 @@ pub struct RuntimeSnapshot {
account_selection_policy: AccountSelectionPolicy,
providers: Arc<BTreeSet<ProviderKind>>,
provider_models: Arc<BTreeMap<ProviderKind, BTreeMap<UpstreamModelId, ModelCapabilities>>>,
provider_model_presentations:
Arc<BTreeMap<ProviderKind, BTreeMap<UpstreamModelId, super::ModelPresentation>>>,
model_mappings: Arc<BTreeMap<String, String>>,
provider_catalog_generations: Arc<BTreeMap<ProviderKind, ProviderCatalogGeneration>>,
client_policies: Arc<BTreeMap<ClientApiKeyId, ClientPolicy>>,
@@ -313,17 +319,24 @@ impl RuntimeSnapshot {
let mut model_map =
BTreeMap::<ProviderKind, BTreeMap<UpstreamModelId, ModelCapabilities>>::new();
let mut presentation_map =
BTreeMap::<ProviderKind, BTreeMap<UpstreamModelId, super::ModelPresentation>>::new();
for model in provider_models {
if !provider_set.contains(model.provider()) {
let ProviderModel {
provider,
upstream_model,
capabilities,
presentation,
} = model;
if !provider_set.contains(&provider) {
return Err(RoutingError::NotFound {
entity: "provider",
id: model.provider().to_string(),
id: provider.to_string(),
});
}
let models = model_map.entry(model.provider).or_default();
let upstream_model = model.upstream_model;
let models = model_map.entry(provider.clone()).or_default();
if models
.insert(upstream_model.clone(), model.capabilities)
.insert(upstream_model.clone(), capabilities)
.is_some()
{
return Err(RoutingError::DuplicateEntity {
@@ -331,6 +344,12 @@ impl RuntimeSnapshot {
id: upstream_model.to_string(),
});
}
if let Some(presentation) = presentation {
presentation_map
.entry(provider)
.or_default()
.insert(upstream_model, presentation);
}
}
let mut client_policy_map = BTreeMap::new();
@@ -350,6 +369,7 @@ impl RuntimeSnapshot {
account_selection_policy,
providers: Arc::new(provider_set),
provider_models: Arc::new(model_map),
provider_model_presentations: Arc::new(presentation_map),
model_mappings: Arc::new(BTreeMap::new()),
provider_catalog_generations: Arc::new(BTreeMap::new()),
client_policies: Arc::new(client_policy_map),
@@ -402,6 +422,38 @@ impl RuntimeSnapshot {
models.into_iter().collect()
}
/// 返回 Provider 已明确声明画像的公开模型;没有画像时不猜测 Provider 语义。
#[must_use]
pub fn public_model_profiles_for_provider(
&self,
provider: &ProviderKind,
) -> Vec<super::PublicModelProfile> {
let Some(presentations) = self.provider_model_presentations.get(provider) else {
return Vec::new();
};
let mut profiles = BTreeMap::new();
for (model, presentation) in presentations {
if let Ok(public_model) = PublicModelId::new(model.as_str().to_owned()) {
profiles.insert(public_model, presentation.clone());
}
}
for alias in self.model_mappings.keys() {
let target = self.mapped_model(alias);
let Some(presentation) = presentations.iter().find_map(|(model, presentation)| {
(model.as_str() == target).then_some(presentation)
}) else {
continue;
};
if let Ok(public_model) = PublicModelId::new(alias.clone()) {
profiles.insert(public_model, presentation.clone());
}
}
profiles
.into_iter()
.map(|(model, presentation)| super::PublicModelProfile::new(model, presentation))
.collect()
}
#[must_use]
pub fn public_models(&self) -> Vec<PublicModelId> {
self.providers
@@ -1005,6 +1005,51 @@ fn authenticated_native_continuation_reaches_every_attempt_context() {
));
}
#[test]
fn streaming_native_continuation_delivers_started_event_before_later_output() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, _, _) = coordinator(vec![Script::Stream {
account_id: "acct_one",
items: complete_stream(None),
}]);
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-response").expect("previous response"),
SafeUpstreamValue::new("upstream-response").expect("upstream response"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_one").expect("account"),
);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
Some(ContinuationBinding::Pinned(continuation)),
CancellationToken::new(),
))
.expect("start execution");
let first = block_on(session.next_event())
.expect("first event")
.expect("started event");
assert_eq!(
first.commit_requirement(),
CommitRequirement::CommitBeforeDelivery
);
let events = first.into_provider_events();
assert!(matches!(
events.as_slice(),
[event]
if matches!(event.canonical_facts(), [GatewayEvent::Started(_)])
));
block_on(session.commit_downstream(Some(200))).expect("commit started event");
while block_on(session.next_event())
.expect("remaining event")
.is_some()
{}
}
#[test]
fn native_continuation_replays_owner_before_safely_switching_account() {
let Operation::Generate(generate) = generate_operation() else {
@@ -24,8 +24,8 @@ use gateway_core::routing::snapshot::{
SnapshotSubscriptionError, SnapshotSubscriptionPort, runtime_revision_needs_refresh,
};
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, PublicModelId, RuntimeSnapshot,
UpstreamModelId,
ConfigRevision, ModelCapabilities, ModelPresentation, ProviderKind, PublicModelId,
RuntimeSnapshot, UpstreamModelId,
};
use gateway_core::task::WorkerKind;
@@ -135,13 +135,19 @@ impl Provider for PublishingCatalogProvider {
if self.queries.fetch_add(1, Ordering::SeqCst) == 0 {
self.generation.store(1, Ordering::SeqCst);
}
Ok(vec![ProviderModelCapabilities::new(
UpstreamModelId::new("upstream-model").expect("model"),
ModelCapabilities::new(
std::collections::BTreeSet::from([OperationKind::Generate]),
Ok(vec![
ProviderModelCapabilities::new(
UpstreamModelId::new("upstream-model").expect("model"),
ModelCapabilities::new(
std::collections::BTreeSet::from([OperationKind::Generate]),
None,
),
)
.with_presentation(ModelPresentation::new(
Some("Upstream Model".to_owned()),
None,
),
)])
)),
])
}
async fn execute(
@@ -215,6 +221,15 @@ fn compiler_retries_when_provider_publishes_catalog_during_compilation() {
.map(|generation| generation.get()),
Some(1),
);
let profiles =
snapshot.public_model_profiles_for_provider(&ProviderKind::new("alpha").expect("provider"));
assert_eq!(
profiles
.iter()
.map(|profile| profile.model().as_str())
.collect::<Vec<_>>(),
vec!["public-model", "upstream-model"],
);
}
#[test]
@@ -661,7 +661,8 @@ pub struct OpsErrorRecord {
pub upstream_model_id: Option<String>,
pub upstream_transport: Option<String>,
pub failure_kind: String,
pub status_code: Option<u16>,
pub client_status_code: Option<u16>,
pub upstream_status_code: Option<u16>,
pub provider_error_code: Option<String>,
pub client_response_id: Option<String>,
pub upstream_request_id: Option<String>,
@@ -1627,7 +1628,8 @@ fn admin_ops_error(error: OpsErrorRecord) -> admin_observability::OpsError {
upstream_model_id: error.upstream_model_id,
upstream_transport: error.upstream_transport,
failure_kind: error.failure_kind,
status_code: error.status_code,
client_status_code: error.client_status_code,
upstream_status_code: error.upstream_status_code,
provider_error_code: error.provider_error_code,
client_response_id: error.client_response_id,
upstream_request_id: error.upstream_request_id,
@@ -1978,8 +1980,12 @@ async fn provider_account_metrics(
let row = sqlx::query(
"select count(*)::bigint as total,
count(*) filter (where enabled)::bigint as enabled,
count(*) filter (where not enabled or availability <> 'ready')::bigint
as unavailable,
count(*) filter (
where not (
enabled and availability = 'ready'
and (access_token_expires_at is null or access_token_expires_at > $1)
)
)::bigint as unavailable,
count(*) filter (
where enabled and availability = 'ready'
and (access_token_expires_at is null or access_token_expires_at > $1)
@@ -2922,7 +2928,7 @@ const OPS_ERRORS_CTE: &str = "with errors as (
mr.provider_kind,
mr.provider_account_ref, mr.upstream_model_id, mr.upstream_transport,
coalesce(mr.error_kind, 'failed') as failure_kind,
coalesce(mr.upstream_status_code, mr.client_status_code) as status_code,
mr.client_status_code, mr.upstream_status_code,
mr.provider_error_code, mr.client_response_id, mr.upstream_request_id,
mr.latency_ms, coalesce(mr.error_message, mr.error_kind, 'request failed') as message,
1::integer as occurrence_count,
@@ -2933,7 +2939,8 @@ const OPS_ERRORS_CTE: &str = "with errors as (
select 'ops_event'::text, oe.id, oe.model_request_id, oe.attempt_index,
mr.client_api_key_ref, oe.component, oe.operation,
oe.provider_kind, oe.provider_account_ref,
oe.upstream_model_id, null::text, oe.failure_kind, oe.status_code,
oe.upstream_model_id, null::text, oe.failure_kind,
null::integer as client_status_code, oe.status_code as upstream_status_code,
oe.provider_error_code, mr.client_response_id, oe.upstream_request_id,
oe.latency_ms, oe.message, oe.occurrence_count, oe.created_at,
'ops_event:' || oe.id
@@ -2948,7 +2955,8 @@ async fn list_ops_errors(pool: &PgPool, query: OpsErrorQuery) -> StoreResult<Ops
" select source, event_id, request_id, attempt_index, client_api_key_ref,
component, operation, provider_kind, provider_account_ref,
upstream_model_id, upstream_transport,
failure_kind, status_code, provider_error_code, client_response_id,
failure_kind, client_status_code, upstream_status_code,
provider_error_code, client_response_id,
upstream_request_id, latency_ms, message, occurrence_count, occurred_at,
stable_sort_id
from errors e where e.occurred_at >= ",
@@ -3047,7 +3055,7 @@ fn push_ops_filter(statement: &mut QueryBuilder<Postgres>, filter: &OpsErrorFilt
statement.push_bind(i32::try_from(index).unwrap_or(i32::MAX));
}
if let Some(status) = filter.status_code {
statement.push(" and e.status_code = ");
statement.push(" and e.upstream_status_code = ");
statement.push_bind(i32::from(status));
}
if let Some(search) = &filter.search {
@@ -3148,7 +3156,8 @@ fn ops_error_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<OpsErrorRecord
upstream_model_id: get(row, "upstream_model_id")?,
upstream_transport: get(row, "upstream_transport")?,
failure_kind: get(row, "failure_kind")?,
status_code: optional_status(row, "status_code")?,
client_status_code: optional_status(row, "client_status_code")?,
upstream_status_code: optional_status(row, "upstream_status_code")?,
provider_error_code: get(row, "provider_error_code")?,
client_response_id: get(row, "client_response_id")?,
upstream_request_id: get(row, "upstream_request_id")?,
@@ -1700,10 +1700,8 @@ fn admin_account_status(
.is_some_and(|expires_at| expires_at <= now)
{
AdminAccountStatus::Expired
} else if account.availability == ProviderAccountAvailability::Ready {
AdminAccountStatus::Active
} else {
AdminAccountStatus::Attention
AdminAccountStatus::Active
}
}
@@ -1728,7 +1726,7 @@ fn admin_account_summary(
.count(),
)
.unwrap_or(u64::MAX);
let attention = u64::try_from(
let unavailable = u64::try_from(
accounts
.iter()
.filter(|account| {
@@ -1746,7 +1744,7 @@ fn admin_account_summary(
total,
active,
quota_exhausted,
attention,
unavailable,
}
}
@@ -1797,7 +1795,6 @@ const fn admin_account_status_name(status: AdminAccountStatus) -> &'static str {
AdminAccountStatus::QuotaExhausted => "quota_exhausted",
AdminAccountStatus::Disabled => "disabled",
AdminAccountStatus::Banned => "banned",
AdminAccountStatus::Attention => "attention",
}
}
@@ -1,4 +1,4 @@
//! PostgreSQL 账号 cooldown 事实的可丢失 Redis 热缓存。
//! 请求调度用的可丢失账号 cooldown Redis 存储。
use std::{sync::Arc, time::SystemTime};
@@ -243,7 +243,7 @@ impl ProviderCooldownPort for RedisCredentialCooldownRepository {
}
}
/// PostgreSQL 账号事实到可丢失 cooldown cache 的 write-through adapter。
/// 账号状态写入到可丢失 cooldown cache 的桥接 adapter。
pub struct CooldownCachingProviderAccountStore {
authoritative: Arc<dyn ProviderAccountStore>,
cooldowns: Arc<dyn CredentialCooldownRepository>,
@@ -62,6 +62,57 @@ fn postgres_admin_observability_adapter_implements_terminal_port() {
assert_port::<PgAdminObservabilityStore>();
}
#[tokio::test]
async fn dashboard_account_metrics_should_partition_available_and_unavailable_accounts() {
let Some(database) = TestDatabase::create("dashboard_account_metrics").await else {
return;
};
let now = Utc::now();
sqlx::query(
"insert into provider_accounts (
id, provider_kind, name, upstream_user_id, authentication_kind,
provider_credentials_json, credential_revision, has_refresh_token,
access_token_expires_at, enabled, availability,
availability_observed_at, created_at, updated_at
) values
('acct_available', 'openai', 'available', 'user-available', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', $1, $1, $1),
('acct_expired', 'openai', 'expired', 'user-expired', 'oauth',
'{}'::jsonb, 1, false, $1 - interval '1 day', true, 'ready', $1, $1, $1),
('acct_banned', 'xai', 'banned', 'user-banned', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'banned', $1, $1, $1),
('acct_disabled', 'xai', 'disabled', 'user-disabled', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready', $1, $1, $1)",
)
.bind(now)
.execute(&database.pool)
.await
.expect("seed account metric states");
let repository = PgObservabilityRepository::new(database.pool.clone());
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
.expect("dashboard range");
let metrics = repository
.dashboard_summary(range)
.await
.expect("dashboard summary")
.provider_accounts;
assert_eq!(
(
metrics.total,
metrics.enabled,
metrics.active,
metrics.unavailable,
metrics.expired,
metrics.disabled,
metrics.banned,
),
(4, 3, 1, 3, 1, 1, 1),
);
database.close().await;
}
#[tokio::test]
async fn calculated_usage_billing_facts_keep_only_completed_calculated_costs() {
let Some(database) = TestDatabase::create("calculated_usage_billing_facts").await else {
@@ -589,6 +640,20 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
.await
.expect("ops errors");
assert_eq!(errors.total, 2);
let request_error = errors
.items
.iter()
.find(|error| error.source == "model_request")
.expect("request error");
assert_eq!(request_error.client_status_code, Some(502));
assert_eq!(request_error.upstream_status_code, Some(429));
let attempt_error = errors
.items
.iter()
.find(|error| error.source == "ops_event")
.expect("attempt error");
assert_eq!(attempt_error.client_status_code, None);
assert_eq!(attempt_error.upstream_status_code, Some(429));
database.close().await;
}
@@ -126,13 +126,13 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
beta.availability = ProviderAccountAvailability::Banned;
let mut charlie = account("acct_charlie", "user-charlie");
charlie.email = Some("charlie@example.invalid".to_owned());
let mut attention = account("acct_attention", "user-attention");
attention.email = Some("attention@example.invalid".to_owned());
attention.availability = ProviderAccountAvailability::Invalid;
let mut invalid = account("acct_invalid", "user-invalid");
invalid.email = Some("invalid@example.invalid".to_owned());
invalid.availability = ProviderAccountAvailability::Invalid;
let mut quota_exhausted = account("acct_quota_exhausted", "user-quota-exhausted");
quota_exhausted.email = Some("quota-exhausted@example.invalid".to_owned());
quota_exhausted.availability = ProviderAccountAvailability::QuotaExhausted;
for account in [alpha, beta, charlie, attention, quota_exhausted] {
for account in [alpha, beta, charlie, invalid, quota_exhausted] {
repository
.insert_provider_account(account)
.await
@@ -214,9 +214,8 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
assert_eq!(usage_page.config_revision.get(), 1);
assert_eq!(usage_page.total, 5);
assert_eq!(usage_page.summary.total, 5);
assert_eq!(usage_page.summary.active, 2);
assert_eq!(usage_page.summary.active, 3);
assert_eq!(usage_page.summary.quota_exhausted, 1);
assert_eq!(usage_page.summary.attention, 1);
assert_eq!(
usage_page
.items
@@ -277,19 +276,6 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
.await
.expect("filter banned accounts");
assert_eq!(banned.items[0].id, "acct_beta");
let attention = store
.list_accounts(AccountListQuery {
page: 1,
page_size: PageSize::new(10).expect("page size"),
provider_kind: None,
search: None,
status: Some(AccountStatus::Attention),
sort: None,
})
.await
.expect("filter attention accounts");
assert_eq!(attention.items[0].id, "acct_attention");
database.close().await;
}
@@ -12,9 +12,10 @@ use gateway_core::engine::credential::{
AccountAvailability, CredentialRevision, ProviderAccountId,
};
use gateway_core::provider_ports::{
ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest,
ProviderRefreshCapacityRequest, ProviderRefreshLeaseRequest, ProviderRefreshPolicy,
ProviderRuntimePolicyPort, ProviderStoreError, provider_refresh_retry_at,
ProviderCooldown, ProviderCooldownPort, ProviderLeaseAcquisition, ProviderLeasePort,
ProviderLeaseRequest, ProviderRefreshCapacityRequest, ProviderRefreshLeaseRequest,
ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderStoreError,
provider_refresh_retry_at,
};
use super::catalog::GrokCredentialCatalogService;
@@ -36,6 +37,7 @@ const MAX_REFRESH_EXCLUSIONS: usize = 400;
const MAX_SECRET_BYTES: usize = 64 * 1_024;
const DISCOVERY_NEGATIVE_CACHE_TTL: Duration = Duration::from_secs(5);
const REFRESH_RETRY_DELAY: Duration = Duration::from_secs(10 * 60);
const REFRESH_AMBIGUITY_COOLDOWN: Duration = Duration::from_secs(24 * 60 * 60);
/// 一个到期且已按 revision 读取明文 RT 的 xAI account。
pub struct DueGrokCredential {
@@ -274,6 +276,7 @@ pub struct GrokCredentialRefreshService {
refresher: Arc<dyn GrokCredentialRefresher>,
catalog: Arc<GrokCredentialCatalogService>,
leases: Arc<dyn ProviderLeasePort>,
cooldowns: Arc<dyn ProviderCooldownPort>,
runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
}
@@ -283,6 +286,7 @@ impl GrokCredentialRefreshService {
refresher: Arc<dyn GrokCredentialRefresher>,
catalog: Arc<GrokCredentialCatalogService>,
leases: Arc<dyn ProviderLeasePort>,
cooldowns: Arc<dyn ProviderCooldownPort>,
runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
) -> Self {
Self {
@@ -290,6 +294,7 @@ impl GrokCredentialRefreshService {
refresher,
catalog,
leases,
cooldowns,
runtime_policy,
}
}
@@ -297,9 +302,6 @@ impl GrokCredentialRefreshService {
pub async fn refresh_due(
&self,
) -> Result<Vec<GrokCredentialRefreshOutcome>, GrokCredentialRefreshError> {
if !self.prepare_cycle_if_due().await? {
return Ok(Vec::new());
}
self.refresh_due_excluding(&[]).await
}
@@ -455,8 +457,21 @@ impl GrokCredentialRefreshService {
.into_iter()
.map(|account_id| GrokCredentialRefreshOutcome::Failed { account_id })
.collect::<Vec<_>>();
outcomes.reserve(batch.credentials.len());
let refreshed = stream::iter(batch.credentials.into_iter().map(|credential| async move {
let mut credentials = Vec::with_capacity(batch.credentials.len());
for credential in batch.credentials {
if !self.runtime_cooldown_active(&credential).await {
credentials.push(credential);
}
}
if credentials.is_empty() {
return Ok(outcomes);
}
self.refresher
.prepare_cycle()
.await
.map_err(|_| GrokCredentialRefreshError::Preparation)?;
outcomes.reserve(credentials.len());
let refreshed = stream::iter(credentials.into_iter().map(|credential| async move {
let account_id = credential.account_id.clone();
(
account_id,
@@ -539,14 +554,11 @@ impl GrokCredentialRefreshService {
)
.await
}
Err(GrokRefreshFailure::Ambiguous) => self
.persist_terminal_failure(
credential,
GrokCredentialAvailability::Invalid,
"refresh_ambiguous",
)
.await
.map(|_| GrokCredentialRefreshOutcome::Ambiguous { account_id }),
Err(GrokRefreshFailure::Ambiguous) => {
self.record_runtime_cooldown(&credential, REFRESH_AMBIGUITY_COOLDOWN)
.await;
Ok(GrokCredentialRefreshOutcome::Ambiguous { account_id })
}
Err(GrokRefreshFailure::Transient) => {
if self
.persist_backoff(&credential, "refresh-transient")
@@ -644,6 +656,36 @@ impl GrokCredentialRefreshService {
})
}
async fn runtime_cooldown_active(&self, credential: &DueGrokCredential) -> bool {
let Ok(Some(cooldown)) = self.cooldowns.read(credential.account_id()).await else {
return false;
};
if cooldown.credential_revision() != credential.credential_revision() {
if cooldown.credential_revision() < credential.credential_revision() {
let _ = self
.cooldowns
.clear(credential.account_id(), credential.credential_revision())
.await;
}
return false;
}
cooldown.until() > SystemTime::now()
}
async fn record_runtime_cooldown(&self, credential: &DueGrokCredential, duration: Duration) {
let cooldown_until = SystemTime::now()
.checked_add(duration)
.unwrap_or_else(SystemTime::now);
let _ = self
.cooldowns
.put_if_later(ProviderCooldown::new(
credential.account_id.clone(),
credential.credential_revision,
cooldown_until,
))
.await;
}
async fn persist_terminal_failure(
&self,
credential: DueGrokCredential,
@@ -721,10 +763,12 @@ impl GrokCredentialRecovery for GrokCredentialRefreshService {
GrokCredentialRefreshOutcome::Refreshed { .. }
| GrokCredentialRefreshOutcome::Stale { .. },
) => GrokCredentialRecoveryOutcome::Recovered,
Ok(
GrokCredentialRefreshOutcome::Invalidated { .. }
| GrokCredentialRefreshOutcome::Ambiguous { .. },
) => GrokCredentialRecoveryOutcome::Rejected,
Ok(GrokCredentialRefreshOutcome::Invalidated { .. }) => {
GrokCredentialRecoveryOutcome::Rejected
}
Ok(GrokCredentialRefreshOutcome::Ambiguous { .. }) => {
GrokCredentialRecoveryOutcome::Unavailable
}
Ok(
GrokCredentialRefreshOutcome::Transient { .. }
| GrokCredentialRefreshOutcome::Rejected { .. }
@@ -7,10 +7,11 @@ use std::time::{Duration, SystemTime};
use chrono::Utc;
use gateway_core::engine::credential::{
AccountCandidate, AccountFeedbackStats, AccountSelectionContext, AccountSelector,
ProviderAccount, ProviderAccountId,
};
use gateway_core::provider_ports::{
ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest,
ProviderSchedulingLeaseRequest,
ProviderCooldown, ProviderCooldownPort, ProviderLeaseAcquisition, ProviderLeasePort,
ProviderLeaseRequest, ProviderSchedulingLeaseRequest,
};
use gateway_core::routing::ProviderKind;
@@ -26,6 +27,7 @@ use crate::{
const DEFAULT_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60);
const MAX_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(24 * 60 * 60);
const STREAM_INTERRUPTION_COOLDOWN: Duration = Duration::from_secs(30);
const UNAUTHORIZED_COOLDOWN: Duration = Duration::from_secs(60);
/// 仅经 Core account port、TTL catalog cache 和 Redis lease 选择一个 OAuth session。
pub struct GrokAccountSessionSelector {
@@ -34,6 +36,7 @@ pub struct GrokAccountSessionSelector {
catalog_cache: Arc<dyn GrokCredentialCatalogCache>,
quota: Arc<GrokCredentialQuotaService>,
scheduling: Arc<dyn ProviderLeasePort>,
cooldowns: Arc<dyn ProviderCooldownPort>,
account_feedback: Arc<AccountFeedbackStats>,
}
@@ -45,6 +48,7 @@ impl GrokAccountSessionSelector {
catalog_cache: Arc<dyn GrokCredentialCatalogCache>,
quota: Arc<GrokCredentialQuotaService>,
scheduling: Arc<dyn ProviderLeasePort>,
cooldowns: Arc<dyn ProviderCooldownPort>,
account_feedback: Arc<AccountFeedbackStats>,
) -> Self {
Self {
@@ -53,6 +57,7 @@ impl GrokAccountSessionSelector {
catalog_cache,
quota,
scheduling,
cooldowns,
account_feedback,
}
}
@@ -94,8 +99,10 @@ impl GrokAccountSessionSelector {
.load_state(&self.provider_kind, &account_ids)
.await
.map_err(|_| GrokSessionSelectorError::Unavailable)?;
let runtime_cooldowns = self.runtime_cooldowns(&catalog_eligible).await;
let mut candidates = catalog_eligible
.into_iter()
.filter(|account| !runtime_cooldowns.contains(account.id()))
.map(|account| {
let health = self
.account_feedback
@@ -195,6 +202,43 @@ impl GrokAccountSessionSelector {
Err(GrokSessionSelectorError::NoEligibleSession)
}
}
async fn runtime_cooldowns(
&self,
accounts: &[ProviderAccount],
) -> std::collections::BTreeSet<ProviderAccountId> {
let now = SystemTime::now();
let mut cooled = std::collections::BTreeSet::new();
for account in accounts {
let Ok(Some(cooldown)) = self.cooldowns.read(account.id()).await else {
continue;
};
if cooldown.credential_revision() != account.revision() {
if cooldown.credential_revision() < account.revision() {
let _ = self.cooldowns.clear(account.id(), account.revision()).await;
}
continue;
}
if cooldown.until() > now {
cooled.insert(account.id().clone());
}
}
cooled
}
async fn record_runtime_cooldown(&self, session: &SelectedGrokSession, duration: Duration) {
let cooldown_until = SystemTime::now()
.checked_add(duration)
.unwrap_or_else(SystemTime::now);
let _ = self
.cooldowns
.put_if_later(ProviderCooldown::new(
session.account_id().clone(),
session.credential_revision(),
cooldown_until,
))
.await;
}
}
impl fmt::Debug for GrokAccountSessionSelector {
@@ -205,6 +249,7 @@ impl fmt::Debug for GrokAccountSessionSelector {
.field("catalog_cache", &"[TTL_CACHE]")
.field("quota", &self.quota)
.field("scheduling", &"[SCHEDULING_PORT]")
.field("cooldowns", &"[COOLDOWN_PORT]")
.finish()
}
}
@@ -221,12 +266,12 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
) -> GrokCredentialFeedbackFuture<'a> {
Box::pin(async move {
let observed_at = Utc::now();
let (availability, reason, cooldown_until) = match failure {
GrokCredentialFailure::Unauthorized => (
GrokCredentialAvailability::Invalid,
"upstream_unauthorized",
None,
),
let persistent = match failure {
GrokCredentialFailure::Unauthorized => {
self.record_runtime_cooldown(session, UNAUTHORIZED_COOLDOWN)
.await;
return;
}
GrokCredentialFailure::QuotaExhausted => (
GrokCredentialAvailability::QuotaExhausted,
"upstream_quota_exhausted",
@@ -236,26 +281,16 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
let retry_after = retry_after
.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)
.min(MAX_RATE_LIMIT_COOLDOWN);
let cooldown_until = chrono::Duration::from_std(retry_after)
.ok()
.and_then(|duration| observed_at.checked_add_signed(duration));
(
GrokCredentialAvailability::Cooldown,
"upstream_rate_limited",
cooldown_until,
)
self.record_runtime_cooldown(session, retry_after).await;
return;
}
GrokCredentialFailure::StreamInterrupted => {
let cooldown_until = chrono::Duration::from_std(STREAM_INTERRUPTION_COOLDOWN)
.ok()
.and_then(|duration| observed_at.checked_add_signed(duration));
(
GrokCredentialAvailability::Cooldown,
"upstream_stream_interrupted",
cooldown_until,
)
self.record_runtime_cooldown(session, STREAM_INTERRUPTION_COOLDOWN)
.await;
return;
}
};
let (availability, reason, cooldown_until) = persistent;
let _ = self
.repository
.update_state(&UpdateGrokCredentialState {
+3
View File
@@ -103,6 +103,7 @@ pub async fn initialize(
ProviderKind::new(XAI_PROVIDER_NAME).map_err(|_| XaiInitializeError::ProviderKind)?;
let accounts: Arc<dyn ProviderAccountStore> = ports.accounts();
let leases = ports.leases();
let cooldowns = ports.cooldowns();
let account_feedback = ports.account_feedback();
let runtime_policy = ports.runtime_policy();
let repository = GrokCredentialRepository::new(Arc::clone(&accounts));
@@ -135,6 +136,7 @@ pub async fn initialize(
catalog_cache,
Arc::clone(&quota),
Arc::clone(&leases),
Arc::clone(&cooldowns),
Arc::clone(&account_feedback),
));
let inference: Arc<dyn GrokInferenceTransport> = Arc::new(
@@ -162,6 +164,7 @@ pub async fn initialize(
refresher,
Arc::clone(&catalog),
Arc::clone(&leases),
cooldowns,
Arc::clone(&runtime_policy),
));
let credential_recovery: Arc<dyn GrokCredentialRecovery> = refresh.clone();
+161 -58
View File
@@ -18,13 +18,16 @@ use gateway_core::engine::provider::{
UpstreamTransport,
};
use gateway_core::engine::{AttemptContext, ContinuationAttempt, UpstreamSendState};
use gateway_core::error::{ContinuationFailure, ProviderError, ProviderErrorKind};
use gateway_core::error::{
ClientVisibleUpstreamError, ContinuationFailure, ProviderError, ProviderErrorKind,
};
use gateway_core::event::{GatewayEvent, ProviderEvent, ProviderResponseObservation, ResponseMeta};
use gateway_core::operation::{
Feature, GenerateRequest, Operation, OperationKind, ProviderSessionState,
};
use gateway_core::routing::{
ModelCapabilities, ProviderCandidate, ProviderKind, SupportLevel, UpstreamModelId,
ModelCapabilities, ModelPresentation, ProviderCandidate, ProviderKind, SupportLevel,
UpstreamModelId,
};
use gateway_core::task::{
ScheduledTask, WorkerContribution, WorkerCycleContext, WorkerDefinitionError, WorkerId,
@@ -35,7 +38,6 @@ use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use url::Url;
use crate::GrokCatalogCapabilityEvidence;
use crate::XaiWireProfileState;
use crate::credential::{
GrokCredentialCatalogError, GrokCredentialCatalogService, GrokCredentialQuotaService,
@@ -54,8 +56,10 @@ use crate::transport::{
GrokRequestEncodeError, GrokResponsesRequest, GrokSessionAffinityKey, GrokSessionSelection,
GrokSessionSelector, GrokSessionSelectorError, SelectedGrokSession,
};
use crate::{GrokCatalogCapabilityEvidence, GrokCatalogModel};
const HTTP_SSE_TRANSPORT: &str = "http_sse";
const DEFAULT_GROK_MODEL: &str = "grok-4.5";
const XAI_SESSION_STATE_MAX_BYTES: usize = 8 * 1024 * 1024;
const XAI_SESSION_OUTPUT_LIMIT: usize = 4_096;
const MIN_REASONING_CIPHERTEXT_BYTES: usize = 50;
@@ -138,38 +142,15 @@ impl Provider for GrokBuildProvider {
async fn query_model_capabilities(
&self,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
let models = self.catalog.query_models().await.map_err(|_| {
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
})?;
let Ok(models) = self.catalog.query_models().await else {
return Ok(vec![default_grok_model_capabilities()?]);
};
if models.is_empty() {
return Ok(vec![default_grok_model_capabilities()?]);
}
Ok(models
.into_iter()
.map(|model| {
let mut operations = BTreeSet::new();
if model.capabilities().responses_api()
== GrokCatalogCapabilityEvidence::DeclaredNative
{
operations.insert(OperationKind::Generate);
}
let capabilities = ModelCapabilities::new(
operations,
model
.limits()
.max_output_tokens()
.map(std::num::NonZeroU64::get),
)
.with_feature(
Feature::Reasoning,
support(model.capabilities().reasoning_effort()),
)
.with_feature(
Feature::Tools,
tool_support(model.capabilities().streaming_tool_calls()),
)
.with_feature(Feature::Vision, SupportLevel::Unknown)
.with_feature(Feature::JsonSchema, SupportLevel::Unknown)
.with_feature(Feature::NativeContinuation, SupportLevel::Native);
ProviderModelCapabilities::new(model.request_model().clone(), capabilities)
})
.map(compile_grok_model_capabilities)
.collect())
}
@@ -403,6 +384,112 @@ fn tool_support(evidence: GrokCatalogCapabilityEvidence) -> SupportLevel {
}
}
fn compile_grok_model_capabilities(model: GrokCatalogModel) -> ProviderModelCapabilities {
let mut operations = BTreeSet::new();
if model.capabilities().responses_api() == GrokCatalogCapabilityEvidence::DeclaredNative {
operations.insert(OperationKind::Generate);
}
let capabilities = ModelCapabilities::new(
operations,
model
.limits()
.max_output_tokens()
.map(std::num::NonZeroU64::get),
)
.with_feature(
Feature::Reasoning,
support(model.capabilities().reasoning_effort()),
)
.with_feature(
Feature::Tools,
tool_support(model.capabilities().streaming_tool_calls()),
)
.with_feature(Feature::Vision, SupportLevel::Unknown)
.with_feature(Feature::JsonSchema, SupportLevel::Unknown)
.with_feature(Feature::NativeContinuation, SupportLevel::Native);
ProviderModelCapabilities::new(model.request_model().clone(), capabilities)
.with_presentation(grok_model_presentation(&model))
}
fn default_grok_model_capabilities() -> Result<ProviderModelCapabilities, ProviderError> {
let model = UpstreamModelId::new(DEFAULT_GROK_MODEL.to_owned())
.map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?;
let capabilities = ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), None)
.with_upstream_feature_validation();
Ok(ProviderModelCapabilities::new(model, capabilities)
.with_presentation(default_grok_model_presentation()))
}
fn default_grok_model_presentation() -> ModelPresentation {
ModelPresentation::new(
Some("Grok 4.5".to_owned()),
Some("xAI Grok 4.5 frontier model with reasoning and vision.".to_owned()),
)
.with_reasoning(
Some("medium".to_owned()),
["low", "medium", "high"].map(str::to_owned).to_vec(),
)
.with_context_window_tokens(Some(500_000))
.with_image_input(true)
.with_agent_tools(true, true)
}
fn grok_model_presentation(model: &GrokCatalogModel) -> ModelPresentation {
let slug = model.request_model().as_str();
let known_grok_4_5 = matches!(
slug,
DEFAULT_GROK_MODEL | "grok-4.5-latest" | "grok-4.5-build-free" | "grok-build-latest"
);
let reasoning_evidence = model.capabilities().reasoning_effort();
let (default_reasoning, reasoning_efforts) = match reasoning_evidence {
GrokCatalogCapabilityEvidence::DeclaredNative => (
Some("medium".to_owned()),
["low", "medium", "high"].map(str::to_owned).to_vec(),
),
GrokCatalogCapabilityEvidence::DeclaredUnsupported => {
(Some("none".to_owned()), vec!["none".to_owned()])
}
GrokCatalogCapabilityEvidence::Unknown if known_grok_4_5 => (
Some("medium".to_owned()),
["low", "medium", "high"].map(str::to_owned).to_vec(),
),
GrokCatalogCapabilityEvidence::Unknown => (None, Vec::new()),
};
let context_window_tokens = model
.limits()
.context_window_tokens()
.map(std::num::NonZeroU64::get)
.or(known_grok_4_5.then_some(500_000));
let tool_evidence = model.capabilities().streaming_tool_calls();
ModelPresentation::new(
model
.display_name()
.map(str::to_owned)
.or_else(|| known_grok_4_5.then(|| "Grok 4.5".to_owned())),
model
.metadata()
.description()
.map(str::to_owned)
.or_else(|| {
known_grok_4_5
.then(|| "xAI Grok 4.5 frontier model with reasoning and vision.".to_owned())
}),
)
.with_reasoning(default_reasoning, reasoning_efforts)
.with_context_window_tokens(context_window_tokens)
.with_image_input(known_grok_4_5)
.with_agent_tools(
tool_evidence != GrokCatalogCapabilityEvidence::DeclaredUnsupported,
tool_evidence == GrokCatalogCapabilityEvidence::DeclaredNative
|| (known_grok_4_5 && tool_evidence == GrokCatalogCapabilityEvidence::Unknown),
)
.with_search_tool(
model.capabilities().backend_search() == GrokCatalogCapabilityEvidence::DeclaredNative,
)
.with_hidden(model.metadata().hidden().unwrap_or(false))
}
#[derive(Clone, Serialize, Deserialize)]
struct XaiSessionState {
account_id: String,
@@ -692,16 +779,12 @@ fn attach_xai_session_update(
.into_iter()
.map(XaiReplayItem::ClientInput),
);
transcript.extend(
capture
.output_items
.into_values()
.filter_map(|item| portable_output_item(item, false))
.map(|item| XaiReplayItem::AccountOutput {
account_id: capture.account_id.clone(),
item,
}),
);
transcript.extend(capture.output_items.into_values().map(|item| {
XaiReplayItem::AccountOutput {
account_id: capture.account_id.clone(),
item,
}
}));
let state = XaiSessionState {
account_id: capture.account_id,
session_id: capture.session_id,
@@ -722,7 +805,11 @@ fn capture_output_item(event: &ProviderEvent, capture: &mut GrokSessionCapture)
.or_else(|| wire.data().get("type").and_then(Value::as_str));
if event_type == Some("response.output_item.done")
&& capture.output_items.len() < XAI_SESSION_OUTPUT_LIMIT
&& let Some(item) = wire.data().get("item").cloned()
&& let Some(item) = wire
.data()
.get("item")
.cloned()
.and_then(|item| portable_output_item(item, false))
{
let index = wire
.data()
@@ -743,16 +830,17 @@ fn capture_output_item(event: &ProviderEvent, capture: &mut GrokSessionCapture)
.and_then(Value::as_array)
.filter(|output| !output.is_empty())
{
capture.output_items.clear();
capture.output_items.extend(
output
.iter()
.take(XAI_SESSION_OUTPUT_LIMIT)
.enumerate()
.filter_map(|(index, item)| {
u32::try_from(index).ok().map(|index| (index, item.clone()))
}),
);
for (index, item) in output.iter().take(XAI_SESSION_OUTPUT_LIMIT).enumerate() {
let Some(index) = u32::try_from(index).ok() else {
break;
};
if capture.output_items.contains_key(&index) {
continue;
}
if let Some(item) = portable_output_item(item.clone(), false) {
capture.output_items.insert(index, item);
}
}
}
}
@@ -1351,8 +1439,10 @@ async fn recover_or_record_failure(
.await
{
GrokCredentialRecoveryOutcome::Recovered => error.with_same_account_retry(),
GrokCredentialRecoveryOutcome::Rejected
| GrokCredentialRecoveryOutcome::Unavailable => error,
GrokCredentialRecoveryOutcome::Rejected => error,
GrokCredentialRecoveryOutcome::Unavailable => {
record_failure(selector, session, error).await
}
};
}
record_failure(selector, session, error).await
@@ -1415,10 +1505,23 @@ fn ensure_sent_context(context: &AttemptContext) -> Result<(), ProviderError> {
fn map_request_error(error: GrokRequestEncodeError) -> ProviderError {
let kind = match error {
GrokRequestEncodeError::InvalidProtocolPayload
| GrokRequestEncodeError::InvalidRequestNormalization => ProviderErrorKind::InvalidRequest,
| GrokRequestEncodeError::InvalidRequestNormalization
| GrokRequestEncodeError::InvalidRequestField { .. } => ProviderErrorKind::InvalidRequest,
GrokRequestEncodeError::Serialization => ProviderErrorKind::Protocol,
};
provider_error(kind, UpstreamSendState::NotSent)
let provider_error = provider_error(kind, UpstreamSendState::NotSent);
if kind != ProviderErrorKind::InvalidRequest {
return provider_error;
}
let detail = ClientVisibleUpstreamError::new(
error.to_string(),
Some("invalid_request_normalization".to_owned()),
Some("invalid_request_error".to_owned()),
);
match detail {
Ok(detail) => provider_error.with_client_visible_upstream_error(detail),
Err(_) => provider_error,
}
}
fn map_selection_error(error: GrokSessionSelectorError) -> ProviderError {
@@ -356,12 +356,6 @@ impl GrokInferenceTransport for ReqwestGrokInferenceTransport {
if !response.status().is_success() {
return Err(classify_inference_status(response).await);
}
if !is_event_stream(&response) {
return Err(GrokInferenceTransportError::new(
GrokInferenceTransportErrorKind::Protocol,
UpstreamSendState::Sent,
));
}
let http_version = upstream_http_version(response.version());
let status_code = response.status().as_u16();
let request_id = upstream_request_id(&response);
@@ -1109,15 +1103,6 @@ fn retry_after(response: &Response) -> Option<Duration> {
.map(Duration::from_secs)
}
fn is_event_stream(response: &Response) -> bool {
response
.headers()
.get(CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.split(';').next())
.is_some_and(|value| value.trim().eq_ignore_ascii_case("text/event-stream"))
}
fn is_json_response(response: &Response) -> bool {
response
.headers()
@@ -1,4 +1,4 @@
use std::collections::BTreeMap;
use std::collections::{BTreeMap, BTreeSet};
use std::fmt;
use gateway_core::operation::GenerateRequest;
@@ -279,6 +279,18 @@ pub enum GrokRequestEncodeError {
/// Responses compatibility fields could not be normalized safely.
#[error("Grok Build request normalization failed")]
InvalidRequestNormalization,
/// 请求中的具体字段无法安全转换为 Grok Build 接受的形态。
#[error("Grok Build request field `{field}` could not be normalized safely")]
InvalidRequestField { field: &'static str },
}
impl GrokRequestEncodeError {
fn at_field(self, field: &'static str) -> Self {
match self {
Self::InvalidRequestNormalization => Self::InvalidRequestField { field },
error => error,
}
}
}
fn sanitize_account_identity(body: &mut Map<String, Value>) {
@@ -548,11 +560,12 @@ fn normalize_responses_request(
}
let text = text
.as_object_mut()
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
.ok_or(GrokRequestEncodeError::InvalidRequestField { field: "text" })?;
if text.get("format").is_none_or(Value::is_null) {
text.insert(
"format".to_owned(),
normalize_response_format(response_format)?,
normalize_response_format(response_format)
.map_err(|error| error.at_field("response_format"))?,
);
}
}
@@ -711,20 +724,22 @@ impl GrokResponseTransform {
.stream_calls
.get_mut(&primary)
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
if matches!(
state.identity.kind,
ToolKind::ToolSearch | ToolKind::Custom | ToolKind::ApplyPatch
) {
state.arguments.push_str(
value
.get("delta")
.and_then(Value::as_str)
.unwrap_or_default(),
);
if state.identity.kind == ToolKind::Custom {
state.last_delta = value.as_object().cloned();
match state.identity.kind {
ToolKind::Custom => {
state.arguments.push_str(
value
.get("delta")
.and_then(Value::as_str)
.unwrap_or_default(),
);
state.last_delta = value.as_object().cloned().map(|mut payload| {
payload.remove("delta");
payload
});
return Ok(Vec::new());
}
return Ok(Vec::new());
ToolKind::ToolSearch | ToolKind::ApplyPatch => return Ok(Vec::new()),
ToolKind::Function => {}
}
}
@@ -767,6 +782,10 @@ impl GrokResponseTransform {
&input,
),
});
if let Some(state) = self.stream_calls.get_mut(&primary) {
state.arguments.clear();
state.last_delta = None;
}
return Ok(output);
}
ToolKind::Function => {}
@@ -783,10 +802,17 @@ impl GrokResponseTransform {
return self.rewrite_apply_patch_done_event(value);
}
let completed_call = (event_type == "response.output_item.done")
.then(|| value.get("item").and_then(Value::as_object))
.flatten()
.and_then(|item| self.primary_stream_call(item));
self.rewrite_response_value(&mut value)?;
if let Some(response) = value.get_mut("response").and_then(Value::as_object_mut) {
self.restore_visible_tools(response);
}
if let Some(primary) = completed_call {
self.take_stream_call(&primary);
}
Ok(vec![GrokTransformedWireEvent {
event_type: event_type.to_owned(),
value,
@@ -847,6 +873,18 @@ impl GrokResponseTransform {
Some(primary)
}
fn primary_stream_call(&self, item: &Map<String, Value>) -> Option<String> {
["id", "call_id"]
.into_iter()
.filter_map(|field| item.get(field).and_then(Value::as_str))
.find_map(|key| self.stream_keys.get(key).cloned())
}
fn take_stream_call(&mut self, primary: &str) -> Option<StreamCallState> {
self.stream_keys.retain(|_, owner| owner != primary);
self.stream_calls.remove(primary)
}
fn rewrite_response_value(&self, value: &mut Value) -> Result<(), GrokRequestEncodeError> {
match value {
Value::Array(values) => {
@@ -953,15 +991,11 @@ impl GrokResponseTransform {
.and_then(Value::as_object)
.cloned()
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
let primary = ["id", "call_id"]
.into_iter()
.find_map(|field| original_item.get(field).and_then(Value::as_str))
.and_then(|key| self.stream_keys.get(key))
.cloned();
let mut added = primary
.as_ref()
.and_then(|key| self.stream_calls.get(key))
.and_then(|state| state.added_payload.clone())
let state = self
.primary_stream_call(&original_item)
.and_then(|primary| self.take_stream_call(&primary));
let mut added = state
.and_then(|state| state.added_payload)
.unwrap_or_else(|| {
Map::from_iter([(
"type".to_owned(),
@@ -1148,20 +1182,27 @@ impl ToolNormalizer {
mut self,
payload: &mut Map<String, Value>,
) -> Result<GrokResponseTransform, GrokRequestEncodeError> {
let (tools, had_tools) = optional_array(payload.get("tools"))?;
let (tools, had_tools) =
optional_array(payload.get("tools")).map_err(|error| error.at_field("tools"))?;
if had_tools {
self.response.visible_tools.clone_from(&tools);
}
let client_search = inspect_tool_search(&tools)?;
self.normalize_client_search_parallel(payload, client_search)?;
let client_search = inspect_tool_search(&tools).map_err(|error| error.at_field("tools"))?;
self.normalize_client_search_parallel(payload, client_search)
.map_err(|error| error.at_field("parallel_tool_calls"))?;
let mut normalized_tools = Vec::with_capacity(tools.len());
for raw_tool in &tools {
normalized_tools.extend(self.normalize_tool(raw_tool, "", client_search, false)?);
normalized_tools.extend(
self.normalize_tool(raw_tool, "", client_search, false)
.map_err(|error| error.at_field("tools"))?,
);
}
if let Some(Value::Array(items)) = payload.get("input") {
let normalized = self.normalize_input_items(items)?;
let normalized = self
.normalize_input_items(items)
.map_err(|error| error.at_field("input"))?;
normalized_tools.extend(normalized.loaded_tools);
self.response.visible_tools.extend(normalized.visible_tools);
payload.insert("input".to_owned(), Value::Array(normalized.items));
@@ -1169,11 +1210,14 @@ impl ToolNormalizer {
.get("input")
.is_some_and(|input| !input.is_null() && !input.is_string())
{
return Err(GrokRequestEncodeError::InvalidRequestNormalization);
return Err(GrokRequestEncodeError::InvalidRequestField { field: "input" });
}
if self.client_search_tool.is_some() {
normalized_tools.push(Value::Object(self.build_client_search_function()?));
normalized_tools.push(Value::Object(
self.build_client_search_function()
.map_err(|error| error.at_field("tools"))?,
));
}
normalized_tools = dedupe_normalized_tools(normalized_tools);
if normalized_tools.is_empty() {
@@ -1184,7 +1228,8 @@ impl ToolNormalizer {
} else {
payload.insert("tools".to_owned(), Value::Array(normalized_tools.clone()));
}
self.normalize_tool_choice(payload, &normalized_tools)?;
self.normalize_tool_choice(payload, &normalized_tools)
.map_err(|error| error.at_field("tool_choice"))?;
Ok(self.response)
}
@@ -1299,6 +1344,11 @@ impl ToolNormalizer {
return Ok(Vec::new());
}
let mut converted = without_defer_loading(tool);
if let Some(parameters) = converted.get("parameters").cloned()
&& let Some(normalized) = normalize_function_parameters_root(&parameters)?
{
converted.insert("parameters".to_owned(), normalized);
}
let alias = self.alias(ToolIdentity::new(ToolKind::Function, namespace, name));
converted.insert("name".to_owned(), Value::String(alias));
Ok(vec![Value::Object(converted)])
@@ -1448,6 +1498,136 @@ fn optional_array(value: Option<&Value>) -> Result<(Vec<Value>, bool), GrokReque
}
}
/// Grok Build 要求 function parameters 的根节点必为 object;Codex 会为可选参数
/// 生成 `object | null`,这里只移除根节点 nullability,嵌套字段保持原样。
fn normalize_function_parameters_root(
value: &Value,
) -> Result<Option<Value>, GrokRequestEncodeError> {
let Some(schema) = value.as_object() else {
return Ok(None);
};
let mut normalized = schema.clone();
let mut changed = false;
if let Some(types) = normalized.get("type").and_then(Value::as_array).cloned() {
let removed_null = types.iter().any(|value| value.as_str() == Some("null"));
if removed_null {
let remaining = types
.into_iter()
.filter(|value| value.as_str() != Some("null"))
.collect::<Vec<_>>();
if remaining.len() != 1 || remaining[0].as_str() != Some("object") {
return Err(invalid_function_parameters_root());
}
normalized.insert("type".to_owned(), Value::String("object".to_owned()));
changed = true;
}
}
for keyword in ["anyOf", "oneOf"] {
let Some(branches) = normalized.get(keyword).and_then(Value::as_array).cloned() else {
continue;
};
let removed_null = branches.iter().any(is_null_only_schema);
if !removed_null {
continue;
}
let remaining = branches
.into_iter()
.filter(|branch| !is_null_only_schema(branch))
.collect::<Vec<_>>();
if remaining.is_empty()
|| remaining.iter().any(|branch| {
branch.as_object().is_none_or(|branch| {
!is_object_root_schema(branch, &normalized, &mut BTreeSet::new())
})
})
{
return Err(invalid_function_parameters_root());
}
if remaining.len() == 1 && normalized.len() == 1 {
normalized = remaining[0]
.as_object()
.cloned()
.ok_or_else(invalid_function_parameters_root)?;
normalized.insert("type".to_owned(), Value::String("object".to_owned()));
} else {
normalized.insert(keyword.to_owned(), Value::Array(remaining));
normalized.insert("type".to_owned(), Value::String("object".to_owned()));
}
changed = true;
}
Ok(changed.then_some(Value::Object(normalized)))
}
fn is_null_only_schema(value: &Value) -> bool {
let Some(schema) = value.as_object() else {
return false;
};
match schema.get("type") {
Some(Value::String(kind)) => kind == "null",
Some(Value::Array(types)) if !types.is_empty() => {
types.iter().all(|kind| kind.as_str() == Some("null"))
}
_ => false,
}
}
fn is_object_root_schema(
schema: &Map<String, Value>,
root: &Map<String, Value>,
visited: &mut BTreeSet<String>,
) -> bool {
match schema.get("type") {
Some(Value::String(kind)) => return kind == "object",
Some(Value::Array(types)) if !types.is_empty() => {
return types.iter().all(|kind| kind.as_str() == Some("object"));
}
Some(_) => return false,
None => {}
}
if schema.contains_key("properties") {
return true;
}
let Some(reference) = schema.get("$ref").and_then(Value::as_str) else {
return false;
};
if !visited.insert(reference.to_owned()) {
return false;
}
resolve_local_schema_ref(root, reference)
.is_some_and(|resolved| is_object_root_schema(resolved, root, visited))
}
fn resolve_local_schema_ref<'a>(
root: &'a Map<String, Value>,
reference: &str,
) -> Option<&'a Map<String, Value>> {
if reference == "#" {
return Some(root);
}
let path = reference.strip_prefix("#/")?;
let mut segments = path.split('/');
let first = decode_json_pointer_segment(segments.next()?);
let mut current = root.get(&first)?;
for segment in segments {
let segment = decode_json_pointer_segment(segment);
current = current.as_object()?.get(&segment)?;
}
current.as_object()
}
fn decode_json_pointer_segment(segment: &str) -> String {
segment.replace("~1", "/").replace("~0", "~")
}
const fn invalid_function_parameters_root() -> GrokRequestEncodeError {
GrokRequestEncodeError::InvalidRequestField {
field: "tools[].parameters",
}
}
fn inspect_tool_search(tools: &[Value]) -> Result<bool, GrokRequestEncodeError> {
let mut client_search = false;
let mut server_search = false;
@@ -20,8 +20,8 @@ use provider_xai::{
};
use crate::support::{
MemoryGrokCatalogCache, MemoryProviderAccountStore, create_input, credential_object,
runtime_policy, seed_input,
MemoryCooldownPort, MemoryGrokCatalogCache, MemoryProviderAccountStore, create_input,
credential_object, runtime_policy, seed_input,
};
const OFFICIAL_FIXTURE: &[u8] =
@@ -169,11 +169,13 @@ async fn fixture_many(
available: lease_available,
calls: AtomicUsize::new(0),
});
let cooldowns = Arc::new(MemoryCooldownPort::default());
let service = GrokCredentialRefreshService::new(
repository.clone(),
refresher_port,
catalog,
leases,
cooldowns,
runtime_policy(),
);
(store, repository, refresher, service)
@@ -380,10 +382,11 @@ async fn banned_failure_marks_account_banned() {
}
#[tokio::test]
async fn ambiguous_refresh_fails_closed_and_is_not_retried() {
async fn ambiguous_refresh_uses_runtime_cooldown_without_persisting_account_state() {
let input = due_input("ambiguous");
let id = input.account_id.clone();
let (store, _, _, service) = fixture(input, [Err(GrokRefreshFailure::Ambiguous)], true).await;
let (store, _, refresher, service) =
fixture(input, [Err(GrokRefreshFailure::Ambiguous)], true).await;
let outcomes = service.refresh_due().await.expect("refresh");
assert!(matches!(
outcomes.as_slice(),
@@ -391,8 +394,9 @@ async fn ambiguous_refresh_fails_closed_and_is_not_retried() {
));
assert_eq!(
store.account(&id).expect("account").availability(),
AccountAvailability::Invalid
AccountAvailability::Unknown
);
assert_eq!(store.account(&id).expect("account").revision().get(), 1);
assert!(
service
.refresh_due()
@@ -400,6 +404,7 @@ async fn ambiguous_refresh_fails_closed_and_is_not_retried() {
.expect("second cycle")
.is_empty()
);
assert_eq!(refresher.prepare_calls.load(Ordering::SeqCst), 1);
}
#[tokio::test]
@@ -23,7 +23,8 @@ use provider_xai::{
};
use crate::support::{
MemoryGrokCatalogCache, MemoryProviderAccountStore, account_id, create_input, seed_input,
MemoryCooldownPort, MemoryGrokCatalogCache, MemoryProviderAccountStore, account_id,
create_input, seed_input,
};
struct SchedulingCoordinator {
@@ -88,6 +89,7 @@ struct SelectorFixture {
cache: Arc<MemoryGrokCatalogCache>,
selector: GrokAccountSessionSelector,
coordinator: Arc<SchedulingCoordinator>,
cooldowns: Arc<MemoryCooldownPort>,
feedback: Arc<AccountFeedbackStats>,
}
@@ -137,6 +139,7 @@ impl SelectorFixture {
signals: Mutex::new(signals),
denied: Mutex::new(BTreeSet::new()),
});
let cooldowns = Arc::new(MemoryCooldownPort::default());
let catalog_cache: Arc<dyn GrokCredentialCatalogCache> = cache.clone();
let lease_port: Arc<dyn ProviderLeasePort> = coordinator.clone();
let quota = Arc::new(crate::support::grok_quota_service(
@@ -150,6 +153,7 @@ impl SelectorFixture {
catalog_cache,
quota,
lease_port,
cooldowns.clone(),
Arc::clone(&feedback),
);
Self {
@@ -157,6 +161,7 @@ impl SelectorFixture {
cache,
selector,
coordinator,
cooldowns,
feedback,
}
}
@@ -257,7 +262,7 @@ async fn required_account_overrides_smart_selection_without_fallback() {
}
#[tokio::test]
async fn unauthorized_feedback_marks_only_selected_account_invalid() {
async fn unauthorized_feedback_records_runtime_cooldown_without_persisting_account_state() {
let fixture = SelectorFixture::new(&["feedback-a", "feedback-b"]).await;
let session = fixture
.selector
@@ -275,7 +280,7 @@ async fn unauthorized_feedback_marks_only_selected_account_invalid() {
.account(&selected)
.expect("selected")
.availability(),
AccountAvailability::Invalid
AccountAvailability::Ready
);
let other = [account_id("feedback-a"), account_id("feedback-b")]
.into_iter()
@@ -285,16 +290,78 @@ async fn unauthorized_feedback_marks_only_selected_account_invalid() {
fixture.store.account(&other).expect("other").availability(),
AccountAvailability::Ready
);
assert!(
fixture
.cooldowns
.cooldown(&selected)
.is_some_and(|cooldown| cooldown.until() > SystemTime::now())
);
}
#[tokio::test]
async fn rate_limit_feedback_applies_bounded_cooldown() {
let fixture = SelectorFixture::new(&["rate-limit"]).await;
async fn cooldown_for_superseded_credential_does_not_block_rotated_account() {
let fixture = SelectorFixture::new(&["revision-fence"]).await;
let id = account_id("revision-fence");
let session = fixture
.selector
.select(fixture.request(BTreeSet::new()))
.await
.expect("session");
fixture
.selector
.record_failure(
&session,
GrokCredentialFailure::RateLimited {
retry_after: Some(Duration::from_secs(60)),
},
)
.await;
let current = fixture
.store
.load_credential(&id, CredentialRevision::new(1).expect("revision"))
.await
.expect("current credential");
let prepared = GrokCredentialAdmin
.prepare_rotation(&RotateManagedGrokCredential {
current,
secret: provider_xai::GrokOAuthSecret {
access_token: provider_xai::SecretValue::new("new-access"),
refresh_token: provider_xai::SecretValue::new("new-refresh"),
id_token: None,
scope: provider_xai::OFFICIAL_SCOPES.join(" "),
},
verified_account: crate::support::profile("subject-revision-fence"),
next_refresh_at: chrono::Utc::now() + chrono::Duration::minutes(30),
})
.expect("rotate");
assert!(matches!(
fixture
.store
.compare_and_swap_credential(prepared.credential)
.await
.expect("persist rotation"),
CredentialCasOutcome::Updated(revision) if revision.get() == 2
));
let next = fixture
.selector
.select(fixture.request(BTreeSet::new()))
.await
.expect("replacement credential should not inherit cooldown");
assert_eq!(next.account_id(), &id);
assert!(fixture.cooldowns.cooldown(&id).is_none());
}
#[tokio::test]
async fn rate_limit_feedback_records_runtime_cooldown_without_persisting_account_state() {
let fixture = SelectorFixture::new(&["rate-limit", "available"]).await;
let session = fixture
.selector
.select(fixture.request(BTreeSet::new()))
.await
.expect("session");
let selected = session.account_id().clone();
fixture
.selector
.record_failure(
@@ -304,20 +371,24 @@ async fn rate_limit_feedback_applies_bounded_cooldown() {
},
)
.await;
let account = fixture
.store
.account(session.account_id())
.expect("account");
assert_eq!(account.availability(), AccountAvailability::Cooldown);
let account = fixture.store.account(&selected).expect("account");
assert_eq!(account.availability(), AccountAvailability::Ready);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
fixture
.cooldowns
.cooldown(&selected)
.is_some_and(|cooldown| cooldown.until() > SystemTime::now())
);
let next = fixture
.selector
.select(fixture.request(BTreeSet::new()))
.await
.expect("runtime cooldown should leave the other account available");
assert_ne!(next.account_id(), &selected);
}
#[tokio::test]
async fn interrupted_stream_feedback_applies_account_cooldown() {
async fn interrupted_stream_feedback_records_runtime_cooldown_without_persisting_account_state() {
let fixture = SelectorFixture::new(&["stream-interrupted"]).await;
let session = fixture
.selector
@@ -333,12 +404,20 @@ async fn interrupted_stream_feedback_applies_account_cooldown() {
.store
.account(session.account_id())
.expect("account");
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert_eq!(account.availability(), AccountAvailability::Ready);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
fixture
.cooldowns
.cooldown(session.account_id())
.is_some_and(|cooldown| cooldown.until() > SystemTime::now())
);
assert!(matches!(
fixture
.selector
.select(fixture.request(BTreeSet::new()))
.await,
Err(GrokSessionSelectorError::NoEligibleSession)
));
}
#[tokio::test]
@@ -33,7 +33,8 @@ use provider_xai::{
GrokCredentialFeedbackFuture, GrokCredentialRecovery, GrokCredentialRecoveryOutcome,
GrokCredentialRepository, GrokInferenceRequest, GrokInferenceResponse, GrokInferenceTransport,
GrokInferenceTransportError, GrokInferenceTransportErrorKind, GrokInferenceTransportFuture,
GrokModelCatalogRequest, GrokModelCatalogTransport, GrokModelCatalogTransportFuture,
GrokModelCatalogRequest, GrokModelCatalogTransport, GrokModelCatalogTransportError,
GrokModelCatalogTransportErrorKind, GrokModelCatalogTransportFuture,
GrokModelCatalogTransportResponse, GrokSessionBinding, GrokSessionSelection,
GrokSessionSelector, GrokSessionSelectorError, GrokSessionSelectorFuture, SecretValue,
SelectedGrokSession,
@@ -352,6 +353,18 @@ impl GrokModelCatalogTransport for CatalogWithoutToolMetadataTransport {
}
}
struct UnavailableCatalogTransport;
impl GrokModelCatalogTransport for UnavailableCatalogTransport {
fn execute(&self, _: GrokModelCatalogRequest) -> GrokModelCatalogTransportFuture<'_> {
Box::pin(async {
Err(GrokModelCatalogTransportError::new(
GrokModelCatalogTransportErrorKind::Unavailable,
))
})
}
}
struct StubRecovery {
calls: AtomicUsize,
outcome: GrokCredentialRecoveryOutcome,
@@ -506,6 +519,19 @@ fn operation_with_reasoning_effort(effort: &str) -> Operation {
Operation::Generate(GenerateRequest::from_protocol_payload(payload))
}
fn operation_with_invalid_tools() -> Operation {
let payload = ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("client-model")),
("input".to_owned(), json!("hello")),
("tools".to_owned(), json!({"not": "an array"})),
]),
)
.expect("OpenAI payload");
Operation::Generate(GenerateRequest::from_protocol_payload(payload))
}
fn compaction_operation() -> Operation {
let payload = ProtocolPayload::json_object(
"openai",
@@ -1047,6 +1073,38 @@ async fn first_unauthorized_should_refresh_and_request_one_same_account_retry()
assert!(selector.feedback.lock().expect("feedback").is_empty());
}
#[tokio::test]
async fn unavailable_unauthorized_recovery_records_temporary_credential_feedback() {
let selector = StubSelector::success();
let transport = StubInferenceTransport::error(
GrokInferenceTransportError::new(
GrokInferenceTransportErrorKind::Unauthorized,
UpstreamSendState::Sent,
)
.with_status(401)
.with_credential_recovery(),
);
let recovery = StubRecovery::new(GrokCredentialRecoveryOutcome::Unavailable);
let provider = provider_with_recovery(selector.clone(), transport, recovery.clone()).await;
let mut stream = provider
.execute(
provider_request("xai"),
context(CancellationToken::new(), None),
)
.await
.expect("stream");
let error = next_provider_error(&mut stream).await;
assert_eq!(error.kind(), ProviderErrorKind::Unauthorized);
assert!(!error.retries_same_account());
assert_eq!(recovery.calls.load(Ordering::SeqCst), 1);
assert_eq!(
selector.feedback.lock().expect("feedback").as_slice(),
&[GrokCredentialFailure::Unauthorized]
);
}
#[tokio::test]
async fn explicit_http_429_marks_provider_error_replay_safe() {
let error = mapped_transport_error(
@@ -1628,6 +1686,37 @@ async fn provider_rejects_target_owned_by_other_provider() {
assert_eq!(selector.calls.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn provider_should_expose_safe_request_normalization_field_before_selection() {
let selector = StubSelector::success();
let provider = provider(selector.clone(), StubInferenceTransport::success()).await;
let error = match provider
.execute(
provider_request_with_operation("xai", operation_with_invalid_tools()),
context(CancellationToken::new(), None),
)
.await
{
Ok(_) => panic!("invalid tools must fail"),
Err(error) => error,
};
assert_eq!(error.kind(), ProviderErrorKind::InvalidRequest);
assert_eq!(
error.client_visible_upstream_error().map(|detail| (
detail.message(),
detail.code(),
detail.error_type()
)),
Some((
"Grok Build request field `tools` could not be normalized safely",
Some("invalid_request_normalization"),
Some("invalid_request_error"),
)),
);
assert_eq!(selector.calls.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn provider_compiles_realtime_catalog_capabilities() {
let provider = provider(StubSelector::success(), StubInferenceTransport::success()).await;
@@ -1645,6 +1734,41 @@ async fn provider_compiles_realtime_catalog_capabilities() {
))
.is_some()
);
let presentation = capabilities[0]
.presentation()
.expect("Grok model presentation");
assert_eq!(presentation.display_name(), Some("Grok 4.5"));
assert_eq!(
presentation.supported_reasoning_efforts(),
["low", "medium", "high"]
);
assert_eq!(presentation.context_window_tokens(), Some(1_000_000));
assert!(presentation.agent_tools());
}
#[tokio::test]
async fn provider_publishes_default_codex_profile_when_catalog_is_unavailable() {
let provider = provider_with_catalog_transport(
StubSelector::success(),
StubInferenceTransport::success(),
StubRecovery::new(GrokCredentialRecoveryOutcome::Unavailable),
Arc::new(UnavailableCatalogTransport),
)
.await;
let capabilities = provider
.query_model_capabilities()
.await
.expect("fallback capabilities");
assert_eq!(capabilities.len(), 1);
assert_eq!(capabilities[0].upstream_model().as_str(), MODEL);
let presentation = capabilities[0]
.presentation()
.expect("fallback Grok model presentation");
assert_eq!(presentation.display_name(), Some("Grok 4.5"));
assert_eq!(presentation.context_window_tokens(), Some(500_000));
assert!(presentation.agent_tools());
assert!(presentation.parallel_tool_calls());
}
#[tokio::test]
+58 -3
View File
@@ -15,9 +15,9 @@ use gateway_core::engine::credential::{
};
use gateway_core::error::{StoreError, StoreErrorKind};
use gateway_core::provider_ports::{
ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderSessionAffinityKey,
ProviderSessionAffinityPort, ProviderSessionExclusionPort, ProviderSessionExclusions,
ProviderStoreError,
ProviderCooldown, ProviderCooldownPort, ProviderRefreshPolicy, ProviderRuntimePolicyPort,
ProviderSessionAffinityKey, ProviderSessionAffinityPort, ProviderSessionExclusionPort,
ProviderSessionExclusions, ProviderStoreError,
};
use gateway_core::routing::ProviderKind;
use provider_xai::{
@@ -550,6 +550,17 @@ pub struct TestSessionAffinity;
pub struct TestSessionExclusions;
#[derive(Default)]
pub struct MemoryCooldownPort {
cooldowns: Mutex<BTreeMap<ProviderAccountId, ProviderCooldown>>,
}
impl MemoryCooldownPort {
pub fn cooldown(&self, account_id: &ProviderAccountId) -> Option<ProviderCooldown> {
lock(&self.cooldowns).get(account_id).cloned()
}
}
impl ProviderSessionAffinityPort for TestSessionAffinity {
fn load<'a>(
&'a self,
@@ -613,6 +624,50 @@ impl ProviderSessionExclusionPort for TestSessionExclusions {
}
}
impl ProviderCooldownPort for MemoryCooldownPort {
fn put_if_later(
&self,
cooldown: ProviderCooldown,
) -> futures::future::BoxFuture<'_, Result<bool, ProviderStoreError>> {
Box::pin(async move {
let mut cooldowns = lock(&self.cooldowns);
let should_write = cooldowns.get(cooldown.account_id()).is_none_or(|current| {
current.credential_revision() < cooldown.credential_revision()
|| (current.credential_revision() == cooldown.credential_revision()
&& current.until() < cooldown.until())
});
if should_write {
cooldowns.insert(cooldown.account_id().clone(), cooldown);
}
Ok(should_write)
})
}
fn read<'a>(
&'a self,
account_id: &'a ProviderAccountId,
) -> futures::future::BoxFuture<'a, Result<Option<ProviderCooldown>, ProviderStoreError>> {
Box::pin(async move { Ok(self.cooldown(account_id)) })
}
fn clear<'a>(
&'a self,
account_id: &'a ProviderAccountId,
through_revision: CredentialRevision,
) -> futures::future::BoxFuture<'a, Result<bool, ProviderStoreError>> {
Box::pin(async move {
let mut cooldowns = lock(&self.cooldowns);
let should_remove = cooldowns
.get(account_id)
.is_some_and(|current| current.credential_revision() <= through_revision);
if should_remove {
cooldowns.remove(account_id);
}
Ok(should_remove)
})
}
}
impl ProviderRuntimePolicyPort for StaticRuntimePolicy {
fn load_refresh_policy(
&self,
@@ -114,6 +114,42 @@ fn decoder_should_normalize_text_usage_and_completion() {
));
}
#[test]
fn decoder_should_emit_each_complete_text_event_without_waiting_for_the_stream_end() {
let mut decoder = GrokCanonicalDecoder::new("grok-4.5");
let created = decoder
.push(
concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_live\",\"model\":\"grok-4.5\"}}\n\n",
)
.as_bytes(),
)
.expect("created event");
let delta = decoder
.push(
concat!(
"event: response.output_text.delta\n",
"data: {\"type\":\"response.output_text.delta\",\"output_index\":0,\"content_index\":0,\"delta\":\"live\"}\n\n",
)
.as_bytes(),
)
.expect("text delta");
assert!(
created
.iter()
.filter_map(ProviderEvent::wire_event)
.any(|event| event.event_type() == Some("response.created"))
);
assert!(delta.iter().any(|event| {
event.wire_event().is_some_and(|wire| {
wire.event_type() == Some("response.output_text.delta")
&& wire.data().get("delta") == Some(&serde_json::json!("live"))
})
}));
}
#[test]
fn decoder_should_preserve_image_tool_tokens_in_canonical_usage() {
let body = concat!(
@@ -416,6 +452,72 @@ fn decoder_should_restore_custom_apply_patch_arguments_as_raw_input() {
}));
}
#[test]
fn decoder_should_buffer_only_custom_arguments_until_the_done_event() {
let request = tool_request(serde_json::json!({
"model": "client",
"input": "render",
"tools": [{"type": "custom", "name": "render"}]
}));
let mut decoder = GrokCanonicalDecoder::for_request("grok-4.5", &request);
decoder
.push(
concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_render\",\"model\":\"grok-4.5\"}}\n\n",
)
.as_bytes(),
)
.expect("custom response created");
let added = decoder
.push(
concat!(
"event: response.output_item.added\n",
"data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"id\":\"item_render\",\"type\":\"function_call\",\"call_id\":\"call_render\",\"name\":\"render\"}}\n\n",
)
.as_bytes(),
)
.expect("custom item added");
let first_delta = decoder
.push(
concat!(
"event: response.function_call_arguments.delta\n",
"data: {\"type\":\"response.function_call_arguments.delta\",\"item_id\":\"item_render\",\"call_id\":\"call_render\",\"output_index\":0,\"delta\":\"{\\\"input\\\":\\\"hel\"}\n\n",
)
.as_bytes(),
)
.expect("first custom argument delta");
let second_delta = decoder
.push(
concat!(
"event: response.function_call_arguments.delta\n",
"data: {\"type\":\"response.function_call_arguments.delta\",\"item_id\":\"item_render\",\"call_id\":\"call_render\",\"output_index\":0,\"delta\":\"lo\\\"}\"}\n\n",
)
.as_bytes(),
)
.expect("second custom argument delta");
let done = decoder
.push(
concat!(
"event: response.function_call_arguments.done\n",
"data: {\"type\":\"response.function_call_arguments.done\",\"item_id\":\"item_render\",\"call_id\":\"call_render\",\"output_index\":0,\"arguments\":\"{\\\"input\\\":\\\"hello\\\"}\"}\n\n",
)
.as_bytes(),
)
.expect("custom argument done");
assert!(wire_events(&added).iter().any(|event| {
event.event_type() == Some("response.output_item.added")
&& event.data().pointer("/item/type") == Some(&serde_json::json!("custom_tool_call"))
}));
assert!(first_delta.is_empty());
assert!(second_delta.is_empty());
assert!(wire_events(&done).iter().any(|event| {
event.event_type() == Some("response.custom_tool_call_input.done")
&& event.data().get("input") == Some(&serde_json::json!("hello"))
}));
}
#[test]
fn decoder_should_coalesce_reasoning_item_part_and_summary_index() {
let body = concat!(
@@ -318,27 +318,28 @@ async fn inference_transport_should_classify_reasoning_decode_rejections() {
}
#[tokio::test]
async fn inference_transport_should_reject_success_with_non_sse_content_type() {
async fn inference_transport_should_not_reject_sse_bytes_only_because_content_type_is_wrong() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/responses"))
.respond_with(ResponseTemplate::new(200).set_body_raw("{}", "application/json"))
.respond_with(ResponseTemplate::new(200).set_body_raw(
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\"}}\n\n",
"application/json",
))
.expect(1)
.mount(&server)
.await;
let origin = Url::parse(&server.uri()).expect("wiremock origin");
let error = inference_transport(&origin)
let chunks = inference_transport(&origin)
.execute(inference_request(&origin))
.await
.expect_err("successful non-SSE response must fail closed");
.expect("body syntax, rather than response metadata, determines the stream")
.into_body()
.collect::<Vec<_>>()
.await;
assert_eq!(
(error.kind(), error.send_state()),
(
GrokInferenceTransportErrorKind::Protocol,
UpstreamSendState::Sent,
)
);
assert_eq!(chunks.len(), 1);
assert!(chunks[0].is_ok());
}
#[tokio::test]
@@ -273,6 +273,104 @@ fn explicit_session_should_merge_client_tools_into_the_native_cache_route() {
assert_eq!(body.pointer("/tool_choice"), Some(&json!("auto")));
}
#[test]
fn function_parameters_should_remove_only_nullable_object_roots() {
let request = raw_request(json!({
"model": "client",
"input": "hello",
"tools": [{
"type": "function",
"name": "automation_update",
"parameters": {
"anyOf": [
{
"type": "object",
"properties": {
"title": {
"anyOf": [{"type": "string"}, {"type": "null"}]
}
},
"required": ["title"],
"additionalProperties": false
},
{"type": "null"}
]
}
}]
}));
let encoded = GrokResponsesRequest::encode(&request, "grok-4.5", &client_key())
.expect("nullable object root");
let body = Value::Object(encoded.body().clone());
assert_eq!(
body.pointer("/tools/0/parameters/type"),
Some(&json!("object"))
);
assert_eq!(body.pointer("/tools/0/parameters/anyOf"), None);
assert_eq!(
body.pointer("/tools/0/parameters/properties/title/anyOf"),
Some(&json!([{"type": "string"}, {"type": "null"}]))
);
}
#[test]
fn function_parameters_should_keep_local_object_refs_after_removing_null() {
let request = raw_request(json!({
"model": "client",
"input": "hello",
"tools": [{
"type": "function",
"name": "lookup",
"parameters": {
"$defs": {
"Args": {
"type": "object",
"properties": {"query": {"type": "string"}}
}
},
"anyOf": [{"$ref": "#/$defs/Args"}, {"type": "null"}]
}
}]
}));
let encoded = GrokResponsesRequest::encode(&request, "grok-4.5", &client_key())
.expect("nullable local object ref");
let body = Value::Object(encoded.body().clone());
assert_eq!(
body.pointer("/tools/0/parameters/type"),
Some(&json!("object"))
);
assert_eq!(
body.pointer("/tools/0/parameters/anyOf"),
Some(&json!([{"$ref": "#/$defs/Args"}]))
);
}
#[test]
fn function_parameters_should_report_the_invalid_nullable_root_field() {
let request = raw_request(json!({
"model": "client",
"input": "hello",
"tools": [{
"type": "function",
"name": "invalid",
"parameters": {
"anyOf": [{"type": "object"}, {"type": "string"}, {"type": "null"}]
}
}]
}));
assert_eq!(
GrokResponsesRequest::encode(&request, "grok-4.5", &client_key())
.expect_err("nullable non-object root"),
GrokRequestEncodeError::InvalidRequestField {
field: "tools[].parameters"
}
);
}
#[test]
fn explicit_session_should_enable_cache_after_codex_additional_tools_normalization() {
let request = raw_request(json!({
@@ -721,7 +819,7 @@ fn compaction_history_should_reject_missing_empty_or_non_string_plaintext() {
assert_eq!(
GrokResponsesRequest::encode(&request, "grok-4.5", &client_key())
.expect_err("invalid compaction plaintext"),
GrokRequestEncodeError::InvalidRequestNormalization
GrokRequestEncodeError::InvalidRequestField { field: "input" }
);
}
}
@@ -765,27 +863,42 @@ fn tool_search_history_should_load_returned_tools_at_the_original_turn() {
#[test]
fn unsupported_or_ambiguous_tool_contracts_should_fail_before_upstream_io() {
for body in [
json!({"model": "client", "input": "x", "tools": [{"type": "future_tool"}]}),
json!({"model": "client", "input": "x", "tools": [
{"type": "tool_search", "execution": "client"},
{"type": "tool_search", "execution": "server"}
]}),
json!({"model": "client", "input": "x", "tools": [
{"type": "shell"}, {"type": "local_shell"}
]}),
json!({"model": "client", "input": [{
"type": "apply_patch_call",
"call_id": "patch_1",
"operation": {"type": "update_file", "path": "a.txt"}
}]}),
json!({"model": "client", "input": [{"type": "compaction_trigger"}]}),
for (body, field) in [
(
json!({"model": "client", "input": "x", "tools": [{"type": "future_tool"}]}),
"tools",
),
(
json!({"model": "client", "input": "x", "tools": [
{"type": "tool_search", "execution": "client"},
{"type": "tool_search", "execution": "server"}
]}),
"tools",
),
(
json!({"model": "client", "input": "x", "tools": [
{"type": "shell"}, {"type": "local_shell"}
]}),
"tools",
),
(
json!({"model": "client", "input": [{
"type": "apply_patch_call",
"call_id": "patch_1",
"operation": {"type": "update_file", "path": "a.txt"}
}]}),
"input",
),
(
json!({"model": "client", "input": [{"type": "compaction_trigger"}]}),
"input",
),
] {
let request = raw_request(body);
assert_eq!(
GrokResponsesRequest::encode(&request, "grok-4.5", &client_key())
.expect_err("invalid tool contract"),
GrokRequestEncodeError::InvalidRequestNormalization
GrokRequestEncodeError::InvalidRequestField { field }
);
}
}
@@ -33,8 +33,8 @@ const overviewItems = computed(() => [
icon: Gauge,
},
{
label: '需处理',
value: formatCount(props.summary.attention),
label: '不可用账号',
value: formatCount(props.summary.unavailable),
caption: '过期 / 禁用 / 封禁',
tone: 'danger',
icon: AlertTriangle,
@@ -136,7 +136,6 @@ const requestBars = computed(() => {
}
})
})
const hasDailyRequests = computed(() => requestBars.value.some(bar => bar.requestCount > 0))
</script>
<template>
@@ -198,7 +197,6 @@ const hasDailyRequests = computed(() => requestBars.value.some(bar => bar.reques
</strong>
</div>
<div
v-if="hasDailyRequests"
class="flex items-stretch gap-px"
:class="[trackShapeClass, !isCompact ? 'mt-2' : 'mt-auto']"
role="img"
@@ -217,11 +215,6 @@ const hasDailyRequests = computed(() => requestBars.value.some(bar => bar.reques
/>
</span>
</div>
<span
v-else
class="block font-mono tabular-nums text-(--cp-text-muted)"
:class="[valueClass, !isCompact ? 'mt-2' : 'mt-auto']"
>—</span>
</template>
<div v-else :class="headerClass">
@@ -273,7 +273,10 @@ export function useAccountConnectionTest(options: AccountConnectionTestOptions)
if (event.type === 'error') {
applyAccountStatus(event.accountStatus)
connectionTestError.value = event.error || '测试连接失败'
appendConnectionTestLog(connectionTestError.value, 'danger')
const detail = event.providerErrorCode || event.providerErrorType
? { code: event.providerErrorCode, type: event.providerErrorType }
: undefined
appendConnectionTestLog(connectionTestError.value, 'danger', detail)
finishConnectionTest('error')
clearConnectionTestRun()
}
@@ -24,7 +24,6 @@ import { useAccountOnboarding } from './useAccountOnboarding'
type AccountRow = Awaited<ReturnType<typeof getAccounts>>['items'][number]
const attentionAccountStatuses = new Set(['expired', 'disabled', 'banned'])
const accountStatusSortRank: Record<string, number> = {
active: 0,
quota_exhausted: 1,
@@ -250,15 +249,19 @@ export function useAccountMutations(options: {
+ Number(status === 'quota_exhausted')
- Number(current.status === 'quota_exhausted'),
),
attention: Math.max(
unavailable: Math.max(
0,
options.accountSummary.value.attention
+ Number(attentionAccountStatuses.has(status))
- Number(attentionAccountStatuses.has(current.status)),
options.accountSummary.value.unavailable
+ Number(isUnavailableAccountStatus(status))
- Number(isUnavailableAccountStatus(current.status)),
),
}
}
function isUnavailableAccountStatus(status: string) {
return status === 'expired' || status === 'disabled' || status === 'banned'
}
function sortAccountsByStatus(rows: AccountRow[]) {
if (options.sort.value?.key !== 'status')
return rows
@@ -17,7 +17,7 @@ export function useAccountsQuery() {
total: 0,
active: 0,
quotaExhausted: 0,
attention: 0,
unavailable: 0,
})
const query = usePagedQuery({
@@ -39,8 +39,8 @@ const emptyCards: DashboardSummary['cards'] = {
credentials: {
total: '0',
totalValue: 0,
enabled: '0',
enabledValue: 0,
available: '0',
availableValue: 0,
unavailable: '0',
unavailableValue: 0,
},
@@ -147,9 +147,9 @@ function metricCards(
tone: 'normal',
details: [
{
label: '启用',
value: credentials.enabled,
tone: credentials.enabledValue > 0 ? 'success' : 'normal',
label: '可用',
value: credentials.available,
tone: credentials.availableValue > 0 ? 'success' : 'normal',
},
{
label: '不可用',
@@ -57,7 +57,7 @@ function latencyDisplay(value: unknown) {
<template v-if="record">
<section class="rounded-(--cp-card-radius) bg-(--cp-bg-subtle) px-4 py-3.5">
<div class="flex flex-wrap items-center gap-3">
<UsageStatusCodeBadge :status-code="record.statusCode" />
<UsageStatusCodeBadge :status-code="record.upstreamStatusCode" />
<span class="text-[12px] font-bold text-(--cp-text-secondary)">
客户端 {{ display(record.clientStatusCode) }}
</span>
@@ -88,9 +88,9 @@ function showDetail(record: Awaited<ReturnType<typeof getOpsErrors>>['items'][nu
@page-change="handlePageChange"
@page-size-change="handlePageSizeChange"
>
<template #statusCode="{ row }">
<template #upstreamStatusCode="{ row }">
<UsageStatusCodeBadge
:status-code="typeof row.statusCode === 'number' ? row.statusCode : null"
:status-code="typeof row.upstreamStatusCode === 'number' ? row.upstreamStatusCode : null"
/>
</template>
<template #failureClass="{ row }">
+1 -1
View File
@@ -108,7 +108,7 @@ export const opsErrorColumns = [
cellClass:
'whitespace-nowrap font-mono text-[12px] font-[650] tabular-nums text-(--cp-text-secondary)',
},
{ key: 'statusCode', label: '状态码', width: '96px', align: 'center' as const },
{ key: 'upstreamStatusCode', label: '上游状态', width: '96px', align: 'center' as const },
{
key: 'failureClass',
label: '失败分类',