diff --git a/backend/crates/providers/xai/src/credential/selector.rs b/backend/crates/providers/xai/src/credential/selector.rs index b6ba828f9..d255229aa 100644 --- a/backend/crates/providers/xai/src/credential/selector.rs +++ b/backend/crates/providers/xai/src/credential/selector.rs @@ -4,7 +4,7 @@ use std::fmt; use std::sync::Arc; use std::time::{Duration, SystemTime}; -use chrono::Utc; +use chrono::{DateTime, Utc}; use gateway_core::engine::credential::{ AccountAvailability, AccountAvailabilityPolicy, AccountCandidate, AccountFeedbackStats, AccountSelectionContext, AccountSelector, ProviderAccount, ProviderAccountId, @@ -17,7 +17,7 @@ use gateway_core::provider_ports::{ use gateway_core::routing::{ProviderKind, UpstreamModelId}; use super::catalog::{GrokCatalogScope, GrokCredentialCatalogCache, GrokCredentialQuotaService}; -use super::repository::GrokCredentialRepository; +use super::repository::{GrokCredentialRepository, GrokCredentialRepositoryError}; use super::types::{GrokCredentialAvailability, UpdateGrokCredentialState}; use crate::{ GrokCredentialFailure, GrokCredentialFeedbackFuture, GrokSessionBinding, GrokSessionSelection, @@ -29,7 +29,6 @@ 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); -const PAYMENT_REQUIRED_COOLDOWN: Duration = Duration::from_secs(60); const MODEL_ACCESS_DENIED_COOLDOWN: Duration = Duration::from_secs(5 * 60); const MODEL_QUOTA_COOLDOWN: Duration = Duration::from_secs(24 * 60 * 60); const MAX_MODEL_COOLDOWN: Duration = Duration::from_secs(24 * 60 * 60); @@ -369,6 +368,66 @@ impl GrokAccountSessionSelector { )) .await; } + + async fn persist_account_failure( + &self, + session: &SelectedGrokSession, + availability: GrokCredentialAvailability, + reason: &'static str, + cooldown_until: Option>, + observed_at: chrono::DateTime, + ) { + let update = |expected_revision| UpdateGrokCredentialState { + account_id: session.account_id().clone(), + expected_revision, + availability, + availability_reason: Some(reason.to_owned()), + cooldown_until, + observed_at, + }; + match self + .repository + .update_state(&update(session.credential_revision())) + .await + { + Ok(()) => return, + Err( + GrokCredentialRepositoryError::Conflict + | GrokCredentialRepositoryError::StaleCredentialRevision, + ) => {} + Err(error) => { + tracing::warn!( + account_id = %session.account_id(), + error = %error, + "xAI account failure state write failed" + ); + return; + } + } + let current = match self.repository.load_current(session.account_id()).await { + Ok(current) if current.account.enabled() => current, + Ok(_) => return, + Err(error) => { + tracing::warn!( + account_id = %session.account_id(), + error = %error, + "xAI account failure state reload failed" + ); + return; + } + }; + if let Err(error) = self + .repository + .update_state(&update(current.account.revision())) + .await + { + tracing::warn!( + account_id = %session.account_id(), + error = %error, + "xAI account failure state retry failed" + ); + } + } } impl fmt::Debug for GrokAccountSessionSelector { @@ -415,15 +474,11 @@ impl GrokSessionSelector for GrokAccountSessionSelector { "upstream_free_quota_exhausted", None, ), - GrokCredentialFailure::PaymentRequired { retry_after } => { - let retry_after = bounded_cooldown( - retry_after, - PAYMENT_REQUIRED_COOLDOWN, - MAX_RATE_LIMIT_COOLDOWN, - ); - self.record_runtime_cooldown(session, retry_after).await; - return; - } + GrokCredentialFailure::PaymentRequired { .. } => ( + GrokCredentialAvailability::QuotaExhausted, + "upstream_payment_quota_exhausted", + None, + ), GrokCredentialFailure::ModelQuotaExhausted { upstream_model, retry_after, @@ -452,26 +507,31 @@ impl GrokSessionSelector for GrokAccountSessionSelector { .unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN) .min(MAX_RATE_LIMIT_COOLDOWN); self.record_runtime_cooldown(session, retry_after).await; - return; + ( + GrokCredentialAvailability::Cooldown, + "upstream_rate_limited", + persistent_cooldown_until(observed_at, retry_after), + ) } GrokCredentialFailure::StreamInterrupted => { self.record_runtime_cooldown(session, STREAM_INTERRUPTION_COOLDOWN) .await; - return; + ( + GrokCredentialAvailability::Cooldown, + "upstream_stream_interrupted", + persistent_cooldown_until(observed_at, STREAM_INTERRUPTION_COOLDOWN), + ) } }; let (availability, reason, cooldown_until) = persistent; - let _ = self - .repository - .update_state(&UpdateGrokCredentialState { - account_id: session.account_id().clone(), - expected_revision: session.credential_revision(), - availability, - availability_reason: Some(reason.to_owned()), - cooldown_until, - observed_at, - }) - .await; + self.persist_account_failure( + session, + availability, + reason, + cooldown_until, + observed_at, + ) + .await; }) } @@ -493,7 +553,6 @@ impl GrokSessionSelector for GrokAccountSessionSelector { account.availability(), AccountAvailability::Unknown | AccountAvailability::Cooldown - | AccountAvailability::QuotaExhausted | AccountAvailability::Expired ) { @@ -540,6 +599,15 @@ fn bounded_cooldown( .min(maximum) } +fn persistent_cooldown_until( + observed_at: DateTime, + duration: Duration, +) -> Option> { + chrono::Duration::from_std(duration) + .ok() + .map(|duration| observed_at + duration) +} + fn insert_runtime_cooldown( cooldowns: &mut std::collections::BTreeMap, account_id: &ProviderAccountId, diff --git a/backend/crates/providers/xai/src/lib.rs b/backend/crates/providers/xai/src/lib.rs index 110da7cc9..0b947f5c7 100644 --- a/backend/crates/providers/xai/src/lib.rs +++ b/backend/crates/providers/xai/src/lib.rs @@ -4,6 +4,7 @@ mod admin; pub mod config; pub mod credential; mod provider; +mod reasoning_replay; pub mod transport; use std::path::Path; diff --git a/backend/crates/providers/xai/src/provider.rs b/backend/crates/providers/xai/src/provider.rs index 76eb66c6a..5b757bef4 100644 --- a/backend/crates/providers/xai/src/provider.rs +++ b/backend/crates/providers/xai/src/provider.rs @@ -5,7 +5,6 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant, SystemTime}; use async_trait::async_trait; -use base64::Engine as _; use futures::{StreamExt, future::BoxFuture}; use gateway_core::engine::continuation::ContinuationBinding; use gateway_core::engine::credential::{ @@ -48,9 +47,13 @@ use crate::credential::{ GrokCredentialRecovery, GrokCredentialRecoveryOutcome, GrokCredentialRefreshOutcome, GrokCredentialRefreshService, GrokQuotaError, }; +use crate::reasoning_replay::{ + GrokReasoningReplay, GrokReasoningReplayCapture, GrokReasoningReplayKey, + valid_reasoning_ciphertext, +}; use crate::transport::canonical::GrokCanonicalDecoder; use crate::transport::config::XAI_PROVIDER_NAME; -use crate::transport::headers::{GrokClientIdentity, GrokHeader, build_grok_headers}; +use crate::transport::headers::{GrokClientIdentity, build_grok_headers}; use crate::transport::profile::{GROK_CLI_RELEASE_POLL_INTERVAL, GrokCliReleaseService}; use crate::transport::{ GROK_RESPONSES_URL, GrokCompactionDecodeError, GrokCompactionRequest, @@ -66,8 +69,6 @@ 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; -const MIN_REASONING_CIPHERTEXT_ENTROPY: f64 = 0.85; const REASONING_DECODE_FAILED_CODE: &str = "reasoning_decode_failed"; const RESPONSE_NOT_FOUND_CODE: &str = "not_found"; @@ -82,6 +83,7 @@ pub struct GrokBuildProvider { credential_recovery: Arc, account_feedback: Arc, client_identity: GrokClientIdentity, + reasoning_replay: GrokReasoningReplay, wire_profile: XaiWireProfileState, responses_url: Url, } @@ -105,6 +107,7 @@ impl GrokBuildProvider { credential_recovery, account_feedback, client_identity: GrokClientIdentity::new(), + reasoning_replay: GrokReasoningReplay::new(), wire_profile, responses_url, }) @@ -217,6 +220,38 @@ impl GrokBuildProvider { selected.account_id(), request_input.as_slice(), )?; + let inherited_replay_session_id = context + .continuation() + .is_none() + .then(|| { + previous_session + .as_ref() + .and_then(|state| state.session_id.as_deref()) + }) + .flatten(); + let reasoning_replay_key = upstream_request + .reasoning_replay_session_id() + .or(inherited_replay_session_id) + .and_then(|session_id| { + self.reasoning_replay.key( + candidate.upstream_model().as_str(), + session_id, + selected.account_id().as_str(), + ) + }); + if !upstream_request.has_previous_response_id() + && let (Some(key), Some(input)) = ( + reasoning_replay_key.as_ref(), + upstream_request.replay_input_items(), + ) + && let Some(input) = self.reasoning_replay.apply(key, &input) + { + upstream_request + .set_replay_input(input) + .map_err(map_request_error)?; + } + let reasoning_replay_capture = + reasoning_replay_key.map(|key| self.reasoning_replay.capture(key)); let session_capture = (!matches!( context.continuation(), Some(ContinuationBinding::Pinned(_) | ContinuationBinding::External(_)) @@ -244,6 +279,7 @@ impl GrokBuildProvider { context, session: Arc::clone(&selected), session_capture, + reasoning_replay_capture, }, ); let stream = ProviderStream::new(metadata, events, selected); @@ -274,7 +310,7 @@ impl GrokBuildProvider { .map(|previous| ProviderAccountId::new(previous.account_id.clone())) .transpose() .map_err(|_| protocol_not_sent())?; - let upstream_session_id = previous_session + let inherited_session_id = previous_session .as_ref() .and_then(|previous| previous.session_id.clone()); let upstream_request = GrokCompactionRequest::encode( @@ -283,6 +319,12 @@ impl GrokBuildProvider { context.client_api_key_ref(), ) .map_err(map_request_error)?; + let explicit_replay_session_id = upstream_request + .reasoning_replay_session_id() + .map(str::to_owned); + let upstream_session_id = inherited_session_id + .clone() + .or_else(|| explicit_replay_session_id.clone()); let selected = Arc::new( select_grok_session( self.selector.as_ref(), @@ -293,6 +335,16 @@ impl GrokBuildProvider { ) .await?, ); + let reasoning_replay_key = explicit_replay_session_id + .as_deref() + .or(inherited_session_id.as_deref()) + .and_then(|session_id| { + self.reasoning_replay.key( + candidate.upstream_model().as_str(), + session_id, + selected.account_id().as_str(), + ) + }); let allows_account_state_mutation = selected.allows_account_state_mutation(); let metadata = provider_call_metadata(candidate, &selected)?; let events = cold_compaction_http_sse_stream( @@ -308,6 +360,8 @@ impl GrokBuildProvider { upstream_session_id, context, session: Arc::clone(&selected), + reasoning_replay: self.reasoning_replay.clone(), + reasoning_replay_key, }, ); let stream = ProviderStream::new(metadata, events, selected); @@ -715,46 +769,6 @@ fn portable_output_item(mut item: Value, strip_opaque: bool) -> Option { Some(item) } -fn valid_reasoning_ciphertext(value: &str) -> bool { - if value.is_empty() - || value != value.trim() - || value.len() > XAI_SESSION_STATE_MAX_BYTES - || value.starts_with("gAAAA") - || value.contains('=') - { - return false; - } - let Ok(decoded) = base64::engine::general_purpose::STANDARD_NO_PAD.decode(value) else { - return false; - }; - decoded.len() >= MIN_REASONING_CIPHERTEXT_BYTES - && byte_entropy_ratio(&decoded) >= MIN_REASONING_CIPHERTEXT_ENTROPY -} - -fn byte_entropy_ratio(value: &[u8]) -> f64 { - if value.is_empty() { - return 0.0; - } - let mut counts = [0_u32; 256]; - for byte in value { - counts[usize::from(*byte)] += 1; - } - let size = value.len() as f64; - let entropy = counts - .into_iter() - .filter(|count| *count > 0) - .map(|count| { - let probability = f64::from(count) / size; - -probability * probability.log2() - }) - .sum::(); - let symbols = value.len().min(256); - if symbols <= 1 { - return 0.0; - } - entropy / (symbols as f64).log2() -} - fn has_readable_reasoning(item: &Map) -> bool { ["summary", "content"].into_iter().any(|field| { item.get(field) @@ -898,6 +912,7 @@ struct GrokStreamAttempt { context: AttemptContext, session: Arc, session_capture: Option, + reasoning_replay_capture: Option, } struct GrokCompactionStreamAttempt { @@ -910,6 +925,8 @@ struct GrokCompactionStreamAttempt { upstream_session_id: Option, context: AttemptContext, session: Arc, + reasoning_replay: GrokReasoningReplay, + reasoning_replay_key: Option, } struct AcceptedGrokInference { @@ -1044,9 +1061,11 @@ fn cold_compaction_http_sse_stream( upstream_session_id, context, session, + reasoning_replay, + reasoning_replay_key, } = attempt; Box::pin(async_stream::try_stream! { - let mut headers = build_grok_headers( + let headers = build_grok_headers( &wire_profile, &session, &client_identity, @@ -1055,12 +1074,6 @@ fn cold_compaction_http_sse_stream( None, &upstream_model, ); - if let Some(upstream_session_id) = upstream_session_id.as_deref() { - headers.push(GrokHeader::sensitive( - "x-grok-session-id", - crate::credential::SecretValue::new(upstream_session_id.to_owned()), - )); - } let body = request.to_json_bytes().map_err(map_request_error)?; let inference_request = GrokInferenceRequest::new( responses_url, @@ -1147,6 +1160,9 @@ fn cold_compaction_http_sse_stream( if upstream_completed && session.allows_account_state_mutation() { selector.record_success(&session).await; } + if let Some(key) = reasoning_replay_key.as_ref() { + reasoning_replay.clear(key); + } yield ProviderEvent::canonical_with_wire(vec![GatewayEvent::Started(started)], created); yield ProviderEvent::wire(output_done); let mut terminal_facts = facts.accounting; @@ -1245,6 +1261,7 @@ fn cold_http_sse_stream( context, session, mut session_capture, + mut reasoning_replay_capture, } = attempt; Box::pin(async_stream::try_stream! { if context.cancellation().is_cancelled() { @@ -1364,6 +1381,9 @@ fn cold_http_sse_stream( .flat_map(ProviderEvent::canonical_facts) .any(|event| matches!(event, GatewayEvent::Completed(_))); attach_xai_session_update(&mut events, &mut session_capture)?; + if let Some(capture) = reasoning_replay_capture.as_mut() { + capture.observe(&events); + } if completed && session.allows_account_state_mutation() { selector.record_success(&session).await; } @@ -1389,6 +1409,9 @@ fn cold_http_sse_stream( .flat_map(ProviderEvent::canonical_facts) .any(|event| matches!(event, GatewayEvent::Completed(_))); attach_xai_session_update(&mut final_events, &mut session_capture)?; + if let Some(capture) = reasoning_replay_capture.as_mut() { + capture.observe(&final_events); + } if completed && session.allows_account_state_mutation() { selector.record_success(&session).await; } @@ -1871,7 +1894,9 @@ const WORKER_LEASE_TTL: Duration = Duration::from_secs(15 * 60); const WORKER_LEASE_RENEWAL: Duration = Duration::from_secs(5 * 60); const OAUTH_REFRESH_INTERVAL: Duration = Duration::from_secs(30); const QUOTA_CATALOG_INTERVAL: Duration = Duration::from_secs(5 * 60); -const QUOTA_PERIODIC_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(30 * 60); +// grok2api 对 Free/未知额度耗尽采用 24 小时恢复探测；当前公共账号状态没有 +// 独立的 paid period-end 字段，因此 Build 账号统一使用同一保守下限。 +const QUOTA_PERIODIC_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60); const CLI_RELEASE_WORKER_OWNER: &str = "xai-cli-release"; pub(crate) fn worker_contributions( diff --git a/backend/crates/providers/xai/src/reasoning_replay.rs b/backend/crates/providers/xai/src/reasoning_replay.rs new file mode 100644 index 000000000..8f27dfceb --- /dev/null +++ b/backend/crates/providers/xai/src/reasoning_replay.rs @@ -0,0 +1,687 @@ +//! Grok Build 的账号隔离 reasoning replay。 + +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::{Arc, Mutex, MutexGuard}; +use std::time::{Duration, Instant}; + +use base64::Engine as _; +use gateway_core::event::ProviderEvent; +use serde_json::{Map, Value}; + +const REASONING_REPLAY_TTL: Duration = Duration::from_secs(60 * 60); +const REASONING_REPLAY_MAX_ENTRIES: usize = 1_024; +const REASONING_REPLAY_MAX_BYTES: usize = 8 * 1024 * 1024; +const REASONING_REPLAY_MAX_OUTPUT_ITEMS: usize = 4_096; +const MIN_REASONING_CIPHERTEXT_BYTES: usize = 50; +const MIN_REASONING_CIPHERTEXT_ENTROPY: f64 = 0.85; + +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)] +pub(crate) struct GrokReasoningReplayKey { + model: String, + session_id: String, + account_id: String, +} + +#[derive(Clone)] +pub(crate) struct GrokReasoningReplay { + state: Arc>, +} + +#[derive(Default)] +struct GrokReasoningReplayState { + entries: BTreeMap, +} + +struct GrokReasoningReplayEntry { + stored_at: Instant, + items: Vec, +} + +pub(crate) struct GrokReasoningReplayCapture { + replay: GrokReasoningReplay, + key: GrokReasoningReplayKey, + output_items: BTreeMap, + captured_bytes: usize, + truncated: bool, + finished: bool, +} + +impl Default for GrokReasoningReplay { + fn default() -> Self { + Self::new() + } +} + +impl GrokReasoningReplay { + pub(crate) fn new() -> Self { + Self { + state: Arc::new(Mutex::new(GrokReasoningReplayState::default())), + } + } + + pub(crate) fn key( + &self, + model: &str, + session_id: &str, + account_id: &str, + ) -> Option { + let model = model.trim(); + let session_id = session_id.trim(); + let account_id = account_id.trim(); + if model.is_empty() || session_id.is_empty() || account_id.is_empty() { + return None; + } + Some(GrokReasoningReplayKey { + model: model.to_owned(), + session_id: session_id.to_owned(), + account_id: account_id.to_owned(), + }) + } + + pub(crate) fn apply( + &self, + key: &GrokReasoningReplayKey, + input: &[Value], + ) -> Option> { + if input.is_empty() { + return None; + } + let cached = self.read(key)?; + let filtered = filter_replay_items_for_input(input, cached)?; + Some(insert_replay_items(input, &filtered)) + } + + pub(crate) fn capture(&self, key: GrokReasoningReplayKey) -> GrokReasoningReplayCapture { + GrokReasoningReplayCapture { + replay: self.clone(), + key, + output_items: BTreeMap::new(), + captured_bytes: 0, + truncated: false, + finished: false, + } + } + + pub(crate) fn clear(&self, key: &GrokReasoningReplayKey) { + self.lock_state().entries.remove(key); + } + + fn read(&self, key: &GrokReasoningReplayKey) -> Option> { + let now = Instant::now(); + let mut state = self.lock_state(); + prune_expired(&mut state, now); + state.entries.get(key).map(|entry| entry.items.clone()) + } + + fn store_output(&self, key: GrokReasoningReplayKey, output: Vec) { + let normalized = match normalize_replay_items(output) { + ReplayNormalization::Store(items) => items, + ReplayNormalization::Delete => { + self.clear(&key); + return; + } + ReplayNormalization::Oversized => return, + }; + let now = Instant::now(); + let mut state = self.lock_state(); + prune_expired(&mut state, now); + if !state.entries.contains_key(&key) + && state.entries.len() >= REASONING_REPLAY_MAX_ENTRIES + && let Some(oldest) = state + .entries + .iter() + .min_by_key(|(_, entry)| entry.stored_at) + .map(|(key, _)| key.clone()) + { + state.entries.remove(&oldest); + } + state.entries.insert( + key, + GrokReasoningReplayEntry { + stored_at: now, + items: normalized, + }, + ); + } + + fn lock_state(&self) -> MutexGuard<'_, GrokReasoningReplayState> { + match self.state.lock() { + Ok(state) => state, + Err(poisoned) => poisoned.into_inner(), + } + } +} + +impl GrokReasoningReplayCapture { + pub(crate) fn observe(&mut self, events: &[ProviderEvent]) { + if self.finished { + return; + } + let mut completed = false; + for event in events { + let Some(wire) = event.wire_event() else { + continue; + }; + let payload = wire.data(); + let event_type = wire + .event_type() + .or_else(|| payload.get("type").and_then(Value::as_str)); + if event_type == Some("response.output_item.done") + && let Some(item) = payload.get("item").cloned() + { + let index = payload + .get("output_index") + .and_then(Value::as_u64) + .and_then(|value| u32::try_from(value).ok()) + .unwrap_or_else(|| u32::try_from(self.output_items.len()).unwrap_or(u32::MAX)); + self.capture_item(index, item); + } + if event_type == Some("response.completed") { + capture_terminal_output(payload, self); + completed = true; + } + } + if completed { + self.finished = true; + if self.truncated { + self.output_items.clear(); + return; + } + let output = std::mem::take(&mut self.output_items) + .into_values() + .collect(); + self.replay.store_output(self.key.clone(), output); + } + } + + fn capture_item(&mut self, index: u32, item: Value) { + if self.truncated || self.output_items.contains_key(&index) { + return; + } + if self.output_items.len() >= REASONING_REPLAY_MAX_OUTPUT_ITEMS { + self.truncate(); + return; + } + let Some(captured_bytes) = serde_json::to_vec(&item) + .ok() + .and_then(|encoded| self.captured_bytes.checked_add(encoded.len())) + .filter(|captured_bytes| *captured_bytes <= REASONING_REPLAY_MAX_BYTES) + else { + self.truncate(); + return; + }; + self.captured_bytes = captured_bytes; + self.output_items.insert(index, item); + } + + fn truncate(&mut self) { + self.truncated = true; + self.output_items.clear(); + self.captured_bytes = 0; + } +} + +fn capture_terminal_output(payload: &Value, capture: &mut GrokReasoningReplayCapture) { + let Some(output) = payload + .get("response") + .and_then(Value::as_object) + .and_then(|response| response.get("output")) + .and_then(Value::as_array) + else { + return; + }; + if output.len() > REASONING_REPLAY_MAX_OUTPUT_ITEMS { + capture.truncate(); + return; + } + for (index, item) in output + .iter() + .take(REASONING_REPLAY_MAX_OUTPUT_ITEMS) + .enumerate() + { + let Some(index) = u32::try_from(index).ok() else { + break; + }; + capture.capture_item(index, item.clone()); + } +} + +fn prune_expired(state: &mut GrokReasoningReplayState, now: Instant) { + state + .entries + .retain(|_, entry| now.saturating_duration_since(entry.stored_at) < REASONING_REPLAY_TTL); +} + +enum ReplayNormalization { + Store(Vec), + Delete, + Oversized, +} + +fn normalize_replay_items(items: Vec) -> ReplayNormalization { + let mut normalized = Vec::new(); + let mut has_anchor = false; + for item in items { + let Some(item) = normalize_replay_item(item) else { + continue; + }; + if matches!( + item.get("type").and_then(Value::as_str), + Some("reasoning" | "function_call" | "custom_tool_call") + ) { + has_anchor = true; + } + normalized.push(item); + } + if !has_anchor || normalized.is_empty() { + return ReplayNormalization::Delete; + } + match serde_json::to_vec(&normalized) { + Ok(encoded) if encoded.len() <= REASONING_REPLAY_MAX_BYTES => { + ReplayNormalization::Store(normalized) + } + _ => ReplayNormalization::Oversized, + } +} + +fn normalize_replay_item(item: Value) -> Option { + let object = item.as_object()?; + match object.get("type").and_then(Value::as_str)?.trim() { + "reasoning" => normalize_reasoning_item(object), + "message" => normalize_assistant_message_item(object), + "function_call" => normalize_function_call_item(object), + "custom_tool_call" => normalize_custom_tool_call_item(object), + _ => None, + } +} + +fn normalize_reasoning_item(item: &Map) -> Option { + let encrypted = item.get("encrypted_content")?.as_str()?; + if !valid_reasoning_ciphertext(encrypted) { + return None; + } + Some(Value::Object(Map::from_iter([ + ("type".to_owned(), Value::String("reasoning".to_owned())), + ("summary".to_owned(), Value::Array(Vec::new())), + ("content".to_owned(), Value::Null), + ( + "encrypted_content".to_owned(), + Value::String(encrypted.to_owned()), + ), + ]))) +} + +fn normalize_assistant_message_item(item: &Map) -> Option { + if !item + .get("role") + .and_then(Value::as_str) + .is_some_and(|role| role.trim().eq_ignore_ascii_case("assistant")) + { + return None; + } + let parts = item + .get("content")? + .as_array()? + .iter() + .filter_map(normalize_assistant_part) + .collect::>(); + if parts.is_empty() { + return None; + } + Some(Value::Object(Map::from_iter([ + ("type".to_owned(), Value::String("message".to_owned())), + ("role".to_owned(), Value::String("assistant".to_owned())), + ("content".to_owned(), Value::Array(parts)), + ]))) +} + +fn normalize_assistant_part(part: &Value) -> Option { + let part = part.as_object()?; + match part.get("type").and_then(Value::as_str)?.trim() { + "output_text" => Some(Value::Object(Map::from_iter([ + ("type".to_owned(), Value::String("output_text".to_owned())), + ( + "text".to_owned(), + Value::String(part.get("text")?.as_str()?.to_owned()), + ), + ]))), + "refusal" => Some(Value::Object(Map::from_iter([ + ("type".to_owned(), Value::String("refusal".to_owned())), + ( + "refusal".to_owned(), + Value::String(part.get("refusal")?.as_str()?.to_owned()), + ), + ]))), + _ => None, + } +} + +fn normalize_function_call_item(item: &Map) -> Option { + let call_id = nonempty_string(item, "call_id")?; + let name = nonempty_string(item, "name")?; + let arguments = item.get("arguments")?.as_str()?.to_owned(); + Some(Value::Object(Map::from_iter([ + ("type".to_owned(), Value::String("function_call".to_owned())), + ("call_id".to_owned(), Value::String(call_id)), + ("name".to_owned(), Value::String(name)), + ("arguments".to_owned(), Value::String(arguments)), + ]))) +} + +fn normalize_custom_tool_call_item(item: &Map) -> Option { + let call_id = nonempty_string(item, "call_id")?; + let name = nonempty_string(item, "name")?; + let input = item.get("input")?.clone(); + let status = item + .get("status") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or("completed") + .to_owned(); + Some(Value::Object(Map::from_iter([ + ( + "type".to_owned(), + Value::String("custom_tool_call".to_owned()), + ), + ("status".to_owned(), Value::String(status)), + ("call_id".to_owned(), Value::String(call_id)), + ("name".to_owned(), Value::String(name)), + ("input".to_owned(), input), + ]))) +} + +fn nonempty_string(item: &Map, field: &str) -> Option { + item.get(field) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn filter_replay_items_for_input(input: &[Value], cached: Vec) -> Option> { + let last_assistant = input.iter().rev().find_map(assistant_message); + let cached_assistant = cached.iter().find_map(assistant_message); + let assistant_matches = match (last_assistant, cached_assistant) { + (Some(current), Some(previous)) => { + if !assistant_content_equal(current, previous) { + return None; + } + true + } + _ => false, + }; + + let mut existing_calls = BTreeSet::new(); + let mut existing_outputs = BTreeMap::new(); + let mut existing_encrypted = BTreeSet::new(); + for item in input { + let Some(item) = item.as_object() else { + continue; + }; + match item.get("type").and_then(Value::as_str).map(str::trim) { + Some("reasoning") => { + if let Some(encrypted) = item + .get("encrypted_content") + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + { + existing_encrypted.insert(encrypted.to_owned()); + } + } + Some("function_call_output" | "custom_tool_call_output") => { + if let Some(call_id) = item.get("call_id").and_then(Value::as_str) { + for candidate in comparable_call_ids(call_id) { + existing_outputs.insert(candidate, call_id.trim().to_owned()); + } + } + } + Some(kind @ ("function_call" | "custom_tool_call")) => { + if let Some(call_id) = item.get("call_id").and_then(Value::as_str) { + existing_calls.extend(replay_call_keys(kind, call_id)); + } + } + _ => {} + } + } + + let mut filtered = Vec::new(); + for mut item in cached { + let kind = item + .get("type") + .and_then(Value::as_str) + .map(str::trim) + .unwrap_or_default() + .to_owned(); + match kind.as_str() { + "reasoning" => { + if item + .get("encrypted_content") + .and_then(Value::as_str) + .is_some_and(|encrypted| existing_encrypted.contains(encrypted)) + { + continue; + } + } + "message" if assistant_matches => continue, + "message" => {} + "function_call" | "custom_tool_call" => { + let Some(call_id) = item.get("call_id").and_then(Value::as_str) else { + continue; + }; + let keys = replay_call_keys(&kind, call_id); + if keys.is_empty() || keys.iter().any(|key| existing_calls.contains(key)) { + continue; + } + let output_call_id = comparable_call_ids(call_id) + .into_iter() + .find_map(|candidate| existing_outputs.get(&candidate).cloned()); + let Some(output_call_id) = output_call_id else { + continue; + }; + existing_calls.extend(keys); + if output_call_id != call_id + && let Some(item) = item.as_object_mut() + { + item.insert("call_id".to_owned(), Value::String(output_call_id)); + } + } + _ => continue, + } + filtered.push(item); + } + (!filtered.is_empty()).then_some(filtered) +} + +fn assistant_message(item: &Value) -> Option<&Map> { + let item = item.as_object()?; + let item_type = item + .get("type") + .and_then(Value::as_str) + .map(str::trim) + .unwrap_or_default(); + let role = item.get("role")?.as_str()?.trim(); + ((item_type.is_empty() || item_type == "message") && role.eq_ignore_ascii_case("assistant")) + .then_some(item) +} + +fn assistant_content_equal(left: &Map, right: &Map) -> bool { + let Some(left) = assistant_parts(left.get("content")) else { + return false; + }; + let Some(right) = assistant_parts(right.get("content")) else { + return false; + }; + left == right +} + +fn assistant_parts(content: Option<&Value>) -> Option> { + match content? { + Value::String(text) => Some(vec![("output_text".to_owned(), text.clone())]), + Value::Array(parts) => { + let mut result = Vec::new(); + for part in parts { + let part = part.as_object()?; + let kind = part.get("type")?.as_str()?.trim(); + let value = match kind { + "output_text" => part.get("text")?.as_str()?, + "refusal" => part.get("refusal")?.as_str()?, + _ => return None, + }; + result.push((kind.to_owned(), value.to_owned())); + } + (!result.is_empty()).then_some(result) + } + _ => None, + } +} + +fn comparable_call_ids(call_id: &str) -> Vec { + let call_id = call_id.trim(); + if call_id.is_empty() { + return Vec::new(); + } + if let Some(upstream) = call_id + .strip_prefix("toolu_") + .filter(|value| !value.is_empty()) + { + return vec![call_id.to_owned(), upstream.to_owned()]; + } + vec![call_id.to_owned(), format!("toolu_{call_id}")] +} + +fn replay_call_keys(item_type: &str, call_id: &str) -> Vec { + if !matches!(item_type, "function_call" | "custom_tool_call") { + return Vec::new(); + } + comparable_call_ids(call_id) + .into_iter() + .map(|call_id| format!("{item_type}\0{call_id}")) + .collect() +} + +fn insert_replay_items(input: &[Value], replay: &[Value]) -> Vec { + let insert_at = replay_insert_index(input, replay); + let mut result = Vec::new(); + for (index, item) in input.iter().enumerate() { + if index == insert_at { + result.extend(replay.iter().cloned()); + } + result.push(item.clone()); + } + if insert_at == input.len() { + result.extend(replay.iter().cloned()); + } + result +} + +fn replay_insert_index(input: &[Value], replay: &[Value]) -> usize { + let replay_call_ids = replay + .iter() + .filter_map(Value::as_object) + .filter(|item| { + matches!( + item.get("type").and_then(Value::as_str), + Some("function_call" | "custom_tool_call") + ) + }) + .filter_map(|item| item.get("call_id").and_then(Value::as_str)) + .flat_map(comparable_call_ids) + .collect::>(); + if !replay_call_ids.is_empty() { + for (index, item) in input.iter().enumerate() { + let Some(item) = item.as_object() else { + continue; + }; + if !matches!( + item.get("type").and_then(Value::as_str), + Some("function_call_output" | "custom_tool_call_output") + ) { + continue; + } + let call_ids = item + .get("call_id") + .and_then(Value::as_str) + .map(comparable_call_ids) + .unwrap_or_default(); + if call_ids.is_empty() + || call_ids + .iter() + .any(|call_id| replay_call_ids.contains(call_id)) + { + return index; + } + } + } + if let Some(index) = input + .iter() + .rposition(|item| assistant_message(item).is_some()) + { + return index; + } + input + .iter() + .position(should_insert_replay_before) + .unwrap_or(input.len()) +} + +fn should_insert_replay_before(item: &Value) -> bool { + let Some(item) = item.as_object() else { + return true; + }; + let item_type = item + .get("type") + .and_then(Value::as_str) + .map(str::trim) + .unwrap_or_default(); + let role = item + .get("role") + .and_then(Value::as_str) + .map(str::trim) + .unwrap_or_default() + .to_ascii_lowercase(); + if role.is_empty() || (!item_type.is_empty() && item_type != "message") { + return true; + } + !matches!(role.as_str(), "system" | "developer") +} + +pub(crate) fn valid_reasoning_ciphertext(value: &str) -> bool { + if value.is_empty() + || value != value.trim() + || value.len() > REASONING_REPLAY_MAX_BYTES + || value.starts_with("gAAAA") + || value.contains('=') + { + return false; + } + let Ok(decoded) = base64::engine::general_purpose::STANDARD_NO_PAD.decode(value) else { + return false; + }; + decoded.len() >= MIN_REASONING_CIPHERTEXT_BYTES + && byte_entropy_ratio(&decoded) >= MIN_REASONING_CIPHERTEXT_ENTROPY +} + +fn byte_entropy_ratio(value: &[u8]) -> f64 { + if value.is_empty() { + return 0.0; + } + let mut counts = [0_u32; 256]; + for byte in value { + counts[usize::from(*byte)] += 1; + } + let size = value.len() as f64; + let entropy = counts + .into_iter() + .filter(|count| *count > 0) + .map(|count| { + let probability = f64::from(count) / size; + -probability * probability.log2() + }) + .sum::(); + let symbols = value.len().min(256); + if symbols <= 1 { + return 0.0; + } + entropy / (symbols as f64).log2() +} diff --git a/backend/crates/providers/xai/src/transport/compaction.rs b/backend/crates/providers/xai/src/transport/compaction.rs index 85a03330a..79c338d9e 100644 --- a/backend/crates/providers/xai/src/transport/compaction.rs +++ b/backend/crates/providers/xai/src/transport/compaction.rs @@ -42,6 +42,7 @@ If the prior conversation contains a note about files at /tmp/compaction/segment pub struct GrokCompactionRequest { body: Map, affinity: Option, + reasoning_replay_session_id: Option, } impl GrokCompactionRequest { @@ -60,6 +61,8 @@ impl GrokCompactionRequest { upstream_model, client_api_key_ref, )?; + let reasoning_replay_session_id = + normalized.reasoning_replay_session_id().map(str::to_owned); let mut body = normalized.body().clone(); let mut input = normalized.input_items(); input.push(summary_prompt_item()); @@ -95,6 +98,7 @@ impl GrokCompactionRequest { Ok(Self { body, affinity: normalized.affinity().cloned(), + reasoning_replay_session_id, }) } @@ -110,6 +114,12 @@ impl GrokCompactionRequest { self.affinity.as_ref() } + /// 返回用于在成功 compact 后清除 reasoning replay 的显式会话身份。 + #[must_use] + pub(crate) fn reasoning_replay_session_id(&self) -> Option<&str> { + self.reasoning_replay_session_id.as_deref() + } + /// 序列化上游请求正文。 /// /// # Errors diff --git a/backend/crates/providers/xai/src/transport/headers.rs b/backend/crates/providers/xai/src/transport/headers.rs index 9732c41b1..6b738dde6 100644 --- a/backend/crates/providers/xai/src/transport/headers.rs +++ b/backend/crates/providers/xai/src/transport/headers.rs @@ -142,11 +142,11 @@ pub fn build_grok_headers( "x-grok-conv-id", SecretValue::new(upstream_session_id.to_owned()), )); + headers.push(GrokHeader::sensitive( + "x-grok-session-id", + SecretValue::new(upstream_session_id.to_owned()), + )); if let Some(turn_index) = turn_index { - headers.push(GrokHeader::sensitive( - "x-grok-session-id", - SecretValue::new(upstream_session_id.to_owned()), - )); headers.push(GrokHeader::public("x-grok-turn-idx", turn_index)); } } diff --git a/backend/crates/providers/xai/src/transport/request.rs b/backend/crates/providers/xai/src/transport/request.rs index 6f82ff4a4..e084a82b8 100644 --- a/backend/crates/providers/xai/src/transport/request.rs +++ b/backend/crates/providers/xai/src/transport/request.rs @@ -68,6 +68,7 @@ const GROK_INTERNAL_HISTORY_KEYS: &[&str] = pub struct GrokResponsesRequest { body: Map, session_id: Option, + reasoning_replay_session_id: Option, affinity: Option, response_transform: GrokResponseTransform, } @@ -85,6 +86,23 @@ impl GrokResponsesRequest { self.session_id.as_deref() } + /// 返回只有显式会话身份才启用的 reasoning replay scope。 + pub(crate) fn reasoning_replay_session_id(&self) -> Option<&str> { + self.reasoning_replay_session_id.as_deref() + } + + pub(crate) fn has_previous_response_id(&self) -> bool { + self.body + .get("previous_response_id") + .and_then(Value::as_str) + .map(str::trim) + .is_some_and(|value| !value.is_empty()) + } + + pub(crate) fn replay_input_items(&self) -> Option> { + self.body.get("input").and_then(Value::as_array).cloned() + } + /// 返回与会话一致、额外绑定模型的账号亲和键。 #[must_use] pub const fn affinity(&self) -> Option<&GrokSessionAffinityKey> { @@ -149,6 +167,9 @@ impl GrokResponsesRequest { return; }; self.session_id = Some(session_id.to_owned()); + if self.reasoning_replay_session_id.is_some() { + self.reasoning_replay_session_id = Some(session_id.to_owned()); + } self.affinity = None; self.body.insert( "prompt_cache_key".to_owned(), @@ -159,6 +180,7 @@ impl GrokResponsesRequest { pub(crate) fn clear_session(&mut self) { self.body.remove("prompt_cache_key"); self.session_id = None; + self.reasoning_replay_session_id = None; self.affinity = None; } @@ -207,13 +229,16 @@ impl GrokResponsesRequest { ); sanitize_account_identity(&mut body); sanitize_foreign_client_metadata(&mut body); - let response_transform = normalize_responses_request(&mut body)?; + let mut response_transform = normalize_responses_request(&mut body)?; + response_transform.observe_client_cache_tools(); if enable_cache_route { - enable_grok_prompt_cache_route(&mut body); + enable_grok_prompt_cache_route(&mut body, upstream_model, &mut response_transform); } let (session_id, affinity) = identity.map_or((None, None), |(session_id, affinity)| { (Some(session_id), Some(affinity)) }); + let reasoning_replay_session_id = + session_seed.is_some().then(|| session_id.clone()).flatten(); match session_id.as_ref() { Some(session_id) => { body.insert( @@ -230,6 +255,7 @@ impl GrokResponsesRequest { Ok(Self { body, session_id, + reasoning_replay_session_id, affinity, response_transform, }) @@ -321,23 +347,61 @@ fn sanitize_foreign_client_metadata(body: &mut Map) { } } -fn enable_grok_prompt_cache_route(body: &mut Map) { +fn enable_grok_prompt_cache_route( + body: &mut Map, + upstream_model: &str, + response: &mut GrokResponseTransform, +) { let tools = body .entry("tools".to_owned()) .or_insert_with(|| Value::Array(Vec::new())); let Value::Array(tools) = tools else { return; }; - // 缓存路由不得改变模型可调用的工具集：客户端带工具时无法既保留其 - // tool_choice 语义又阻止注入的原生搜索被调用，因此直接不注入。 - if !tools.is_empty() { + response.observe_upstream_cache_tools(tools); + if is_cache_media_model(upstream_model) || has_tool_type(tools, "image_generation") { return; } - for tool in GROK_CACHE_ROUTE_TOOLS { - tools.push(json_object([("type", Value::String((*tool).to_owned()))])); + if tools.is_empty() { + for tool in GROK_CACHE_ROUTE_TOOLS { + tools.push(json_object([("type", Value::String((*tool).to_owned()))])); + response.mark_injected_cache_tool(tool); + } + // 无客户端工具时以 none 选中缓存路由，同时不授予搜索能力。 + body.insert("tool_choice".to_owned(), Value::String("none".to_owned())); + } else if !has_tool_type(tools, "x_search") { + tools.push(json_object([( + "type", + Value::String("x_search".to_owned()), + )])); + response.mark_injected_cache_tool("x_search"); + } + append_x_search_to_allowed_tools(body); +} + +fn is_cache_media_model(model: &str) -> bool { + let model = model.trim().to_ascii_lowercase(); + ["image", "imagine", "video"] + .into_iter() + .any(|marker| model.contains(marker)) +} + +fn append_x_search_to_allowed_tools(body: &mut Map) { + let Some(choice) = body.get_mut("tool_choice").and_then(Value::as_object_mut) else { + return; + }; + if choice.get("type").and_then(Value::as_str) != Some("allowed_tools") { + return; + } + let Some(allowed) = choice.get_mut("tools").and_then(Value::as_array_mut) else { + return; + }; + if !has_tool_type(allowed, "x_search") { + allowed.push(json_object([( + "type", + Value::String("x_search".to_owned()), + )])); } - // 禁用工具调用，使原生搜索声明只承担缓存路由作用。 - body.insert("tool_choice".to_owned(), Value::String("none".to_owned())); } fn explicit_session_seed(request: &GenerateRequest, body: &Map) -> Option { @@ -674,6 +738,11 @@ pub(crate) struct GrokResponseTransform { aliases: BTreeMap, visible_tools: Vec, legacy_local_shell: bool, + filter_x_search: bool, + injected_tool_types: BTreeSet, + client_declared_tools: BTreeSet, + dropped_output_indexes: BTreeSet, + dropped_item_ids: BTreeSet, stream_calls: BTreeMap, stream_keys: BTreeMap, } @@ -698,7 +767,42 @@ impl GrokTransformedWireEvent { impl GrokResponseTransform { #[must_use] pub(crate) fn is_empty(&self) -> bool { - self.aliases.is_empty() && self.visible_tools.is_empty() && !self.legacy_local_shell + self.aliases.is_empty() + && self.visible_tools.is_empty() + && !self.legacy_local_shell + && !self.filter_x_search + && self.injected_tool_types.is_empty() + } + + fn observe_client_cache_tools(&mut self) { + for tool in &self.visible_tools { + let kind = tool.get("type").and_then(Value::as_str).unwrap_or_default(); + if kind == "x_search" { + self.filter_x_search = true; + } + if matches!(kind, "function" | "custom") + && let Some(name) = tool + .get("name") + .and_then(Value::as_str) + .map(str::trim) + .filter(|name| !name.is_empty()) + { + self.client_declared_tools.insert(name.to_owned()); + } + } + } + + fn observe_upstream_cache_tools(&mut self, tools: &[Value]) { + if has_tool_type(tools, "x_search") { + self.filter_x_search = true; + } + } + + fn mark_injected_cache_tool(&mut self, kind: &str) { + self.injected_tool_types.insert(kind.to_owned()); + if kind == "x_search" { + self.filter_x_search = true; + } } pub(crate) fn rewrite_stream_event( @@ -706,6 +810,20 @@ impl GrokResponseTransform { event_type: &str, mut value: Value, ) -> Result, GrokRequestEncodeError> { + if self.cache_filter_enabled() { + if value + .get("item") + .is_some_and(|item| self.is_internal_cache_call(item)) + { + self.record_dropped_cache_item(&value); + return Ok(Vec::new()); + } + self.filter_cache_envelope(&mut value)?; + if self.references_dropped_cache_item(&value) { + return Ok(Vec::new()); + } + self.compact_cache_output_index(&mut value); + } if event_type == "response.output_item.added" && let Some(item) = value.get("item").and_then(Value::as_object) && let Some(primary) = self.remember_stream_call(item) @@ -1031,9 +1149,145 @@ impl GrokResponseTransform { ]) } + fn cache_filter_enabled(&self) -> bool { + self.filter_x_search || !self.injected_tool_types.is_empty() + } + + fn is_internal_cache_call(&self, item: &Value) -> bool { + let Some(item) = item.as_object() else { + return false; + }; + let kind = string_field(item, "type").trim(); + if kind == "web_search_call" { + return self.injected_tool_types.contains("web_search"); + } + if !matches!(kind, "custom_tool_call" | "function_call") { + return false; + } + if string_field(item, "call_id").trim().starts_with("xs_call") { + return true; + } + let name = string_field(item, "name").trim(); + matches!( + name, + "x_user_search" | "x_semantic_search" | "x_keyword_search" | "x_thread_fetch" + ) && string_field(item, "namespace").trim().is_empty() + && !self.client_declared_tools.contains(name) + } + + fn record_dropped_cache_item(&mut self, payload: &Value) { + if let Some(index) = payload.get("output_index").and_then(Value::as_u64) { + self.dropped_output_indexes.insert(index); + } + let Some(item) = payload.get("item").and_then(Value::as_object) else { + return; + }; + for field in ["id", "call_id"] { + if let Some(value) = item + .get(field) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + { + self.dropped_item_ids.insert(value.to_owned()); + } + } + } + + fn references_dropped_cache_item(&self, payload: &Value) -> bool { + if payload + .get("output_index") + .and_then(Value::as_u64) + .is_some_and(|index| self.dropped_output_indexes.contains(&index)) + { + return true; + } + ["item_id", "call_id"].into_iter().any(|field| { + payload + .get(field) + .and_then(Value::as_str) + .map(str::trim) + .is_some_and(|value| self.dropped_item_ids.contains(value)) + }) + } + + fn compact_cache_output_index(&self, payload: &mut Value) { + let Some(index) = payload.get("output_index").and_then(Value::as_u64) else { + return; + }; + let removed_before = self + .dropped_output_indexes + .iter() + .filter(|dropped| **dropped < index) + .count() as u64; + if removed_before > 0 + && let Some(payload) = payload.as_object_mut() + { + payload.insert( + "output_index".to_owned(), + Value::from(index.saturating_sub(removed_before)), + ); + } + } + + fn filter_cache_envelope(&self, payload: &mut Value) -> Result<(), GrokRequestEncodeError> { + let Some(payload) = payload.as_object_mut() else { + return Ok(()); + }; + self.filter_cache_output(payload)?; + self.filter_injected_cache_tools(payload)?; + if let Some(response) = payload.get_mut("response").and_then(Value::as_object_mut) { + self.filter_cache_output(response)?; + self.filter_injected_cache_tools(response)?; + } + Ok(()) + } + + fn filter_cache_output( + &self, + envelope: &mut Map, + ) -> Result<(), GrokRequestEncodeError> { + let Some(output) = envelope.get_mut("output") else { + return Ok(()); + }; + let output = output + .as_array_mut() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + output.retain(|item| !self.is_internal_cache_call(item)); + Ok(()) + } + + fn filter_injected_cache_tools( + &self, + envelope: &mut Map, + ) -> Result<(), GrokRequestEncodeError> { + if self.injected_tool_types.is_empty() { + return Ok(()); + } + let Some(tools) = envelope.get_mut("tools") else { + return Ok(()); + }; + let tools = tools + .as_array_mut() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + tools.retain(|tool| { + tool.get("type") + .and_then(Value::as_str) + .is_none_or(|kind| !self.injected_tool_types.contains(kind.trim())) + }); + if tools.is_empty() { + envelope.remove("tools"); + } + Ok(()) + } + fn restore_visible_tools(&self, response: &mut Map) { if response.contains_key("tools") { - response.insert("tools".to_owned(), Value::Array(self.visible_tools.clone())); + if self.visible_tools.is_empty() { + response.remove("tools"); + } else { + response.insert("tools".to_owned(), Value::Array(self.visible_tools.clone())); + } } } } @@ -1125,6 +1379,8 @@ impl fmt::Debug for GrokResponseTransform { .field("alias_count", &self.aliases.len()) .field("visible_tool_count", &self.visible_tools.len()) .field("legacy_local_shell", &self.legacy_local_shell) + .field("filter_x_search", &self.filter_x_search) + .field("injected_tool_types", &self.injected_tool_types) .finish() } } @@ -1982,6 +2238,9 @@ impl ToolNormalizer { ]), ); } + "allowed_tools" => { + self.normalize_allowed_tools_choice(payload, object, normalized_tools)?; + } "function" => self.normalize_function_tool_choice(payload, object)?, _ => { if let Some(hosted_kind) = normalize_hosted_tool_choice_kind(&kind) { @@ -2004,6 +2263,94 @@ impl ToolNormalizer { Ok(()) } + fn normalize_allowed_tools_choice( + &self, + payload: &mut Map, + mut object: Map, + normalized_tools: &[Value], + ) -> Result<(), GrokRequestEncodeError> { + if !matches!( + object.get("mode").and_then(Value::as_str), + Some("auto" | "required") + ) { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + let allowed = object + .get_mut("tools") + .and_then(Value::as_array_mut) + .filter(|tools| !tools.is_empty()) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + for allowed_tool in allowed { + self.normalize_allowed_tool_reference(allowed_tool, normalized_tools)?; + } + payload.insert("tool_choice".to_owned(), Value::Object(object)); + Ok(()) + } + + fn normalize_allowed_tool_reference( + &self, + allowed_tool: &mut Value, + normalized_tools: &[Value], + ) -> Result<(), GrokRequestEncodeError> { + let object = allowed_tool + .as_object_mut() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + let kind = string_field(object, "type").trim().to_owned(); + match kind.as_str() { + "function" => { + let choice = match object.get_mut("function") { + Some(Value::Object(function)) => function, + Some(_) => return Err(GrokRequestEncodeError::InvalidRequestNormalization), + None => object, + }; + rewrite_namespace_choice(choice, &self.identity_aliases)?; + let name = string_field(choice, "name").trim(); + if name.is_empty() || !has_named_tool(normalized_tools, "function", name) { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + } + "custom" => { + let identity = ToolIdentity::new( + ToolKind::Custom, + string_field(object, "namespace").trim(), + string_field(object, "name").trim(), + ); + let alias = self + .identity_aliases + .get(&identity) + .cloned() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + object.insert("type".to_owned(), Value::String("function".to_owned())); + object.insert("name".to_owned(), Value::String(alias)); + object.remove("namespace"); + } + "apply_patch" | "tool_search" => { + let identity = if kind == "apply_patch" { + ToolIdentity::new(ToolKind::ApplyPatch, "", "apply_patch") + } else { + ToolIdentity::new(ToolKind::ToolSearch, "", "tool_search") + }; + let alias = self + .identity_aliases + .get(&identity) + .cloned() + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + object.clear(); + object.insert("type".to_owned(), Value::String("function".to_owned())); + object.insert("name".to_owned(), Value::String(alias)); + } + _ => { + let hosted_kind = normalize_hosted_tool_choice_kind(&kind) + .ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?; + if !has_tool_type(normalized_tools, hosted_kind) { + return Err(GrokRequestEncodeError::InvalidRequestNormalization); + } + object.insert("type".to_owned(), Value::String(hosted_kind.to_owned())); + } + } + Ok(()) + } + fn normalize_function_tool_choice( &self, payload: &mut Map, @@ -2080,6 +2427,13 @@ fn has_tool_type(tools: &[Value], kind: &str) -> bool { .any(|tool| tool.pointer("/type").and_then(Value::as_str) == Some(kind)) } +fn has_named_tool(tools: &[Value], kind: &str, name: &str) -> bool { + tools.iter().any(|tool| { + tool.pointer("/type").and_then(Value::as_str) == Some(kind) + && tool.pointer("/name").and_then(Value::as_str) == Some(name) + }) +} + fn tools_of_type(tools: &[Value], kind: &str) -> Vec { tools .iter() diff --git a/backend/crates/providers/xai/tests/credential/selector.rs b/backend/crates/providers/xai/tests/credential/selector.rs index c41293995..40e837713 100644 --- a/backend/crates/providers/xai/tests/credential/selector.rs +++ b/backend/crates/providers/xai/tests/credential/selector.rs @@ -324,7 +324,7 @@ async fn unauthorized_feedback_records_runtime_cooldown_without_persisting_accou } #[tokio::test] -async fn cooldown_for_superseded_credential_does_not_block_rotated_account() { +async fn account_scoped_cooldown_survives_credential_rotation() { let fixture = SelectorFixture::new(&["revision-fence"]).await; let id = account_id("revision-fence"); let session = fixture @@ -369,17 +369,25 @@ async fn cooldown_for_superseded_credential_does_not_block_rotated_account() { 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()); + assert!(matches!( + fixture + .selector + .select(fixture.request(BTreeSet::new())) + .await, + Err(GrokSessionSelectorError::NoEligibleSession) + )); + let account = fixture.store.account(&id).expect("rotated account"); + assert_eq!(account.revision().get(), 2); + assert_eq!(account.availability(), AccountAvailability::Cooldown); + assert!( + account + .cooldown_until() + .is_some_and(|until| until > SystemTime::now()) + ); } #[tokio::test] -async fn rate_limit_feedback_records_runtime_cooldown_without_persisting_account_state() { +async fn rate_limit_feedback_persists_the_grok2api_cooldown_state() { let fixture = SelectorFixture::new(&["rate-limit", "available"]).await; let session = fixture .selector @@ -397,7 +405,12 @@ async fn rate_limit_feedback_records_runtime_cooldown_without_persisting_account ) .await; let account = fixture.store.account(&selected).expect("account"); - assert_eq!(account.availability(), AccountAvailability::Ready); + assert_eq!(account.availability(), AccountAvailability::Cooldown); + assert!( + account + .cooldown_until() + .is_some_and(|until| until > SystemTime::now()) + ); assert!( fixture .cooldowns @@ -413,7 +426,35 @@ async fn rate_limit_feedback_records_runtime_cooldown_without_persisting_account } #[tokio::test] -async fn unknown_payment_required_feedback_only_records_runtime_account_cooldown() { +async fn successful_request_clears_the_persisted_cooldown_state() { + let fixture = SelectorFixture::new(&["cooldown-success"]).await; + 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(5)), + }, + ) + .await; + + fixture.selector.record_success(&session).await; + + let account = fixture + .store + .account(session.account_id()) + .expect("account"); + assert_eq!(account.availability(), AccountAvailability::Ready); + assert!(account.cooldown_until().is_none()); +} + +#[tokio::test] +async fn payment_required_feedback_persists_account_quota_exhaustion() { let fixture = SelectorFixture::new(&["payment-required", "available"]).await; let session = fixture .selector @@ -437,9 +478,9 @@ async fn unknown_payment_required_feedback_only_records_runtime_account_cooldown .account(&selected) .expect("selected account") .availability(), - AccountAvailability::Ready + AccountAvailability::QuotaExhausted ); - assert!(fixture.cooldowns.cooldown(&selected).is_some()); + assert!(fixture.cooldowns.cooldown(&selected).is_none()); let next = fixture .selector .select(fixture.request(BTreeSet::new())) @@ -538,7 +579,7 @@ async fn model_access_feedback_reports_model_cooldown_without_blocking_the_accou } #[tokio::test] -async fn interrupted_stream_feedback_records_runtime_cooldown_without_persisting_account_state() { +async fn interrupted_stream_feedback_persists_the_grok2api_cooldown_state() { let fixture = SelectorFixture::new(&["stream-interrupted"]).await; let session = fixture .selector @@ -554,7 +595,12 @@ async fn interrupted_stream_feedback_records_runtime_cooldown_without_persisting .store .account(session.account_id()) .expect("account"); - assert_eq!(account.availability(), AccountAvailability::Ready); + assert_eq!(account.availability(), AccountAvailability::Cooldown); + assert!( + account + .cooldown_until() + .is_some_and(|until| until > SystemTime::now()) + ); assert!( fixture .cooldowns @@ -595,6 +641,75 @@ async fn quota_feedback_uses_common_quota_exhausted_state() { ); } +#[tokio::test] +async fn quota_feedback_should_follow_the_account_across_a_credential_rotation() { + let fixture = SelectorFixture::new(&["quota-rotation"]).await; + let id = account_id("quota-rotation"); + let session = fixture + .selector + .select(fixture.request(BTreeSet::new())) + .await + .expect("session"); + 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-quota-rotation"), + next_refresh_at: chrono::Utc::now() + chrono::Duration::minutes(30), + }) + .expect("rotate"); + fixture + .store + .compare_and_swap_credential(prepared.credential) + .await + .expect("persist rotation"); + + fixture + .selector + .record_failure(&session, GrokCredentialFailure::FreeQuotaExhausted) + .await; + + assert_eq!( + fixture.store.account(&id).expect("account").availability(), + AccountAvailability::QuotaExhausted + ); +} + +#[tokio::test] +async fn ordinary_success_should_not_clear_a_confirmed_quota_exhaustion() { + let fixture = SelectorFixture::new(&["quota-success"]).await; + let session = fixture + .selector + .select(fixture.request(BTreeSet::new())) + .await + .expect("session"); + fixture + .selector + .record_failure(&session, GrokCredentialFailure::FreeQuotaExhausted) + .await; + + fixture.selector.record_success(&session).await; + + assert_eq!( + fixture + .store + .account(session.account_id()) + .expect("account") + .availability(), + AccountAvailability::QuotaExhausted + ); +} + #[tokio::test] async fn excluded_account_is_never_selected_again() { let fixture = SelectorFixture::new(&["excluded"]).await; diff --git a/backend/crates/providers/xai/tests/provider/contract.rs b/backend/crates/providers/xai/tests/provider/contract.rs index 5d1a6125e..eea9e8c22 100644 --- a/backend/crates/providers/xai/tests/provider/contract.rs +++ b/backend/crates/providers/xai/tests/provider/contract.rs @@ -106,6 +106,15 @@ const SUCCESS_SSE: &[u8] = concat!( ) .as_bytes(); +fn replay_ciphertext(seed: u8) -> String { + use base64::Engine as _; + + let bytes = (0_u16..128) + .map(|value| (value as u8).wrapping_add(seed)) + .collect::>(); + base64::engine::general_purpose::STANDARD_NO_PAD.encode(bytes) +} + fn stateful_sse(encrypted_content: &str) -> Vec { format!( concat!( @@ -121,6 +130,25 @@ fn stateful_sse(encrypted_content: &str) -> Vec { .into_bytes() } +fn stateful_sse_with_assistant(encrypted_content: &str, assistant: &str) -> Vec { + let assistant = serde_json::to_string(assistant).expect("assistant text"); + format!( + concat!( + "event: response.created\n", + "data: {{\"type\":\"response.created\",\"response\":{{\"id\":\"resp_state\",\"model\":\"grok-4.5\"}}}}\n\n", + "event: response.output_item.done\n", + "data: {{\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{{\"id\":\"reason_account_bound\",\"type\":\"reasoning\",\"status\":\"completed\",\"summary\":[],\"content\":null,\"encrypted_content\":\"{}\"}}}}\n\n", + "event: response.output_item.done\n", + "data: {{\"type\":\"response.output_item.done\",\"output_index\":1,\"item\":{{\"id\":\"message_account_bound\",\"type\":\"message\",\"role\":\"assistant\",\"status\":\"completed\",\"content\":[{{\"type\":\"output_text\",\"text\":{assistant}}}]}}}}\n\n", + "event: response.completed\n", + "data: {{\"type\":\"response.completed\",\"response\":{{\"id\":\"resp_state\",\"model\":\"grok-4.5\",\"status\":\"completed\"}}}}\n\n" + ), + encrypted_content, + assistant = assistant, + ) + .into_bytes() +} + fn custom_apply_patch_sse(patch: &str) -> Vec { let arguments = serde_json::to_string(&json!({"patch": patch})).expect("patch arguments"); let arguments_json = serde_json::to_string(&arguments).expect("arguments JSON string"); @@ -278,6 +306,54 @@ impl GrokSessionSelector for StubSelector { } } +struct SequencedAccountSelector { + accounts: Mutex>, +} + +impl SequencedAccountSelector { + fn new( + accounts: impl IntoIterator, + ) -> Arc { + Arc::new(Self { + accounts: Mutex::new(accounts.into_iter().collect()), + }) + } +} + +impl GrokSessionSelector for SequencedAccountSelector { + fn select(&self, request: GrokSessionSelection) -> GrokSessionSelectorFuture<'_> { + let selected = self.accounts.lock().expect("account sequence").pop_front(); + Box::pin(async move { + let id = selected.ok_or(GrokSessionSelectorError::NoEligibleSession)?; + if request.excluded_accounts().contains(&id) + || request + .required_account() + .is_some_and(|required| required != &id) + { + return Err(GrokSessionSelectorError::NoEligibleSession); + } + SelectedGrokSession::new( + id, + CredentialRevision::new(1).expect("revision"), + SecretValue::new("oauth-access"), + SecretValue::new("verified-user"), + Some(SecretValue::new("user@example.com")), + GrokSessionBinding::new("acct_provider").expect("binding"), + (), + ) + .map_err(|_| GrokSessionSelectorError::InvalidSession) + }) + } + + fn record_failure<'a>( + &'a self, + _: &'a SelectedGrokSession, + _: GrokCredentialFailure, + ) -> GrokCredentialFeedbackFuture<'a> { + Box::pin(async {}) + } +} + enum InferenceMode { Success, SuccessWithMetrics(GrokInferenceTransportMetrics), @@ -481,7 +557,7 @@ async fn provider_with_recovery( } async fn provider_with_catalog_transport( - selector: Arc, + selector: Arc, transport: Arc, recovery: Arc, catalog_transport: Arc, @@ -569,6 +645,19 @@ fn operation() -> Operation { Operation::Generate(GenerateRequest::from_protocol_payload(payload)) } +fn reasoning_replay_operation(session: &str, input: serde_json::Value) -> Operation { + let payload = ProtocolPayload::json_object( + "openai", + Map::from_iter([ + ("model".to_owned(), json!("client-model")), + ("prompt_cache_key".to_owned(), json!(session)), + ("input".to_owned(), input), + ]), + ) + .expect("OpenAI payload"); + Operation::Generate(GenerateRequest::from_protocol_payload(payload)) +} + fn operation_with_reasoning_effort(effort: &str) -> Operation { let payload = ProtocolPayload::json_object( "openai", @@ -621,6 +710,26 @@ fn compaction_operation() -> Operation { Operation::Generate(GenerateRequest::from_protocol_payload(payload)) } +fn compaction_operation_with_prompt_cache(session: &str) -> Operation { + let payload = ProtocolPayload::json_object( + "openai", + Map::from_iter([ + ("model".to_owned(), json!("client-model")), + ("prompt_cache_key".to_owned(), json!(session)), + ( + "input".to_owned(), + json!([ + {"type": "message", "role": "user", "content": "history"}, + {"type": "compaction_trigger"} + ]), + ), + ("stream".to_owned(), json!(true)), + ]), + ) + .expect("OpenAI payload"); + Operation::Generate(GenerateRequest::from_protocol_payload(payload)) +} + fn compaction_operation_with_state(state: ProviderSessionState) -> Operation { let payload = ProtocolPayload::json_object( "openai", @@ -737,6 +846,18 @@ fn context_with_continuation_attempt( context(CancellationToken::new(), Some(continuation)).with_continuation_attempt(attempt) } +async fn execute_successfully(provider: &GrokBuildProvider, operation: Operation) { + let mut stream = provider + .execute( + provider_request_with_operation("xai", operation), + context(CancellationToken::new(), None), + ) + .await + .expect("provider stream"); + let events = stream.by_ref().collect::>().await; + assert!(events.iter().all(Result::is_ok)); +} + #[tokio::test] async fn request_observation_preserves_raw_xai_reasoning_effort() { let provider = provider(StubSelector::success(), StubInferenceTransport::success()).await; @@ -1300,7 +1421,7 @@ async fn explicit_quota_http_429_marks_provider_error_replay_safe() { } #[tokio::test] -async fn unknown_http_402_is_retryable_without_account_quota_feedback() { +async fn unknown_http_402_is_retryable_with_payment_quota_feedback() { let selector = StubSelector::success(); let transport = StubInferenceTransport::error( GrokInferenceTransportError::new( @@ -1750,6 +1871,190 @@ async fn external_previous_response_is_rejected_before_selection() { assert_eq!(selector.calls.load(Ordering::SeqCst), 0); } +#[tokio::test] +async fn reasoning_replay_is_reused_only_for_the_same_explicit_session() { + let encrypted_content = replay_ciphertext(0); + let transport = StubInferenceTransport::sequence([ + InferenceMode::SuccessBody(stateful_sse(&encrypted_content)), + InferenceMode::Success, + InferenceMode::Success, + ]); + let provider = provider(StubSelector::success(), transport.clone()).await; + + execute_successfully( + &provider, + reasoning_replay_operation( + "session-a", + json!([{"type":"message","role":"user","content":"first"}]), + ), + ) + .await; + execute_successfully( + &provider, + reasoning_replay_operation( + "session-b", + json!([{"type":"message","role":"user","content":"other"}]), + ), + ) + .await; + execute_successfully( + &provider, + reasoning_replay_operation( + "session-a", + json!([{"type":"message","role":"user","content":"next"}]), + ), + ) + .await; + + let requests = transport.requests.lock().expect("requests"); + let other: serde_json::Value = + serde_json::from_slice(requests[1].body()).expect("other session body"); + let replayed: serde_json::Value = + serde_json::from_slice(requests[2].body()).expect("replayed body"); + assert!( + other["input"] + .as_array() + .is_some_and(|input| input.iter().all(|item| item["type"] != "reasoning")) + ); + assert_eq!( + replayed.pointer("/input/0/encrypted_content"), + Some(&json!(encrypted_content)) + ); +} + +#[tokio::test] +async fn reasoning_replay_is_scoped_to_the_selected_account() { + let encrypted_content = replay_ciphertext(1); + let transport = StubInferenceTransport::sequence([ + InferenceMode::SuccessBody(stateful_sse(&encrypted_content)), + InferenceMode::Success, + InferenceMode::Success, + ]); + let selector = SequencedAccountSelector::new([ + account_id("replay-a"), + account_id("replay-b"), + account_id("replay-a"), + ]); + let provider = provider_with_catalog_transport( + selector, + transport.clone(), + StubRecovery::new(GrokCredentialRecoveryOutcome::Unavailable), + Arc::new(StaticCatalogTransport), + ) + .await; + + for prompt in ["first", "different account", "owner again"] { + execute_successfully( + &provider, + reasoning_replay_operation( + "shared-session", + json!([{"type":"message","role":"user","content":prompt}]), + ), + ) + .await; + } + + let requests = transport.requests.lock().expect("requests"); + let different_account: serde_json::Value = + serde_json::from_slice(requests[1].body()).expect("different account request body"); + let owner_again: serde_json::Value = + serde_json::from_slice(requests[2].body()).expect("owner request body"); + assert!( + different_account["input"] + .as_array() + .is_some_and(|input| input.iter().all(|item| item["type"] != "reasoning")) + ); + assert_eq!( + owner_again.pointer("/input/0/encrypted_content"), + Some(&json!(encrypted_content)) + ); +} + +#[tokio::test] +async fn reasoning_replay_rejects_mismatched_assistant_history() { + let encrypted_content = replay_ciphertext(2); + let transport = StubInferenceTransport::sequence([ + InferenceMode::SuccessBody(stateful_sse_with_assistant( + &encrypted_content, + "cached answer", + )), + InferenceMode::Success, + ]); + let provider = provider(StubSelector::success(), transport.clone()).await; + + execute_successfully( + &provider, + reasoning_replay_operation( + "assistant-session", + json!([{"type":"message","role":"user","content":"first"}]), + ), + ) + .await; + execute_successfully( + &provider, + reasoning_replay_operation( + "assistant-session", + json!([ + {"type":"message","role":"user","content":"first"}, + {"type":"message","role":"assistant","content":[{"type":"output_text","text":"different answer"}]}, + {"type":"message","role":"user","content":"next"} + ]), + ), + ) + .await; + + let requests = transport.requests.lock().expect("requests"); + let body: serde_json::Value = serde_json::from_slice(requests[1].body()).expect("second body"); + assert!( + body["input"] + .as_array() + .is_some_and(|input| input.iter().all(|item| item["type"] != "reasoning")) + ); +} + +#[tokio::test] +async fn successful_compaction_clears_reasoning_replay_for_the_session() { + let encrypted_content = replay_ciphertext(3); + let summary = valid_compaction_summary("reasoning replay cleared"); + let transport = StubInferenceTransport::sequence([ + InferenceMode::SuccessBody(stateful_sse(&encrypted_content)), + InferenceMode::SuccessBody(compaction_sse(&summary, None)), + InferenceMode::Success, + ]); + let provider = provider(StubSelector::success(), transport.clone()).await; + + execute_successfully( + &provider, + reasoning_replay_operation( + "compact-session", + json!([{"type":"message","role":"user","content":"first"}]), + ), + ) + .await; + execute_successfully( + &provider, + compaction_operation_with_prompt_cache("compact-session"), + ) + .await; + execute_successfully( + &provider, + reasoning_replay_operation( + "compact-session", + json!([{"type":"message","role":"user","content":"after compact"}]), + ), + ) + .await; + + let requests = transport.requests.lock().expect("requests"); + let body: serde_json::Value = + serde_json::from_slice(requests[2].body()).expect("post-compaction body"); + assert!( + body["input"] + .as_array() + .is_some_and(|input| input.iter().all(|item| item["type"] != "reasoning")) + ); +} + #[tokio::test] async fn connection_state_inherits_session_and_recovers_reasoning_on_pinned_account() { use base64::Engine as _; diff --git a/backend/crates/providers/xai/tests/transport/canonical.rs b/backend/crates/providers/xai/tests/transport/canonical.rs index f9f2a6e12..7554354a7 100644 --- a/backend/crates/providers/xai/tests/transport/canonical.rs +++ b/backend/crates/providers/xai/tests/transport/canonical.rs @@ -642,6 +642,76 @@ fn decoder_should_preserve_image_and_hosted_tool_items_without_inventing_facts() ); } +#[test] +fn decoder_should_hide_injected_search_calls_and_compact_later_output_indexes() { + let request = tool_request(serde_json::json!({ + "model": "client", + "prompt_cache_key": "cache-session", + "input": "hello" + })); + let body = concat!( + "event: response.created\n", + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_cache\",\"tools\":[{\"type\":\"web_search\"},{\"type\":\"x_search\"}]}}\n\n", + "event: response.output_item.added\n", + "data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"id\":\"item_xs\",\"call_id\":\"xs_call_1\",\"name\":\"x_semantic_search\",\"arguments\":\"{}\"}}\n\n", + "event: response.function_call_arguments.delta\n", + "data: {\"type\":\"response.function_call_arguments.delta\",\"output_index\":0,\"item_id\":\"item_xs\",\"call_id\":\"xs_call_1\",\"delta\":\"{}\"}\n\n", + "event: response.output_item.done\n", + "data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"id\":\"item_xs\",\"call_id\":\"xs_call_1\",\"name\":\"x_semantic_search\",\"arguments\":\"{}\"}}\n\n", + "event: response.output_item.added\n", + "data: {\"type\":\"response.output_item.added\",\"output_index\":1,\"item\":{\"type\":\"image_generation_call\",\"id\":\"image_visible\"}}\n\n", + "event: response.completed\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_cache\",\"status\":\"completed\",\"tools\":[{\"type\":\"web_search\"},{\"type\":\"x_search\"}],\"output\":[{\"type\":\"function_call\",\"call_id\":\"xs_call_1\",\"name\":\"x_semantic_search\",\"arguments\":\"{}\"},{\"type\":\"image_generation_call\",\"id\":\"image_visible\"}]}}\n\n", + ); + + let events = GrokCanonicalDecoder::for_request("grok-4.5", &request) + .push(body.as_bytes()) + .expect("filtered cache route response"); + let wire = wire_events(&events); + + assert_eq!(wire.len(), 3); + assert_eq!( + wire[1].data().get("output_index"), + Some(&serde_json::json!(0)) + ); + assert_eq!( + wire[2].data().pointer("/response/output/0/id"), + Some(&serde_json::json!("image_visible")) + ); + assert_eq!(wire[2].data().pointer("/response/output/1"), None); + assert_eq!(wire[2].data().pointer("/response/tools"), None); +} + +#[test] +fn decoder_should_preserve_an_internal_search_name_declared_by_the_client() { + let request = tool_request(serde_json::json!({ + "model": "client", + "prompt_cache_key": "cache-session", + "input": "hello", + "tools": [{ + "type": "function", + "name": "x_user_search", + "parameters": {"type": "object"} + }] + })); + let body = concat!( + "event: response.created\n", + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_declared\"}}\n\n", + "event: response.output_item.added\n", + "data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"id\":\"item_client\",\"call_id\":\"call_client\",\"name\":\"x_user_search\",\"arguments\":\"{}\"}}\n\n", + "event: response.completed\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_declared\",\"status\":\"completed\"}}\n\n", + ); + + let events = GrokCanonicalDecoder::for_request("grok-4.5", &request) + .push(body.as_bytes()) + .expect("client-declared search function"); + + assert!(wire_events(&events).iter().any(|event| { + event.data().pointer("/item/name") == Some(&serde_json::json!("x_user_search")) + })); +} + #[test] fn decoder_should_tolerate_mismatched_event_identity_and_preserve_the_wire_event() { let body = concat!( diff --git a/backend/crates/providers/xai/tests/transport/headers.rs b/backend/crates/providers/xai/tests/transport/headers.rs index a8cf9f74a..e3c7ba496 100644 --- a/backend/crates/providers/xai/tests/transport/headers.rs +++ b/backend/crates/providers/xai/tests/transport/headers.rs @@ -91,7 +91,7 @@ fn headers_should_bind_identity_to_the_selected_oauth_account() { } #[test] -fn headers_should_keep_openai_cache_requests_out_of_shell_session_mode() { +fn headers_should_send_both_stable_session_headers_without_inventing_a_turn_index() { let headers = build_grok_headers( &crate::support::xai_wire_profile(), &selected_session(), @@ -109,7 +109,7 @@ fn headers_should_keep_openai_cache_requests_out_of_shell_session_mode() { }; assert_eq!(value("x-grok-conv-id"), Some("session-fixture")); - assert_eq!(value("x-grok-session-id"), None); + assert_eq!(value("x-grok-session-id"), Some("session-fixture")); assert_eq!(value("x-grok-turn-idx"), None); } diff --git a/backend/crates/providers/xai/tests/transport/request.rs b/backend/crates/providers/xai/tests/transport/request.rs index 1ad82edfa..c68f75634 100644 --- a/backend/crates/providers/xai/tests/transport/request.rs +++ b/backend/crates/providers/xai/tests/transport/request.rs @@ -256,7 +256,7 @@ fn explicit_session_should_enable_the_noop_native_cache_route() { } #[test] -fn explicit_session_should_not_add_invokable_cache_tools_to_client_tool_requests() { +fn explicit_session_should_add_only_internal_x_search_to_client_tool_requests() { let request = raw_request(json!({ "model": "client", "prompt_cache_key": "conversation-42", @@ -273,15 +273,15 @@ fn explicit_session_should_not_add_invokable_cache_tools_to_client_tool_requests GrokResponsesRequest::encode(&request, "grok-4.5", &client_key()).expect("request"); let body = Value::Object(encoded.body().clone()); - // 缓存路由不得扩大模型可调用的工具集：带客户端工具时不注入原生搜索。 assert_eq!( - body.pointer("/tools"), - Some(&json!([{ + body.pointer("/tools/0"), + Some(&json!({ "type": "function", "name": "read_file", "parameters": {"type": "object"} - }])) + })) ); + assert_eq!(body.pointer("/tools/1"), Some(&json!({"type": "x_search"}))); assert_eq!(body.pointer("/tool_choice"), Some(&json!("auto"))); } @@ -384,7 +384,7 @@ fn function_parameters_should_report_the_invalid_nullable_root_field() { } #[test] -fn explicit_session_should_not_inject_cache_tools_after_codex_additional_tools_normalization() { +fn explicit_session_should_add_x_search_after_codex_additional_tools_normalization() { let request = raw_request(json!({ "model": "client", "prompt_cache_key": "conversation-42", @@ -417,11 +417,38 @@ fn explicit_session_should_not_inject_cache_tools_after_codex_additional_tools_n Some(&json!(["patch"])) ); assert_eq!(body.pointer("/tools/1/name"), Some(&json!("read_file"))); - // 归一化装载的会话工具视同客户端工具：不再追加可被调用的原生搜索。 - assert_eq!(body.pointer("/tools/2"), None); + assert_eq!(body.pointer("/tools/2"), Some(&json!({"type": "x_search"}))); assert_eq!(body.pointer("/tool_choice"), Some(&json!("auto"))); } +#[test] +fn explicit_session_should_add_x_search_to_allowed_tools_choice() { + let request = raw_request(json!({ + "model": "client", + "prompt_cache_key": "conversation-42", + "input": "first", + "tools": [{ + "type": "function", + "name": "read_file", + "parameters": {"type": "object"} + }], + "tool_choice": { + "type": "allowed_tools", + "mode": "auto", + "tools": [{"type": "function", "name": "read_file"}] + } + })); + + let encoded = + GrokResponsesRequest::encode(&request, "grok-4.5", &client_key()).expect("request"); + let body = Value::Object(encoded.body().clone()); + + assert_eq!( + body.pointer("/tool_choice/tools/1"), + Some(&json!({"type": "x_search"})) + ); +} + #[test] fn explicit_session_should_not_duplicate_existing_native_cache_tools() { let request = raw_request(json!({