diff --git a/backend/crates/gateway-admin/src/use_case/observability.rs b/backend/crates/gateway-admin/src/use_case/observability.rs index f3437ce15..5c758e21d 100644 --- a/backend/crates/gateway-admin/src/use_case/observability.rs +++ b/backend/crates/gateway-admin/src/use_case/observability.rs @@ -492,14 +492,14 @@ impl DefaultObservabilityService { } } - /// 逐条尽力把 calculated 总额升级为完整分解;单条脏数据(非法 Provider kind - /// 或费用规则失败)只保留该条已存的总额,不影响整页返回。 + /// 逐条尽力把可校验的总额升级为完整分解;单条脏数据(非法 Provider kind、 + /// 不支持的来源或费用规则失败)只保留该条已存的总额,不影响整页返回。 fn enrich_billing(&self, records: &mut [crate::model::observability::UsageRecord]) { for record in records { let Some(UsageBilling::Total { source, total }) = record.billing.as_ref() else { continue; }; - if source != "calculated" { + if !matches!(source.as_str(), "calculated" | "provider_reported") { continue; } let (Some(provider), Some(upstream_model_id)) = ( diff --git a/backend/crates/gateway-admin/tests/use_case/observability.rs b/backend/crates/gateway-admin/tests/use_case/observability.rs index ecf5ffe54..1725cf101 100644 --- a/backend/crates/gateway-admin/tests/use_case/observability.rs +++ b/backend/crates/gateway-admin/tests/use_case/observability.rs @@ -327,9 +327,19 @@ async fn usage_records_should_tolerate_records_that_fail_billing_enrichment() { let store = Arc::new(FixtureObservabilityStore::new(observation_range(now))); let invalid_kind = "x".repeat(65); store.replace_usage_records(vec![ - calculated_total_record("request_invalid_kind", Some(&invalid_kind), now), - calculated_total_record("request_unregistered_kind", Some("anthropic"), now), - calculated_total_record("request_enriched", Some("openai"), now), + total_record( + "request_invalid_kind", + Some(&invalid_kind), + "calculated", + now, + ), + total_record( + "request_unregistered_kind", + Some("anthropic"), + "calculated", + now, + ), + total_record("request_enriched", Some("openai"), "calculated", now), ]); let services = observability_services_with_calculated_billing(store).await; @@ -354,6 +364,30 @@ async fn usage_records_should_tolerate_records_that_fail_billing_enrichment() { ); } +#[tokio::test] +async fn usage_records_should_enrich_provider_reported_totals_when_pricing_matches() { + let now = Utc::now(); + let store = Arc::new(FixtureObservabilityStore::new(observation_range(now))); + store.replace_usage_records(vec![total_record( + "request_provider_reported", + Some("openai"), + "provider_reported", + now, + )]); + let services = observability_services_with_calculated_billing(store).await; + + let page = services + .observability() + .usage_records(usage_query(now)) + .await + .expect("usage records"); + + assert!(matches!( + page.items[0].billing, + Some(UsageBilling::Calculated(_)) + )); +} + struct FixtureObservabilityStore { trend: Mutex>, overview: Mutex, @@ -567,9 +601,10 @@ fn usage_query(now: DateTime) -> UsageQuery { } } -fn calculated_total_record( +fn total_record( id: &str, provider_kind: Option<&str>, + source: &str, now: DateTime, ) -> UsageRecord { UsageRecord { @@ -613,11 +648,11 @@ fn calculated_total_record( image_input_tokens: None, image_output_tokens: None, total_tokens: Some(1_000), - cost_source: "calculated".to_owned(), + cost_source: source.to_owned(), cost_amount: None, cost_currency: None, billing: Some(UsageBilling::Total { - source: "calculated".to_owned(), + source: source.to_owned(), total: CurrencyCost { currency: "USD".to_owned(), amount: gateway_admin::model::observability::DecimalAmount::from_str("1.25") diff --git a/backend/crates/providers/xai/src/transport/request.rs b/backend/crates/providers/xai/src/transport/request.rs index e084a82b8..868ccec79 100644 --- a/backend/crates/providers/xai/src/transport/request.rs +++ b/backend/crates/providers/xai/src/transport/request.rs @@ -56,6 +56,7 @@ const SESSION_FIELDS: &[&str] = &[ "conversation_id", "conversationId", ]; +const PROTOCOL_CONTEXT_SESSION_FIELDS: &[&str] = &["conversation_id", "session_id", "thread_id"]; const FOREIGN_CLIENT_METADATA_FIELDS: &[&str] = &["x-openai-subagent"]; const MAX_SESSION_SEED_BYTES: usize = 1_024; const GROK_CACHE_ROUTE_TOOLS: &[&str] = &["web_search", "x_search"]; @@ -415,6 +416,17 @@ fn explicit_session_seed(request: &GenerateRequest, body: &Map) - .and_then(Value::as_object) .and_then(first_session_value) }) + .or_else(|| { + PROTOCOL_CONTEXT_SESSION_FIELDS.iter().find_map(|field| { + request + .protocol_payload() + .context() + .get(*field) + .and_then(Value::as_str) + .and_then(valid_session_seed) + .map(ToOwned::to_owned) + }) + }) } fn first_session_value(body: &Map) -> Option { diff --git a/backend/crates/providers/xai/tests/provider/contract.rs b/backend/crates/providers/xai/tests/provider/contract.rs index eea9e8c22..ceaa39805 100644 --- a/backend/crates/providers/xai/tests/provider/contract.rs +++ b/backend/crates/providers/xai/tests/provider/contract.rs @@ -149,6 +149,26 @@ fn stateful_sse_with_assistant(encrypted_content: &str, assistant: &str) -> Vec< .into_bytes() } +fn stateful_sse_with_custom_tool_call(encrypted_content: &str, input: &str) -> Vec { + let arguments = serde_json::to_string(&json!({"input": input})).expect("tool arguments"); + let arguments_json = serde_json::to_string(&arguments).expect("arguments JSON string"); + 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\":\"item_exec\",\"type\":\"function_call\",\"status\":\"completed\",\"call_id\":\"call_exec\",\"name\":\"exec\",\"arguments\":{arguments_json}}}}}\n\n", + "event: response.completed\n", + "data: {{\"type\":\"response.completed\",\"response\":{{\"id\":\"resp_state\",\"model\":\"grok-4.5\",\"status\":\"completed\"}}}}\n\n" + ), + encrypted_content, + arguments_json = arguments_json, + ) + .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"); @@ -658,6 +678,23 @@ fn reasoning_replay_operation(session: &str, input: serde_json::Value) -> Operat Operation::Generate(GenerateRequest::from_protocol_payload(payload)) } +fn contextual_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")), + ("input".to_owned(), input), + ("tools".to_owned(), json!([{"type":"custom","name":"exec"}])), + ]), + ) + .expect("OpenAI payload") + .with_context(Map::from_iter([( + "conversation_id".to_owned(), + json!(session), + )])); + Operation::Generate(GenerateRequest::from_protocol_payload(payload)) +} + fn operation_with_reasoning_effort(effort: &str) -> Operation { let payload = ProtocolPayload::json_object( "openai", @@ -1922,6 +1959,55 @@ async fn reasoning_replay_is_reused_only_for_the_same_explicit_session() { ); } +#[tokio::test] +async fn reasoning_replay_should_pair_context_scoped_custom_call_with_its_output() { + let encrypted_content = replay_ciphertext(4); + let transport = StubInferenceTransport::sequence([ + InferenceMode::SuccessBody(stateful_sse_with_custom_tool_call( + &encrypted_content, + "inspect skills", + )), + InferenceMode::Success, + ]); + let provider = provider(StubSelector::success(), transport.clone()).await; + + execute_successfully( + &provider, + contextual_reasoning_replay_operation( + "conversation-from-header", + json!([{"type":"message","role":"user","content":"find skills"}]), + ), + ) + .await; + execute_successfully( + &provider, + contextual_reasoning_replay_operation( + "conversation-from-header", + json!([{ + "type":"custom_tool_call_output", + "call_id":"call_exec", + "output":[{"type":"input_text","text":"skills found"}] + }]), + ), + ) + .await; + + let requests = transport.requests.lock().expect("requests"); + let replayed: serde_json::Value = + serde_json::from_slice(requests[1].body()).expect("replayed body"); + let item_types = replayed["input"] + .as_array() + .expect("replayed input") + .iter() + .filter_map(|item| item["type"].as_str()) + .collect::>(); + + assert_eq!( + item_types, + ["reasoning", "function_call", "function_call_output"] + ); +} + #[tokio::test] async fn reasoning_replay_is_scoped_to_the_selected_account() { let encrypted_content = replay_ciphertext(1);