From 73e96ef89cfc7d938b4706bea625ac98f7ecc8a3 Mon Sep 17 00:00:00 2001 From: wsp Date: Sun, 16 Aug 2026 21:18:01 +0800 Subject: [PATCH 1/5] fix(responses): stabilize prompt cache routing - Add a provider-neutral request context with a stable cache route key - Send prompt_cache_key consistently across tool, retry, and finalize rounds - Parse cache write token usage without conflating missing and zero values - Add content-free request and response cache diagnostics --- Cargo.lock | 1 + src/crates/adapters/ai-adapters/Cargo.toml | 4 +- src/crates/adapters/ai-adapters/src/client.rs | 35 +++- src/crates/adapters/ai-adapters/src/lib.rs | 4 +- .../src/providers/openai/codex_chatgpt.rs | 9 +- .../src/providers/openai/responses.rs | 165 ++++++++++++++++-- .../src/stream/stream_handler/responses.rs | 76 +++++++- .../ai-adapters/src/stream/types/responses.rs | 31 +++- .../adapters/ai-adapters/src/types/config.rs | 2 +- .../tests/common/stream_test_harness.rs | 1 + .../src/agentic/execution/execution_engine.rs | 89 +++++++++- .../src/agentic/execution/round_executor.rs | 4 +- .../core/src/agentic/execution/types.rs | 3 + src/crates/contracts/core-types/src/ai.rs | 12 ++ src/crates/contracts/core-types/src/lib.rs | 21 +-- 15 files changed, 415 insertions(+), 42 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f2e12d2092..85b5bf9648 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -885,6 +885,7 @@ dependencies = [ "eventsource-stream", "fs2", "futures", + "hex", "keyring-core", "libc", "log", diff --git a/src/crates/adapters/ai-adapters/Cargo.toml b/src/crates/adapters/ai-adapters/Cargo.toml index 772f1e5721..dd676b6a14 100644 --- a/src/crates/adapters/ai-adapters/Cargo.toml +++ b/src/crates/adapters/ai-adapters/Cargo.toml @@ -36,7 +36,8 @@ log = { workspace = true } reqwest = { workspace = true, features = ["http2", "json", "rustls", "socks", "stream"] } serde = { workspace = true } serde_json = { workspace = true } -sha2 = { workspace = true, optional = true } +sha2 = { workspace = true } +hex = { workspace = true } tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } tokio-stream = { workspace = true } tokio-util = { workspace = true } @@ -62,7 +63,6 @@ subscription-auth = [ "dep:keyring-core", "dep:libc", "reqwest/form", - "dep:sha2", "tokio/fs", "tokio/io-util", "tokio/net", diff --git a/src/crates/adapters/ai-adapters/src/client.rs b/src/crates/adapters/ai-adapters/src/client.rs index aa1e1feecc..0832ce3cdb 100644 --- a/src/crates/adapters/ai-adapters/src/client.rs +++ b/src/crates/adapters/ai-adapters/src/client.rs @@ -244,8 +244,37 @@ impl AIClient { messages: Vec, tools: Option>, trace: Option, + ) -> Result { + self.send_message_stream_once_with_request_context(messages, tools, None, trace) + .await + } + + /// Open one model stream without an adapter-owned retry loop, carrying + /// provider-neutral request-scoped facts to adapters that support them. + pub async fn send_message_stream_once_with_request_context( + &self, + messages: Vec, + tools: Option>, + request_context: Option, + trace: Option, ) -> Result { let custom_body = self.config.custom_request_body.clone(); + if matches!( + ApiFormat::parse(&self.config.format)?, + ApiFormat::OpenAIResponses + ) { + return openai::responses::send_stream( + self, + messages, + tools, + custom_body, + 1, + trace, + request_context, + ) + .await; + } + self.send_message_stream_with_extra_body_and_max_attempts( messages, tools, @@ -387,8 +416,10 @@ impl AIClient { openai::chat::send_stream(self, messages, tools, extra_body, max_tries, trace).await } ApiFormat::OpenAIResponses => { - openai::responses::send_stream(self, messages, tools, extra_body, max_tries, trace) - .await + openai::responses::send_stream( + self, messages, tools, extra_body, max_tries, trace, None, + ) + .await } ApiFormat::Anthropic => { anthropic::request::send_stream(self, messages, tools, extra_body, max_tries, trace) diff --git a/src/crates/adapters/ai-adapters/src/lib.rs b/src/crates/adapters/ai-adapters/src/lib.rs index 2d4aff353e..987e8b3c21 100644 --- a/src/crates/adapters/ai-adapters/src/lib.rs +++ b/src/crates/adapters/ai-adapters/src/lib.rs @@ -24,6 +24,6 @@ pub use trace::{ }; pub use types::{ resolve_request_url, AIConfig, ConnectionTestMessageCode, ConnectionTestResult, GeminiResponse, - GeminiUsage, Message, ProxyConfig, ReasoningPresetAction, ReasoningPresetDescriptor, - RemoteModelInfo, ToolCall, ToolDefinition, ToolImageAttachment, + GeminiUsage, Message, ModelRequestContext, ProxyConfig, ReasoningPresetAction, + ReasoningPresetDescriptor, RemoteModelInfo, ToolCall, ToolDefinition, ToolImageAttachment, }; diff --git a/src/crates/adapters/ai-adapters/src/providers/openai/codex_chatgpt.rs b/src/crates/adapters/ai-adapters/src/providers/openai/codex_chatgpt.rs index 60b4772a59..5d27afaca4 100644 --- a/src/crates/adapters/ai-adapters/src/providers/openai/codex_chatgpt.rs +++ b/src/crates/adapters/ai-adapters/src/providers/openai/codex_chatgpt.rs @@ -208,7 +208,14 @@ pub(crate) async fn send_stream( trace, || common::apply_headers(client, client.client.post(&url)), move |response, tx, tx_raw, remaining_ttft_timeout| { - handle_responses_stream(response, tx, tx_raw, remaining_ttft_timeout, idle_timeout) + handle_responses_stream( + response, + tx, + tx_raw, + remaining_ttft_timeout, + idle_timeout, + None, + ) }, ) .await diff --git a/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs b/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs index a0a878c6fb..39b1ea6528 100644 --- a/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs +++ b/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs @@ -4,16 +4,73 @@ use crate::client::{AIClient, StreamResponse}; use crate::providers::shared; use crate::stream::handle_responses_stream; use crate::trace::ModelExchangeTraceConfig; -use crate::types::{Message, ReasoningPresetAction, ToolDefinition}; +use crate::types::{Message, ModelRequestContext, ReasoningPresetAction, ToolDefinition}; use anyhow::{anyhow, Result}; use log::debug; +use sha2::{Digest, Sha256}; -pub(crate) fn try_build_request_body( +const TARGET: &str = "ai::responses_stream_request"; + +fn hash_json(value: &serde_json::Value) -> String { + hex::encode(Sha256::digest( + serde_json::to_vec(value).unwrap_or_default(), + )) +} + +fn hash_text(value: Option<&str>) -> String { + hex::encode(Sha256::digest(value.unwrap_or_default().as_bytes())) +} + +fn short_hash(value: &str) -> &str { + value.get(..12).unwrap_or(value) +} + +fn prompt_cache_key_hash(request_body: &serde_json::Value) -> Option { + request_body + .get("prompt_cache_key") + .and_then(serde_json::Value::as_str) + .map(|key| hex::encode(Sha256::digest(key.as_bytes()))) +} + +fn log_prompt_cache_diagnostics(request_body: &serde_json::Value) { + let input = request_body + .get("input") + .cloned() + .unwrap_or_else(|| serde_json::json!([])); + let tools = request_body + .get("tools") + .cloned() + .unwrap_or_else(|| serde_json::json!([])); + let instructions = request_body + .get("instructions") + .and_then(serde_json::Value::as_str); + let cache_key_hash = prompt_cache_key_hash(request_body); + let input_hash = hash_json(&input); + let tools_hash = hash_json(&tools); + let instructions_hash = hash_text(instructions); + + debug!( + target: TARGET, + "Responses prompt cache diagnostics: cache_key_hash={}, instructions_hash={}, input_hash={}, tools_hash={}, input_items={}, tool_count={}", + cache_key_hash + .as_deref() + .map(short_hash) + .unwrap_or("none"), + short_hash(&instructions_hash), + short_hash(&input_hash), + short_hash(&tools_hash), + input.as_array().map(Vec::len).unwrap_or(0), + tools.as_array().map(Vec::len).unwrap_or(0), + ); +} + +fn try_build_request_body_with_context( client: &AIClient, instructions: Option, response_input: Vec, openai_tools: Option>, extra_body: Option, + request_context: Option<&ModelRequestContext>, ) -> Result { let mut request_body = serde_json::json!({ "model": client.config.model, @@ -29,6 +86,14 @@ pub(crate) fn try_build_request_body( request_body["max_output_tokens"] = serde_json::json!(max_tokens); } + if let Some(prompt_cache_route_key) = request_context + .and_then(|context| context.prompt_cache_route_key.as_deref()) + .map(str::trim) + .filter(|key| !key.is_empty()) + { + request_body["prompt_cache_key"] = serde_json::Value::String(prompt_cache_route_key.into()); + } + let base_reasoning_fields = shared::capture_reasoning_fields(&request_body, &["reasoning"], &[]); let protected_keys = &[ @@ -37,6 +102,7 @@ pub(crate) fn try_build_request_body( "instructions", "stream", "max_output_tokens", + "prompt_cache_key", "tools", ]; let compile = |action: &ReasoningPresetAction, body: &mut serde_json::Value| -> Result { @@ -69,6 +135,7 @@ pub(crate) fn try_build_request_body( "instructions", "stream", "max_output_tokens", + "prompt_cache_key", ], &[], ); @@ -76,11 +143,18 @@ pub(crate) fn try_build_request_body( if let Some(extra) = extra_body { if let Some(extra_obj) = extra.as_object() { shared::merge_extra_body(&mut request_body, extra_obj); - shared::log_extra_body_keys("ai::responses_stream_request", extra_obj); + shared::log_extra_body_keys(TARGET, extra_obj); } } shared::restore_protected_body(&mut request_body, protected_body); + if let Some(prompt_cache_route_key) = request_context + .and_then(|context| context.prompt_cache_route_key.as_deref()) + .map(str::trim) + .filter(|key| !key.is_empty()) + { + request_body["prompt_cache_key"] = serde_json::Value::String(prompt_cache_route_key.into()); + } if let Some(preset) = client.selected_reasoning_preset.as_ref() { shared::reset_reasoning_fields( &mut request_body, @@ -92,20 +166,34 @@ pub(crate) fn try_build_request_body( } shared::log_request_body( - "ai::responses_stream_request", + TARGET, "Responses stream request body (excluding tools):", &request_body, ); - common::attach_tools( - &mut request_body, - openai_tools, - "ai::responses_stream_request", - ); + common::attach_tools(&mut request_body, openai_tools, TARGET); + log_prompt_cache_diagnostics(&request_body); Ok(request_body) } +pub(crate) fn try_build_request_body( + client: &AIClient, + instructions: Option, + response_input: Vec, + openai_tools: Option>, + extra_body: Option, +) -> Result { + try_build_request_body_with_context( + client, + instructions, + response_input, + openai_tools, + extra_body, + None, + ) +} + #[cfg(test)] pub(crate) fn build_request_body( client: &AIClient, @@ -124,6 +212,26 @@ pub(crate) fn build_request_body( .expect("request body should compile") } +#[cfg(test)] +fn build_request_body_with_context( + client: &AIClient, + instructions: Option, + response_input: Vec, + openai_tools: Option>, + extra_body: Option, + request_context: Option<&ModelRequestContext>, +) -> serde_json::Value { + try_build_request_body_with_context( + client, + instructions, + response_input, + openai_tools, + extra_body, + request_context, + ) + .expect("request body should compile") +} + pub(crate) async fn send_stream( client: &AIClient, messages: Vec, @@ -131,6 +239,7 @@ pub(crate) async fn send_stream( extra_body: Option, max_tries: usize, trace: Option, + request_context: Option, ) -> Result { // Codex CLI's ChatGPT-login backend (`chatgpt.com/backend-api/codex`) // speaks a constrained Responses dialect with several extra @@ -153,13 +262,15 @@ pub(crate) async fn send_stream( let (instructions, response_input) = OpenAIMessageConverter::convert_messages_to_responses_input(messages); let openai_tools = common::convert_tools_flat(tools); - let request_body = try_build_request_body( + let request_body = try_build_request_body_with_context( client, instructions, response_input, openai_tools, extra_body, + request_context.as_ref(), )?; + let expected_prompt_cache_key_hash = prompt_cache_key_hash(&request_body); let idle_timeout = client.stream_options.idle_timeout; let ttft_timeout = client.stream_options.ttft_timeout; @@ -172,7 +283,14 @@ pub(crate) async fn send_stream( trace, || common::apply_headers(client, client.client.post(&url)), move |response, tx, tx_raw, remaining_ttft_timeout| { - handle_responses_stream(response, tx, tx_raw, remaining_ttft_timeout, idle_timeout) + handle_responses_stream( + response, + tx, + tx_raw, + remaining_ttft_timeout, + idle_timeout, + expected_prompt_cache_key_hash.clone(), + ) }, ) .await @@ -180,8 +298,8 @@ pub(crate) async fn send_stream( #[cfg(test)] mod tests { - use super::build_request_body; - use crate::types::ToolDefinition; + use super::{build_request_body, build_request_body_with_context}; + use crate::types::{ModelRequestContext, ToolDefinition}; use crate::{client::AIClient, types::AIConfig}; use serde_json::json; @@ -234,4 +352,25 @@ mod tests { assert_eq!(request_body["tools"][0]["type"], json!("function")); assert!(request_body["tools"][0].get("function").is_none()); } + + #[test] + fn attaches_runtime_prompt_cache_key_after_custom_body_merge() { + let client = test_client(); + let request_context = ModelRequestContext { + prompt_cache_route_key: Some("bitfun-pc-v1-stable".to_string()), + }; + let request_body = build_request_body_with_context( + &client, + Some("stable instructions".to_string()), + Vec::new(), + None, + Some(json!({ "prompt_cache_key": "user-override" })), + Some(&request_context), + ); + + assert_eq!( + request_body["prompt_cache_key"], + json!("bitfun-pc-v1-stable") + ); + } } diff --git a/src/crates/adapters/ai-adapters/src/stream/stream_handler/responses.rs b/src/crates/adapters/ai-adapters/src/stream/stream_handler/responses.rs index ea6e4382ef..98e620be0c 100644 --- a/src/crates/adapters/ai-adapters/src/stream/stream_handler/responses.rs +++ b/src/crates/adapters/ai-adapters/src/stream/stream_handler/responses.rs @@ -8,15 +8,46 @@ use anyhow::{anyhow, Result}; use bitfun_agent_stream::ToolCallCompletion; use bitfun_core_types::errors::AiProviderError; use eventsource_stream::Eventsource; -use log::{error, trace}; +use log::{debug, error, trace}; use reqwest::Response; use serde_json::Value; +use sha2::{Digest, Sha256}; use std::collections::HashMap; use std::time::Duration; use tokio::sync::mpsc; const AI_STREAM_RESPONSE_TARGET: &str = "ai::responses_stream_response"; +fn extract_response_prompt_cache_key_hash(response: Option<&Value>) -> Option { + response + .and_then(|response| response.get("prompt_cache_key")) + .and_then(Value::as_str) + .map(|cache_key| hex::encode(Sha256::digest(cache_key.as_bytes()))) +} + +fn log_cache_stream_diagnostics( + response_created_count: usize, + completed_response_id: Option<&str>, + expected_prompt_cache_key_hash: Option<&str>, + response_prompt_cache_key_hash: Option<&str>, +) { + debug!( + target: "ai::responses_cache", + "Responses cache stream diagnostics: response_created_count={}, completed_response_id={}, request_cache_key_hash={}, response_cache_key_hash={}, cache_key_matches={:?}", + response_created_count, + completed_response_id.unwrap_or("none"), + expected_prompt_cache_key_hash + .and_then(|value| value.get(..12)) + .unwrap_or("none"), + response_prompt_cache_key_hash + .and_then(|value| value.get(..12)) + .unwrap_or("none"), + expected_prompt_cache_key_hash + .zip(response_prompt_cache_key_hash) + .map(|(expected, actual)| expected == actual), + ); +} + #[derive(Debug, Default, Clone)] struct InProgressToolCall { call_id: Option, @@ -291,6 +322,7 @@ pub async fn handle_responses_stream( tx_raw_sse: Option>, ttft_timeout: Option, idle_timeout: Option, + expected_prompt_cache_key_hash: Option, ) { let mut stream = response.bytes_stream().eventsource(); // Some providers close the stream after emitting the terminal event and may not send `[DONE]`. @@ -301,6 +333,8 @@ pub async fn handle_responses_stream( let mut tool_call_index_by_id: HashMap = HashMap::new(); let mut stats = StreamStats::new("Responses"); let mut timeout_controller = StreamTimeoutController::new(ttft_timeout, idle_timeout); + let mut response_created_count = 0usize; + let mut response_prompt_cache_key_hash: Option = None; loop { let sse = match next_stream_item(&mut stream, &timeout_controller).await { @@ -394,6 +428,24 @@ pub async fn handle_responses_stream( stats.increment(format!("event:{}", event.kind)); match event.kind.as_str() { + "response.created" => { + response_created_count += 1; + if let Some(response) = event.response.as_ref() { + if let Some(response_id) = response.get("id").and_then(Value::as_str) { + debug!( + target: "ai::responses_cache", + "Responses response.created observed: count={}, response_id={}", + response_created_count, + response_id + ); + } + if let Some(cache_key_hash) = + extract_response_prompt_cache_key_hash(Some(response)) + { + response_prompt_cache_key_hash = Some(cache_key_hash); + } + } + } "response.output_item.added" => { // Track tool calls so we can stream arguments via `response.function_call_arguments.delta`. if let Some(item) = event.item.as_ref() { @@ -503,6 +555,11 @@ pub async fn handle_responses_stream( } // Best-effort: use the final response object to fill any missing tool-call argument tail. if let Some(response_val) = event.response.as_ref() { + if let Some(cache_key_hash) = + extract_response_prompt_cache_key_hash(Some(response_val)) + { + response_prompt_cache_key_hash = Some(cache_key_hash); + } if let Some(output) = response_val.get("output").and_then(Value::as_array) { for (idx, item) in output.iter().enumerate() { if item.get("type").and_then(Value::as_str) != Some("function_call") { @@ -571,6 +628,12 @@ pub async fn handle_responses_stream( &mut stats, unified_response, ); + log_cache_stream_diagnostics( + response_created_count, + Some(response.id.as_str()), + expected_prompt_cache_key_hash.as_deref(), + response_prompt_cache_key_hash.as_deref(), + ); continue; } Some(Err(e)) => { @@ -605,6 +668,11 @@ pub async fn handle_responses_stream( if received_finish_reason { continue; } + if let Some(cache_key_hash) = + extract_response_prompt_cache_key_hash(event.response.as_ref()) + { + response_prompt_cache_key_hash = Some(cache_key_hash); + } match event.response.map(serde_json::from_value::) { Some(Ok(response)) => { received_finish_reason = true; @@ -622,6 +690,12 @@ pub async fn handle_responses_stream( &mut stats, unified_response, ); + log_cache_stream_diagnostics( + response_created_count, + response.id.as_deref(), + expected_prompt_cache_key_hash.as_deref(), + response_prompt_cache_key_hash.as_deref(), + ); continue; } Some(Err(e)) => { diff --git a/src/crates/adapters/ai-adapters/src/stream/types/responses.rs b/src/crates/adapters/ai-adapters/src/stream/types/responses.rs index 8a17e4b607..38b81e6897 100644 --- a/src/crates/adapters/ai-adapters/src/stream/types/responses.rs +++ b/src/crates/adapters/ai-adapters/src/stream/types/responses.rs @@ -32,7 +32,6 @@ pub struct ResponsesCompleted { #[derive(Debug, Deserialize)] pub struct ResponsesDone { #[serde(default)] - #[allow(dead_code)] pub id: Option, #[serde(default)] pub usage: Option, @@ -50,19 +49,23 @@ pub struct ResponsesUsage { #[derive(Debug, Deserialize)] pub struct ResponsesInputTokensDetails { pub cached_tokens: u32, + #[serde(default)] + pub cache_write_tokens: Option, } impl From for UnifiedTokenUsage { fn from(usage: ResponsesUsage) -> Self { + let (cached_content_token_count, cache_creation_token_count) = usage + .input_tokens_details + .map(|details| (Some(details.cached_tokens), details.cache_write_tokens)) + .unwrap_or((None, None)); Self { prompt_token_count: usage.input_tokens, candidates_token_count: usage.output_tokens, total_token_count: usage.total_tokens, reasoning_token_count: None, - cached_content_token_count: usage - .input_tokens_details - .map(|details| details.cached_tokens), - cache_creation_token_count: None, + cached_content_token_count, + cache_creation_token_count, } } } @@ -141,14 +144,14 @@ mod tests { fn responses_cached_tokens_maps_to_cached_content() { let raw = r#"{ "input_tokens": 200, - "input_tokens_details": { "cached_tokens": 80 }, + "input_tokens_details": { "cached_tokens": 80, "cache_write_tokens": 64 }, "output_tokens": 40, "total_tokens": 240 }"#; let usage: ResponsesUsage = serde_json::from_str(raw).expect("valid responses usage"); let unified: UnifiedTokenUsage = usage.into(); assert_eq!(unified.cached_content_token_count, Some(80)); - assert_eq!(unified.cache_creation_token_count, None); + assert_eq!(unified.cache_creation_token_count, Some(64)); } #[test] @@ -160,6 +163,20 @@ mod tests { assert_eq!(unified.cache_creation_token_count, None); } + #[test] + fn responses_explicit_zero_cache_write_is_preserved() { + let raw = r#"{ + "input_tokens": 200, + "input_tokens_details": { "cached_tokens": 80, "cache_write_tokens": 0 }, + "output_tokens": 40, + "total_tokens": 240 + }"#; + let usage: ResponsesUsage = serde_json::from_str(raw).expect("valid responses usage"); + let unified: UnifiedTokenUsage = usage.into(); + assert_eq!(unified.cached_content_token_count, Some(80)); + assert_eq!(unified.cache_creation_token_count, Some(0)); + } + #[test] fn parses_output_text_message_item() { let response = parse_responses_output_item( diff --git a/src/crates/adapters/ai-adapters/src/types/config.rs b/src/crates/adapters/ai-adapters/src/types/config.rs index ff45848cd7..66be61cf8c 100644 --- a/src/crates/adapters/ai-adapters/src/types/config.rs +++ b/src/crates/adapters/ai-adapters/src/types/config.rs @@ -1,5 +1,5 @@ pub use bitfun_core_types::{ - AIConfig, ProxyConfig, ReasoningPresetAction, ReasoningPresetDescriptor, + AIConfig, ModelRequestContext, ProxyConfig, ReasoningPresetAction, ReasoningPresetDescriptor, }; fn append_endpoint(base_url: &str, endpoint: &str) -> String { diff --git a/src/crates/adapters/ai-adapters/tests/common/stream_test_harness.rs b/src/crates/adapters/ai-adapters/tests/common/stream_test_harness.rs index 4f6c1a1ea8..025f963f2d 100644 --- a/src/crates/adapters/ai-adapters/tests/common/stream_test_harness.rs +++ b/src/crates/adapters/ai-adapters/tests/common/stream_test_harness.rs @@ -171,6 +171,7 @@ pub(crate) async fn run_stream_fixture_with_options( Some(tx_raw_sse), options.ttft_timeout, options.idle_timeout, + None, )); } } diff --git a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs index d9999eedbf..c62e155634 100644 --- a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs +++ b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs @@ -66,10 +66,11 @@ use crate::util::types::ToolDefinition; use crate::util::{elapsed_ms_u64, truncate_at_char_boundary}; use bitfun_agent_runtime::output_surface::TOOL_CONTEXT_INLINE_MARKDOWN_IMAGE_DISPLAY_KEY; use bitfun_agent_runtime::permission::PERMISSION_MODE_CONTEXT_KEY; +use bitfun_agent_runtime::prompt_cache::prompt_cache_scope_key; use bitfun_agent_runtime::remote_file_delivery::TOOL_CONTEXT_REMOTE_FILE_DELIVERY_KEY; use bitfun_agent_runtime::thread_goal_tools::ensure_thread_goal_tools; use bitfun_ai_adapters::ModelExchangeTraceConfig; -use bitfun_core_types::SessionModelBindingPolicy; +use bitfun_core_types::{ModelRequestContext, SessionModelBindingPolicy}; use bitfun_runtime_ports::{resolve_permission_mode, PermissionMode, PermissionModeLayers}; use dashmap::DashMap; use log::{debug, error, info, trace, warn}; @@ -494,6 +495,7 @@ struct FinalizeRoundInput<'a> { messages: &'a [Message], prepended_reminders: &'a [&'a str], primary_model_facts: &'a PrimaryModelFacts, + model_request_context: &'a ModelRequestContext, execution_context_vars: &'a HashMap, round_group_id: Option, round_number: usize, @@ -524,6 +526,7 @@ pub struct ExecutionEngine { } impl ExecutionEngine { + const PROVIDER_PROMPT_CACHE_ROUTE_SCHEMA: &'static str = "bitfun-provider-prompt-cache-v1"; const AUTO_COMPRESSION_SAFETY_RESERVE_TOKENS: usize = 10_000; const MAX_COMPRESSION_OVERFLOW_ATTEMPTS: usize = 4; const MAX_MAIN_CONTEXT_OVERFLOW_RECOVERIES: usize = 2; @@ -534,6 +537,49 @@ impl ExecutionEngine { const FINALIZE_USER_FOLLOWUP: &'static str = "Provide a final answer. You MUST not call any tools."; + fn model_request_context( + session_id: &str, + model_binding_fingerprint: &str, + effective_model_name: &str, + current_agent: &dyn crate::agentic::agents::Agent, + ) -> ModelRequestContext { + let prompt_scope = prompt_cache_scope_key( + ¤t_agent.system_prompt_cache_identity(Some(effective_model_name)), + ¤t_agent.user_context_cache_identity(), + ); + Self::model_request_context_from_scope_key( + session_id, + model_binding_fingerprint, + effective_model_name, + &prompt_scope, + ) + } + + fn model_request_context_from_scope_key( + session_id: &str, + model_binding_fingerprint: &str, + effective_model_name: &str, + prompt_scope: &str, + ) -> ModelRequestContext { + let mut hasher = Sha256::new(); + for component in [ + Self::PROVIDER_PROMPT_CACHE_ROUTE_SCHEMA, + session_id, + model_binding_fingerprint, + effective_model_name, + prompt_scope, + ] { + hasher.update(component.as_bytes()); + hasher.update([0]); + } + ModelRequestContext { + prompt_cache_route_key: Some(format!( + "bitfun-pc-v1-{}", + hex::encode(hasher.finalize()) + )), + } + } + async fn context_vars_for_round( &self, base: &HashMap, @@ -1946,6 +1992,7 @@ impl ExecutionEngine { loaded_deferred_tool_specs: Vec::new(), model_config_id: input.primary_model_facts.model_id.clone(), effective_model_name: input.ai_client.config.model.clone(), + model_request_context: input.model_request_context.clone(), primary_model_facts: input.primary_model_facts.clone(), agent_type: input.agent_type, context_vars: round_context_vars, @@ -3390,7 +3437,7 @@ impl ExecutionEngine { } } - let (model_id, _) = self + let (model_id, model_binding_fingerprint) = self .resolve_model_id_for_turn( &session, &agent_type, @@ -3468,6 +3515,12 @@ impl ExecutionEngine { }; Self::validate_frozen_model_contract(&context).await?; Self::validate_frozen_reasoning_contract(&context, ai_client.as_ref())?; + let model_request_context = Self::model_request_context( + &context.session_id, + &model_binding_fingerprint, + &ai_client.config.model, + current_agent.as_ref(), + ); // Primary model vision capability (tools + system prompt appendix; also used below for API message stripping). let primary_model_facts = Self::resolve_primary_model_context( @@ -4029,6 +4082,7 @@ impl ExecutionEngine { loaded_deferred_tool_specs, model_config_id: model_id.clone(), effective_model_name: ai_client.config.model.clone(), + model_request_context: model_request_context.clone(), primary_model_facts: primary_model_facts.clone(), agent_type: agent_type.clone(), context_vars: round_context_vars, @@ -4806,6 +4860,7 @@ impl ExecutionEngine { round_group_id: finalize_round_group_id.clone(), execution_context_vars: &execution_context_vars, primary_model_facts: &primary_model_facts, + model_request_context: &model_request_context, prepended_reminders: &finalize_prepended_reminders, messages: &messages, reminder_text: finalize_reminder, @@ -4837,6 +4892,7 @@ impl ExecutionEngine { round_group_id: finalize_round_group_id.clone(), execution_context_vars: &execution_context_vars, primary_model_facts: &primary_model_facts, + model_request_context: &model_request_context, prepended_reminders: &finalize_prepended_reminders, messages: &messages, reminder_text: finalize_reminder, @@ -6398,6 +6454,35 @@ mod tests { assert_eq!(snapshot.compression_failure_count, 2); } + #[test] + fn provider_prompt_cache_route_key_is_stable_and_changes_with_scope() { + let first = ExecutionEngine::model_request_context_from_scope_key( + "session-1", + "binding-1", + "gpt-5.6-terra", + "scope-1", + ); + let retry = ExecutionEngine::model_request_context_from_scope_key( + "session-1", + "binding-1", + "gpt-5.6-terra", + "scope-1", + ); + let changed = ExecutionEngine::model_request_context_from_scope_key( + "session-1", + "binding-2", + "gpt-5.6-terra", + "scope-1", + ); + + assert_eq!(first, retry); + assert_ne!(first, changed); + assert!(first + .prompt_cache_route_key + .as_deref() + .is_some_and(|key| key.starts_with("bitfun-pc-v1-"))); + } + fn command_result(tool_name: &str, success: bool, exit_code: Option) -> Message { Message::tool_result(ToolResult { tool_id: format!("{}-tool", tool_name), diff --git a/src/crates/assembly/core/src/agentic/execution/round_executor.rs b/src/crates/assembly/core/src/agentic/execution/round_executor.rs index 5818ca3700..9967c0c54a 100644 --- a/src/crates/assembly/core/src/agentic/execution/round_executor.rs +++ b/src/crates/assembly/core/src/agentic/execution/round_executor.rs @@ -394,9 +394,10 @@ impl RoundExecutor { let request_trace_config = trace_config .clone() .map(|config| config.with_round_attempt(attempt_id.clone(), attempt_number)); - let send_future = ai_client.send_message_stream_once( + let send_future = ai_client.send_message_stream_once_with_request_context( ai_messages.clone(), tool_definitions.clone(), + Some(context.model_request_context.clone()), request_trace_config, ); let send_result = tokio::select! { @@ -1699,6 +1700,7 @@ mod tests { loaded_deferred_tool_specs: Vec::new(), model_config_id: "model-1".to_string(), effective_model_name: "model-1".to_string(), + model_request_context: Default::default(), primary_model_facts: tool_runtime::context::PrimaryModelFacts::new( "model-1", "model-1", "openai", true, ), diff --git a/src/crates/assembly/core/src/agentic/execution/types.rs b/src/crates/assembly/core/src/agentic/execution/types.rs index 89820ce545..4732c591ed 100644 --- a/src/crates/assembly/core/src/agentic/execution/types.rs +++ b/src/crates/assembly/core/src/agentic/execution/types.rs @@ -8,6 +8,7 @@ use crate::agentic::workspace::WorkspaceServices; use crate::agentic::WorkspaceBinding; pub use bitfun_agent_runtime::events::FinishReason; use bitfun_agent_tools::LoadedDeferredToolSpec; +use bitfun_core_types::ModelRequestContext; use bitfun_runtime_ports::{ DelegationPolicy, PermissionConstraintLayer, PermissionDelegationContext, PermissionRuntimeCeiling, RemoteExecPort, TerminalPort, @@ -71,6 +72,8 @@ pub struct RoundContext { pub model_config_id: String, /// Provider model name sent in the request. pub effective_model_name: String, + /// Provider-neutral request-scoped facts resolved by the runtime owner. + pub model_request_context: ModelRequestContext, pub primary_model_facts: PrimaryModelFacts, pub agent_type: String, pub context_vars: HashMap, diff --git a/src/crates/contracts/core-types/src/ai.rs b/src/crates/contracts/core-types/src/ai.rs index 6406df7e63..304215364d 100644 --- a/src/crates/contracts/core-types/src/ai.rs +++ b/src/crates/contracts/core-types/src/ai.rs @@ -643,6 +643,18 @@ pub struct AIConfig { pub custom_request_body_mode: Option, } +/// Provider-neutral options that vary per model request rather than per model +/// configuration. +/// +/// Adapters may map these opaque facts to provider-specific request fields. +/// Runtime owners must not place provider field names or raw local identifiers +/// in this contract. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct ModelRequestContext { + /// Stable, opaque routing identity for provider-side prompt-prefix caches. + pub prompt_cache_route_key: Option, +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ToolCall { pub id: String, diff --git a/src/crates/contracts/core-types/src/lib.rs b/src/crates/contracts/core-types/src/lib.rs index 9dd59e4de5..9e0cfe79ce 100644 --- a/src/crates/contracts/core-types/src/lib.rs +++ b/src/crates/contracts/core-types/src/lib.rs @@ -15,16 +15,17 @@ pub mod tool_image_attachment; pub mod worktree; pub use ai::{ - AIConfig, ConnectionTestMessageCode, ConnectionTestResult, Message, ModelsDevCatalogSource, - ModelsDevCatalogStatus, ModelsDevReasoningCatalog, ModelsDevReasoningModel, - ModelsDevReasoningProvider, ModelsDevRefreshResult, ModelsDevRefreshStatus, ProviderCatalog, - ProviderCatalogEndpoint, ProviderCatalogModel, ProviderCatalogModelCapabilities, - ProviderCatalogModelLimits, ProviderCatalogModelPricing, ProviderCatalogModelSource, - ProviderCatalogProvider, ProviderCatalogSource, ProviderCatalogUpstreamProvider, ProxyConfig, - ReasoningCapabilityStatus, ReasoningCatalogBinding, ReasoningCatalogProjection, - ReasoningCatalogProjectionRequest, ReasoningConfig, ReasoningPreset, ReasoningPresetAction, - ReasoningPresetDescriptor, ReasoningPresetSource, RemoteModelInfo, ToolCall, - ToolCallConfirmationDetails, ToolCallRequestInfo, ToolCallResponseInfo, ToolDefinition, + AIConfig, ConnectionTestMessageCode, ConnectionTestResult, Message, ModelRequestContext, + ModelsDevCatalogSource, ModelsDevCatalogStatus, ModelsDevReasoningCatalog, + ModelsDevReasoningModel, ModelsDevReasoningProvider, ModelsDevRefreshResult, + ModelsDevRefreshStatus, ProviderCatalog, ProviderCatalogEndpoint, ProviderCatalogModel, + ProviderCatalogModelCapabilities, ProviderCatalogModelLimits, ProviderCatalogModelPricing, + ProviderCatalogModelSource, ProviderCatalogProvider, ProviderCatalogSource, + ProviderCatalogUpstreamProvider, ProxyConfig, ReasoningCapabilityStatus, + ReasoningCatalogBinding, ReasoningCatalogProjection, ReasoningCatalogProjectionRequest, + ReasoningConfig, ReasoningPreset, ReasoningPresetAction, ReasoningPresetDescriptor, + ReasoningPresetSource, RemoteModelInfo, ToolCall, ToolCallConfirmationDetails, + ToolCallRequestInfo, ToolCallResponseInfo, ToolDefinition, }; pub use errors::{AiErrorDetail, ErrorCategory}; pub use model::{ From 793d0842e110b5415c27ab9408c4e82327000f73 Mon Sep 17 00:00:00 2001 From: wsp Date: Sun, 16 Aug 2026 22:49:53 +0800 Subject: [PATCH 2/5] fix(responses): replay encrypted reasoning state Capture and persist the authoritative Responses output layout, including opaque reasoning state, for subsequent agent rounds. - Replay reasoning, messages, and function calls in original order - Bind replay state to the resolved model runtime fingerprint - Fall back atomically for incompatible fingerprints or layouts - Preserve legacy deserialization and non-Responses behavior - Cover capture, persistence, ordering, and fallback paths This keeps multi-round prompts prefix-stable so provider-side prompt caching can reuse prior reasoning and tool history. --- .../ai-adapters/src/client/healthcheck.rs | 1 + .../src/client/response_aggregator.rs | 1 + .../providers/anthropic/message_converter.rs | 2 + .../src/providers/gemini/message_converter.rs | 4 + .../src/providers/openai/codex_chatgpt.rs | 11 +- .../src/providers/openai/message_converter.rs | 298 ++++++++++++++++++ .../src/providers/openai/responses.rs | 34 +- .../src/stream/stream_handler/responses.rs | 168 +++++++++- .../ai-adapters/src/stream/types/anthropic.rs | 1 + .../ai-adapters/src/stream/types/gemini.rs | 6 + .../ai-adapters/src/stream/types/openai.rs | 4 + .../ai-adapters/src/stream/types/responses.rs | 151 ++++++++- .../stream_replay_regressions.rs | 1 + .../assembly/core/src/agentic/core/message.rs | 80 +++++ .../src/agentic/execution/execution_engine.rs | 9 + .../src/agentic/execution/round_executor.rs | 27 +- .../src/agentic/execution/stream_processor.rs | 6 +- .../image_analysis/image_processing.rs | 4 + .../src/agentic/session/session_manager.rs | 2 + src/crates/contracts/core-types/src/ai.rs | 40 +++ src/crates/contracts/core-types/src/lib.rs | 22 +- src/crates/execution/agent-stream/src/lib.rs | 68 +++- .../execution/agent-stream/src/unified.rs | 15 + 23 files changed, 925 insertions(+), 30 deletions(-) diff --git a/src/crates/adapters/ai-adapters/src/client/healthcheck.rs b/src/crates/adapters/ai-adapters/src/client/healthcheck.rs index f298baeeec..0663066ee1 100644 --- a/src/crates/adapters/ai-adapters/src/client/healthcheck.rs +++ b/src/crates/adapters/ai-adapters/src/client/healthcheck.rs @@ -219,6 +219,7 @@ pub(crate) async fn test_image_input_connection( name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, }]; match client diff --git a/src/crates/adapters/ai-adapters/src/client/response_aggregator.rs b/src/crates/adapters/ai-adapters/src/client/response_aggregator.rs index 3157abcdec..135a8b2a13 100644 --- a/src/crates/adapters/ai-adapters/src/client/response_aggregator.rs +++ b/src/crates/adapters/ai-adapters/src/client/response_aggregator.rs @@ -33,6 +33,7 @@ pub(crate) async fn aggregate_stream_response( tool_call_completion: _, finish_reason: chunk_finish_reason, provider_metadata: chunk_provider_metadata, + model_response_replay: _, } = chunk; if let Some(text) = text { diff --git a/src/crates/adapters/ai-adapters/src/providers/anthropic/message_converter.rs b/src/crates/adapters/ai-adapters/src/providers/anthropic/message_converter.rs index 618a12db6a..aa1a875220 100644 --- a/src/crates/adapters/ai-adapters/src/providers/anthropic/message_converter.rs +++ b/src/crates/adapters/ai-adapters/src/providers/anthropic/message_converter.rs @@ -246,6 +246,7 @@ mod tests { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, }; let (_, messages) = AnthropicMessageConverter::convert_messages(vec![msg]); @@ -288,6 +289,7 @@ mod tests { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, }; let (_, messages) = AnthropicMessageConverter::convert_messages(vec![ diff --git a/src/crates/adapters/ai-adapters/src/providers/gemini/message_converter.rs b/src/crates/adapters/ai-adapters/src/providers/gemini/message_converter.rs index a48118fe13..3dd47b5876 100644 --- a/src/crates/adapters/ai-adapters/src/providers/gemini/message_converter.rs +++ b/src/crates/adapters/ai-adapters/src/providers/gemini/message_converter.rs @@ -673,6 +673,7 @@ mod tests { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, }, Message { role: "tool".to_string(), @@ -684,6 +685,7 @@ mod tests { name: Some("get_weather".to_string()), is_error: None, tool_image_attachments: None, + model_response_replay: None, }, ]; @@ -726,6 +728,7 @@ mod tests { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, }]; let (_, contents) = @@ -764,6 +767,7 @@ mod tests { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, }]; let (_, contents) = GeminiMessageConverter::convert_messages(messages, "gemini-2.5-pro"); diff --git a/src/crates/adapters/ai-adapters/src/providers/openai/codex_chatgpt.rs b/src/crates/adapters/ai-adapters/src/providers/openai/codex_chatgpt.rs index 5d27afaca4..bc6c6e85b5 100644 --- a/src/crates/adapters/ai-adapters/src/providers/openai/codex_chatgpt.rs +++ b/src/crates/adapters/ai-adapters/src/providers/openai/codex_chatgpt.rs @@ -24,7 +24,7 @@ use crate::client::{AIClient, StreamResponse}; use crate::providers::shared; use crate::stream::handle_responses_stream; use crate::trace::ModelExchangeTraceConfig; -use crate::types::{Message, ReasoningPresetAction, ToolDefinition}; +use crate::types::{Message, ModelRequestContext, ReasoningPresetAction, ToolDefinition}; use anyhow::Result; use log::debug; use serde_json::{json, Value}; @@ -184,6 +184,7 @@ pub(crate) async fn send_stream( extra_body: Option, max_tries: usize, trace: Option, + request_context: Option, ) -> Result { let url = client.config.request_url.clone(); debug!( @@ -191,8 +192,14 @@ pub(crate) async fn send_stream( client.config.model, url, max_tries ); + let model_binding_fingerprint = request_context + .as_ref() + .and_then(|context| context.model_binding_fingerprint.as_deref()); let (instructions, response_input) = - OpenAIMessageConverter::convert_messages_to_responses_input(messages); + OpenAIMessageConverter::convert_messages_to_responses_input_with_context( + messages, + model_binding_fingerprint, + ); let tools_flat = common::convert_tools_flat(tools); let request_body = try_build_request_body(client, instructions, response_input, tools_flat, extra_body)?; diff --git a/src/crates/adapters/ai-adapters/src/providers/openai/message_converter.rs b/src/crates/adapters/ai-adapters/src/providers/openai/message_converter.rs index e1d9554007..ccf6fff9e5 100644 --- a/src/crates/adapters/ai-adapters/src/providers/openai/message_converter.rs +++ b/src/crates/adapters/ai-adapters/src/providers/openai/message_converter.rs @@ -1,14 +1,24 @@ //! OpenAI message format converter +use crate::stream::types::responses::OPENAI_RESPONSES_REPLAY_PROTOCOL; use crate::types::{Message, ToolDefinition}; +use bitfun_core_types::ModelResponseReplayItem; use log::{error, warn}; use serde_json::{json, Value}; +use std::collections::{HashMap, HashSet}; pub struct OpenAIMessageConverter; impl OpenAIMessageConverter { pub fn convert_messages_to_responses_input( messages: Vec, + ) -> (Option, Vec) { + Self::convert_messages_to_responses_input_with_context(messages, None) + } + + pub fn convert_messages_to_responses_input_with_context( + messages: Vec, + model_binding_fingerprint: Option<&str>, ) -> (Option, Vec) { let mut instructions = Vec::new(); let mut input = Vec::new(); @@ -27,6 +37,13 @@ impl OpenAIMessageConverter { } } "assistant" => { + if let Some(replay_items) = + Self::convert_assistant_replay_items(&msg, model_binding_fingerprint) + { + input.extend(replay_items); + continue; + } + if let Some(content_items) = Self::convert_message_content_to_responses_items( &msg.role, msg.content.as_deref(), @@ -74,6 +91,105 @@ impl OpenAIMessageConverter { (instructions, input) } + fn convert_assistant_replay_items( + msg: &Message, + model_binding_fingerprint: Option<&str>, + ) -> Option> { + let replay = msg.model_response_replay.as_ref()?; + let current_fingerprint = model_binding_fingerprint.filter(|value| !value.is_empty())?; + if replay.protocol != OPENAI_RESPONSES_REPLAY_PROTOCOL + || replay.model_binding_fingerprint != current_fingerprint + || replay.items.is_empty() + { + return None; + } + + let assistant_message = + Self::convert_message_content_to_responses_items("assistant", msg.content.as_deref()) + .map(|content| { + json!({ + "type": "message", + "role": "assistant", + "content": content, + }) + }); + + let tool_calls = msg.tool_calls.as_deref().unwrap_or_default(); + let mut tool_calls_by_id = HashMap::with_capacity(tool_calls.len()); + for tool_call in tool_calls { + if tool_call.id.is_empty() + || tool_calls_by_id + .insert(tool_call.id.as_str(), tool_call) + .is_some() + { + return None; + } + } + + let mut output = Vec::with_capacity(replay.items.len()); + let mut assistant_message_used = false; + let mut used_tool_calls = HashSet::with_capacity(tool_calls.len()); + let mut saw_opaque_reasoning = false; + + for replay_item in &replay.items { + match replay_item { + ModelResponseReplayItem::OpaqueReasoning { + item_id, + summary, + opaque_state, + } => { + if opaque_state.is_empty() { + return None; + } + saw_opaque_reasoning = true; + let mut item = json!({ + "type": "reasoning", + "summary": summary + .iter() + .map(|part| json!({ + "type": "summary_text", + "text": part.text, + })) + .collect::>(), + "encrypted_content": opaque_state, + }); + if let Some(item_id) = item_id.as_ref().filter(|value| !value.is_empty()) { + item["id"] = Value::String(item_id.clone()); + } + output.push(item); + } + ModelResponseReplayItem::AssistantMessage => { + if assistant_message_used { + return None; + } + output.push(assistant_message.clone()?); + assistant_message_used = true; + } + ModelResponseReplayItem::FunctionCall { call_id } => { + if !used_tool_calls.insert(call_id.as_str()) { + return None; + } + let tool_call = tool_calls_by_id.get(call_id.as_str())?; + output.push(json!({ + "type": "function_call", + "call_id": tool_call.id, + "name": tool_call.name, + "arguments": tool_call.serialized_arguments(), + })); + } + } + } + + if !saw_opaque_reasoning + || assistant_message_used != assistant_message.is_some() + || used_tool_calls.len() != tool_calls.len() + { + return None; + } + + Some(output) + } + pub fn convert_messages(messages: Vec) -> Vec { let mut messages = messages .into_iter() @@ -429,8 +545,179 @@ impl OpenAIMessageConverter { mod tests { use super::OpenAIMessageConverter; use crate::types::{Message, ToolCall, ToolImageAttachment}; + use bitfun_core_types::{ + ModelReasoningSummaryPart, ModelResponseReplay, ModelResponseReplayItem, + }; use serde_json::json; + fn assistant_with_replay( + content: Option<&str>, + tool_calls: Vec, + items: Vec, + fingerprint: &str, + ) -> Message { + Message { + role: "assistant".to_string(), + content: content.map(ToString::to_string), + reasoning_content: Some("readable summary".to_string()), + thinking_signature: None, + tool_calls: (!tool_calls.is_empty()).then_some(tool_calls), + tool_call_id: None, + name: None, + is_error: None, + tool_image_attachments: None, + model_response_replay: Some(ModelResponseReplay { + protocol: "openai_responses".to_string(), + model_binding_fingerprint: fingerprint.to_string(), + items, + }), + } + } + + fn opaque_reasoning(id: &str, state: &str) -> ModelResponseReplayItem { + ModelResponseReplayItem::OpaqueReasoning { + item_id: Some(id.to_string()), + summary: vec![ModelReasoningSummaryPart { + text: format!("summary {id}"), + }], + opaque_state: state.to_string(), + } + } + + #[test] + fn replays_reasoning_before_final_assistant_message_when_fingerprint_matches() { + let message = assistant_with_replay( + Some("done"), + vec![], + vec![ + opaque_reasoning("rs_1", "opaque_1"), + ModelResponseReplayItem::AssistantMessage, + ], + "binding-1", + ); + + let (_, input) = OpenAIMessageConverter::convert_messages_to_responses_input_with_context( + vec![message], + Some("binding-1"), + ); + + assert_eq!(input.len(), 2); + assert_eq!(input[0]["type"], json!("reasoning")); + assert_eq!(input[0]["id"], json!("rs_1")); + assert_eq!(input[0]["encrypted_content"], json!("opaque_1")); + assert_eq!(input[1]["type"], json!("message")); + assert_eq!(input[1]["content"][0]["text"], json!("done")); + } + + #[test] + fn replays_multiple_reasoning_and_function_calls_in_original_order() { + let message = assistant_with_replay( + None, + vec![ + ToolCall { + id: "call_1".to_string(), + name: "one".to_string(), + arguments: json!({"value": 1}), + raw_arguments: None, + }, + ToolCall { + id: "call_2".to_string(), + name: "two".to_string(), + arguments: json!({"value": 2}), + raw_arguments: Some("{\"value\":2}".to_string()), + }, + ], + vec![ + opaque_reasoning("rs_1", "opaque_1"), + ModelResponseReplayItem::FunctionCall { + call_id: "call_2".to_string(), + }, + opaque_reasoning("rs_2", "opaque_2"), + ModelResponseReplayItem::FunctionCall { + call_id: "call_1".to_string(), + }, + ], + "binding-1", + ); + + let (_, input) = OpenAIMessageConverter::convert_messages_to_responses_input_with_context( + vec![message], + Some("binding-1"), + ); + + assert_eq!( + input + .iter() + .map(|item| item["type"].as_str().unwrap()) + .collect::>(), + vec!["reasoning", "function_call", "reasoning", "function_call"] + ); + assert_eq!(input[1]["call_id"], json!("call_2")); + assert_eq!(input[1]["arguments"], json!("{\"value\":2}")); + assert_eq!(input[3]["call_id"], json!("call_1")); + } + + #[test] + fn fingerprint_mismatch_uses_ordinary_responses_conversion() { + let message = assistant_with_replay( + Some("done"), + vec![], + vec![ + opaque_reasoning("rs_1", "opaque_1"), + ModelResponseReplayItem::AssistantMessage, + ], + "binding-old", + ); + + let (_, input) = OpenAIMessageConverter::convert_messages_to_responses_input_with_context( + vec![message], + Some("binding-new"), + ); + + assert_eq!(input.len(), 1); + assert_eq!(input[0]["type"], json!("message")); + assert!(input[0].get("encrypted_content").is_none()); + } + + #[test] + fn invalid_replay_layout_falls_back_atomically() { + let message = assistant_with_replay( + None, + vec![ + ToolCall { + id: "call_1".to_string(), + name: "one".to_string(), + arguments: json!({"value": 1}), + raw_arguments: None, + }, + ToolCall { + id: "call_2".to_string(), + name: "two".to_string(), + arguments: json!({"value": 2}), + raw_arguments: None, + }, + ], + vec![ + opaque_reasoning("rs_1", "opaque_1"), + ModelResponseReplayItem::FunctionCall { + call_id: "call_1".to_string(), + }, + ], + "binding-1", + ); + + let (_, input) = OpenAIMessageConverter::convert_messages_to_responses_input_with_context( + vec![message], + Some("binding-1"), + ); + + assert_eq!(input.len(), 2); + assert!(input.iter().all(|item| item["type"] == "function_call")); + assert!(input + .iter() + .all(|item| item.get("encrypted_content").is_none())); + } + #[test] fn converts_messages_to_responses_input() { let messages = vec![ @@ -452,6 +739,7 @@ mod tests { name: Some("get_weather".to_string()), is_error: None, tool_image_attachments: None, + model_response_replay: None, }, ]; @@ -528,6 +816,7 @@ mod tests { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, }]; let (_, input) = OpenAIMessageConverter::convert_messages_to_responses_input(messages); @@ -552,6 +841,7 @@ mod tests { mime_type: "image/jpeg".to_string(), data_base64: "AAA".to_string(), }]), + model_response_replay: None, }]; let (_, input) = OpenAIMessageConverter::convert_messages_to_responses_input(messages); @@ -583,6 +873,7 @@ mod tests { mime_type: "image/jpeg".to_string(), data_base64: "YmFi".to_string(), }]), + model_response_replay: None, }; let openai = OpenAIMessageConverter::convert_messages(vec![msg]); @@ -613,6 +904,7 @@ mod tests { name: Some("WebFetch".to_string()), is_error: None, tool_image_attachments: None, + model_response_replay: None, }; let openai = OpenAIMessageConverter::convert_messages(vec![msg]); @@ -641,6 +933,7 @@ mod tests { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, }; let openai = OpenAIMessageConverter::convert_messages(vec![msg]); @@ -672,6 +965,7 @@ mod tests { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, }; let openai = OpenAIMessageConverter::convert_messages(vec![msg]); @@ -700,6 +994,7 @@ mod tests { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, }]); assert_eq!(input[0]["content"][0]["type"], json!("input_text")); @@ -731,6 +1026,7 @@ mod tests { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, }]); assert_eq!(input[0]["content"][0]["type"], json!("input_text")); @@ -749,6 +1045,7 @@ mod tests { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, }; let openai = OpenAIMessageConverter::convert_messages(vec![msg]); @@ -773,6 +1070,7 @@ mod tests { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, }; let openai = OpenAIMessageConverter::convert_messages(vec![msg]); diff --git a/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs b/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs index 39b1ea6528..df1c5d8d50 100644 --- a/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs +++ b/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs @@ -248,7 +248,13 @@ pub(crate) async fn send_stream( // self-contained so the standard Responses path stays untouched. if super::codex_chatgpt::is_codex_chatgpt_endpoint(&client.config.request_url) { return super::codex_chatgpt::send_stream( - client, messages, tools, extra_body, max_tries, trace, + client, + messages, + tools, + extra_body, + max_tries, + trace, + request_context, ) .await; } @@ -259,8 +265,14 @@ pub(crate) async fn send_stream( client.config.model, client.config.request_url, max_tries ); + let model_binding_fingerprint = request_context + .as_ref() + .and_then(|context| context.model_binding_fingerprint.as_deref()); let (instructions, response_input) = - OpenAIMessageConverter::convert_messages_to_responses_input(messages); + OpenAIMessageConverter::convert_messages_to_responses_input_with_context( + messages, + model_binding_fingerprint, + ); let openai_tools = common::convert_tools_flat(tools); let request_body = try_build_request_body_with_context( client, @@ -353,11 +365,29 @@ mod tests { assert!(request_body["tools"][0].get("function").is_none()); } + #[test] + fn ordinary_responses_request_does_not_add_encrypted_reasoning_include() { + let request_body = build_request_body( + &test_client(), + None, + vec![json!({ + "type": "message", + "role": "user", + "content": [{ "type": "input_text", "text": "hello" }] + })], + None, + None, + ); + + assert!(request_body.get("include").is_none()); + } + #[test] fn attaches_runtime_prompt_cache_key_after_custom_body_merge() { let client = test_client(); let request_context = ModelRequestContext { prompt_cache_route_key: Some("bitfun-pc-v1-stable".to_string()), + model_binding_fingerprint: Some("binding-1".to_string()), }; let request_body = build_request_body_with_context( &client, diff --git a/src/crates/adapters/ai-adapters/src/stream/stream_handler/responses.rs b/src/crates/adapters/ai-adapters/src/stream/stream_handler/responses.rs index 98e620be0c..55a48a3524 100644 --- a/src/crates/adapters/ai-adapters/src/stream/stream_handler/responses.rs +++ b/src/crates/adapters/ai-adapters/src/stream/stream_handler/responses.rs @@ -1,7 +1,8 @@ use super::stream_stats::StreamStats; use super::{next_stream_item, StreamTimeoutController, StreamTimeoutStage, TimedStreamItem}; use crate::stream::types::responses::{ - parse_responses_output_item, ResponsesCompleted, ResponsesDone, ResponsesStreamEvent, + parse_responses_output_item, parse_responses_replay_capture, ResponsesCompleted, ResponsesDone, + ResponsesStreamEvent, }; use crate::stream::types::unified::UnifiedResponse; use anyhow::{anyhow, Result}; @@ -12,7 +13,7 @@ use log::{debug, error, trace}; use reqwest::Response; use serde_json::Value; use sha2::{Digest, Sha256}; -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::time::Duration; use tokio::sync::mpsc; @@ -48,6 +49,39 @@ fn log_cache_stream_diagnostics( ); } +fn completed_replay_capture( + response: Option<&Value>, + done_items: &BTreeMap, + done_items_complete: bool, +) -> Option { + let terminal_output = response + .and_then(|response| response.get("output")) + .and_then(Value::as_array) + .filter(|output| !output.is_empty()); + if let Some(capture) = terminal_output.and_then(|output| parse_responses_replay_capture(output)) + { + return Some(capture); + } + + if !done_items_complete || done_items.is_empty() { + return None; + } + if terminal_output.is_some_and(|output| output.len() != done_items.len()) { + return None; + } + if done_items + .keys() + .copied() + .enumerate() + .any(|(expected, actual)| expected != actual) + { + return None; + } + + let output = done_items.values().cloned().collect::>(); + parse_responses_replay_capture(&output) +} + #[derive(Debug, Default, Clone)] struct InProgressToolCall { call_id: Option, @@ -331,6 +365,8 @@ pub async fn handle_responses_stream( let mut saw_function_call = false; let mut tool_calls_by_output_index: HashMap = HashMap::new(); let mut tool_call_index_by_id: HashMap = HashMap::new(); + let mut done_items_by_output_index: BTreeMap = BTreeMap::new(); + let mut done_items_complete = true; let mut stats = StreamStats::new("Responses"); let mut timeout_controller = StreamTimeoutController::new(ttft_timeout, idle_timeout); let mut response_created_count = 0usize; @@ -517,6 +553,11 @@ pub async fn handle_responses_stream( let Some(item_value) = event.item else { continue; }; + if let Some(output_index) = event.output_index { + done_items_by_output_index.insert(output_index, item_value.clone()); + } else { + done_items_complete = false; + } // For tool calls, prefer streaming deltas and only use item.done as a tail-filler / fallback. if item_value.get("type").and_then(Value::as_str) == Some("function_call") { @@ -553,6 +594,11 @@ pub async fn handle_responses_stream( if received_finish_reason { continue; } + let model_response_replay = completed_replay_capture( + event.response.as_ref(), + &done_items_by_output_index, + done_items_complete, + ); // Best-effort: use the final response object to fill any missing tool-call argument tail. if let Some(response_val) = event.response.as_ref() { if let Some(cache_key_hash) = @@ -620,6 +666,7 @@ pub async fn handle_responses_stream( saw_function_call, )), finish_reason: Some("stop".to_string()), + model_response_replay, ..Default::default() }; emit_unified_response( @@ -652,6 +699,7 @@ pub async fn handle_responses_stream( saw_function_call, )), finish_reason: Some("stop".to_string()), + model_response_replay, ..Default::default() }; emit_unified_response( @@ -668,6 +716,11 @@ pub async fn handle_responses_stream( if received_finish_reason { continue; } + let model_response_replay = completed_replay_capture( + event.response.as_ref(), + &done_items_by_output_index, + done_items_complete, + ); if let Some(cache_key_hash) = extract_response_prompt_cache_key_hash(event.response.as_ref()) { @@ -682,6 +735,7 @@ pub async fn handle_responses_stream( saw_function_call, )), finish_reason: Some("stop".to_string()), + model_response_replay, ..Default::default() }; emit_unified_response( @@ -713,6 +767,7 @@ pub async fn handle_responses_stream( saw_function_call, )), finish_reason: Some("stop".to_string()), + model_response_replay, ..Default::default() }; emit_unified_response( @@ -796,14 +851,15 @@ pub async fn handle_responses_stream( #[cfg(test)] mod tests { use super::{ - super::stream_stats::StreamStats, extract_api_error, extract_api_error_message, - handle_function_call_arguments_delta, handle_function_call_output_item_done, - responses_completed_tool_call_completion, InProgressToolCall, StreamTimeoutController, + super::stream_stats::StreamStats, completed_replay_capture, extract_api_error, + extract_api_error_message, handle_function_call_arguments_delta, + handle_function_call_output_item_done, responses_completed_tool_call_completion, + InProgressToolCall, StreamTimeoutController, }; use bitfun_agent_stream::ToolCallCompletion; use bitfun_core_types::errors::ErrorCategory; use serde_json::json; - use std::collections::HashMap; + use std::collections::{BTreeMap, HashMap}; use tokio::sync::mpsc; #[test] @@ -956,4 +1012,104 @@ mod tests { assert!(err.to_string().contains("untracked output_index 2")); } + + #[test] + fn completed_output_is_authoritative_for_replay_layout() { + let done_items = BTreeMap::from([( + 0, + json!({ "type": "reasoning", "encrypted_content": "stale", "summary": [] }), + )]); + let response = json!({ + "output": [ + { "type": "reasoning", "encrypted_content": "final", "summary": [] }, + { "type": "message", "role": "assistant", "content": [] } + ] + }); + + let capture = + completed_replay_capture(Some(&response), &done_items, true).expect("terminal capture"); + assert_eq!(capture.items.len(), 2); + assert!(matches!( + &capture.items[0], + bitfun_core_types::ModelResponseReplayItem::OpaqueReasoning { opaque_state, .. } + if opaque_state == "final" + )); + } + + #[test] + fn completed_empty_output_falls_back_to_ordered_done_items() { + let done_items = BTreeMap::from([ + ( + 0, + json!({ "type": "reasoning", "encrypted_content": "opaque", "summary": [] }), + ), + (1, json!({ "type": "function_call", "call_id": "call_1" })), + ]); + let response = json!({ "output": [] }); + + let capture = completed_replay_capture(Some(&response), &done_items, true) + .expect("done-item fallback"); + assert_eq!(capture.items.len(), 2); + } + + #[test] + fn completed_output_without_opaque_state_falls_back_to_equivalent_done_items() { + let done_items = BTreeMap::from([ + ( + 0, + json!({ "type": "reasoning", "encrypted_content": "opaque", "summary": [] }), + ), + ( + 1, + json!({ "type": "message", "role": "assistant", "content": [] }), + ), + ]); + let response = json!({ + "output": [ + { "type": "reasoning", "summary": [] }, + { "type": "message", "role": "assistant", "content": [] } + ] + }); + + let capture = completed_replay_capture(Some(&response), &done_items, true) + .expect("equivalent done-item fallback"); + assert!(matches!( + &capture.items[0], + bitfun_core_types::ModelResponseReplayItem::OpaqueReasoning { opaque_state, .. } + if opaque_state == "opaque" + )); + } + + #[test] + fn completed_output_does_not_fall_back_to_a_shorter_done_item_subset() { + let done_items = BTreeMap::from([( + 0, + json!({ "type": "reasoning", "encrypted_content": "opaque", "summary": [] }), + )]); + let response = json!({ + "output": [ + { "type": "reasoning", "summary": [] }, + { "type": "message", "role": "assistant", "content": [] } + ] + }); + + assert!(completed_replay_capture(Some(&response), &done_items, true).is_none()); + } + + #[test] + fn incomplete_done_item_sequence_is_not_persisted() { + let done_items = BTreeMap::from([ + ( + 0, + json!({ "type": "reasoning", "encrypted_content": "opaque", "summary": [] }), + ), + ( + 2, + json!({ "type": "message", "role": "assistant", "content": [] }), + ), + ]); + + assert!(completed_replay_capture(None, &done_items, true).is_none()); + assert!(completed_replay_capture(None, &done_items, false).is_none()); + } } diff --git a/src/crates/adapters/ai-adapters/src/stream/types/anthropic.rs b/src/crates/adapters/ai-adapters/src/stream/types/anthropic.rs index 7c62dac562..9bcebcc641 100644 --- a/src/crates/adapters/ai-adapters/src/stream/types/anthropic.rs +++ b/src/crates/adapters/ai-adapters/src/stream/types/anthropic.rs @@ -107,6 +107,7 @@ impl From for UnifiedResponse { .map(map_anthropic_stop_reason), finish_reason: value.delta.stop_reason, provider_metadata: None, + model_response_replay: None, } } } diff --git a/src/crates/adapters/ai-adapters/src/stream/types/gemini.rs b/src/crates/adapters/ai-adapters/src/stream/types/gemini.rs index 99456b0d57..aa63ac4632 100644 --- a/src/crates/adapters/ai-adapters/src/stream/types/gemini.rs +++ b/src/crates/adapters/ai-adapters/src/stream/types/gemini.rs @@ -383,6 +383,7 @@ impl GeminiSSEData { tool_call_completion: None, finish_reason: None, provider_metadata: None, + model_response_replay: None, }); continue; } @@ -398,6 +399,7 @@ impl GeminiSSEData { tool_call_completion: None, finish_reason: None, provider_metadata: None, + model_response_replay: None, }); continue; } @@ -416,6 +418,7 @@ impl GeminiSSEData { tool_call_completion: None, finish_reason: None, provider_metadata: None, + model_response_replay: None, }); continue; } @@ -431,6 +434,7 @@ impl GeminiSSEData { tool_call_completion: None, finish_reason: None, provider_metadata: None, + model_response_replay: None, }); continue; } @@ -445,6 +449,7 @@ impl GeminiSSEData { tool_call_completion: None, finish_reason: None, provider_metadata: None, + model_response_replay: None, }); } } @@ -480,6 +485,7 @@ impl GeminiSSEData { tool_call_completion: None, finish_reason: None, provider_metadata: Some(provider_metadata), + model_response_replay: None, }); } diff --git a/src/crates/adapters/ai-adapters/src/stream/types/openai.rs b/src/crates/adapters/ai-adapters/src/stream/types/openai.rs index 0793714916..83d8aac39e 100644 --- a/src/crates/adapters/ai-adapters/src/stream/types/openai.rs +++ b/src/crates/adapters/ai-adapters/src/stream/types/openai.rs @@ -221,6 +221,7 @@ impl OpenAISSEData { tool_call_completion: None, finish_reason: None, provider_metadata: None, + model_response_replay: None, }); } @@ -236,6 +237,7 @@ impl OpenAISSEData { tool_call_completion: None, finish_reason: None, provider_metadata: None, + model_response_replay: None, }); } } @@ -256,6 +258,7 @@ impl OpenAISSEData { tool_call_completion: Some(map_openai_finish_reason(&finish_reason)), finish_reason: Some(finish_reason), provider_metadata: None, + model_response_replay: None, }); return responses; } @@ -270,6 +273,7 @@ impl OpenAISSEData { tool_call_completion: None, finish_reason, provider_metadata: None, + model_response_replay: None, }); } diff --git a/src/crates/adapters/ai-adapters/src/stream/types/responses.rs b/src/crates/adapters/ai-adapters/src/stream/types/responses.rs index 38b81e6897..27282f03d1 100644 --- a/src/crates/adapters/ai-adapters/src/stream/types/responses.rs +++ b/src/crates/adapters/ai-adapters/src/stream/types/responses.rs @@ -1,7 +1,11 @@ use super::unified::{UnifiedResponse, UnifiedTokenUsage, UnifiedToolCall}; +use bitfun_agent_stream::ModelResponseReplayCapture; +use bitfun_core_types::{ModelReasoningSummaryPart, ModelResponseReplayItem}; use serde::Deserialize; use serde_json::Value; +pub const OPENAI_RESPONSES_REPLAY_PROTOCOL: &str = "openai_responses"; + #[derive(Debug, Deserialize)] pub struct ResponsesStreamEvent { #[serde(rename = "type")] @@ -101,6 +105,7 @@ pub fn parse_responses_output_item( tool_call_completion: None, finish_reason: None, provider_metadata: None, + model_response_replay: None, }), "message" => { let text = item_value @@ -126,18 +131,98 @@ pub fn parse_responses_output_item( tool_call_completion: None, finish_reason: None, provider_metadata: None, + model_response_replay: None, + }) + } + _ => None, + } +} + +pub fn parse_responses_replay_capture(output: &[Value]) -> Option { + if output.is_empty() { + return None; + } + + let mut contains_opaque_reasoning = false; + let mut items = Vec::with_capacity(output.len()); + for item in output { + let replay_item = parse_responses_replay_item(item)?; + contains_opaque_reasoning |= + matches!(replay_item, ModelResponseReplayItem::OpaqueReasoning { .. }); + items.push(replay_item); + } + + contains_opaque_reasoning.then(|| ModelResponseReplayCapture { + protocol: OPENAI_RESPONSES_REPLAY_PROTOCOL.to_string(), + items, + }) +} + +fn parse_responses_replay_item(item: &Value) -> Option { + match item.get("type")?.as_str()? { + "reasoning" => { + let opaque_state = item + .get("encrypted_content") + .and_then(Value::as_str) + .filter(|value| !value.is_empty())? + .to_string(); + let item_id = item + .get("id") + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + .map(ToString::to_string); + let summary = match item.get("summary").filter(|summary| !summary.is_null()) { + Some(summary) => parse_reasoning_summary(summary)?, + None => Vec::new(), + }; + Some(ModelResponseReplayItem::OpaqueReasoning { + item_id, + summary, + opaque_state, }) } + "message" + if item + .get("role") + .and_then(Value::as_str) + .is_none_or(|role| role == "assistant") => + { + Some(ModelResponseReplayItem::AssistantMessage) + } + "function_call" => item + .get("call_id") + .and_then(Value::as_str) + .filter(|call_id| !call_id.is_empty()) + .map(|call_id| ModelResponseReplayItem::FunctionCall { + call_id: call_id.to_string(), + }), _ => None, } } +fn parse_reasoning_summary(value: &Value) -> Option> { + value + .as_array()? + .iter() + .map(|part| { + (part.get("type").and_then(Value::as_str) == Some("summary_text")) + .then(|| part.get("text").and_then(Value::as_str)) + .flatten() + .map(|text| ModelReasoningSummaryPart { + text: text.to_string(), + }) + }) + .collect() +} + #[cfg(test)] mod tests { use super::{ - parse_responses_output_item, ResponsesCompleted, ResponsesStreamEvent, ResponsesUsage, + parse_responses_output_item, parse_responses_replay_capture, ResponsesCompleted, + ResponsesStreamEvent, ResponsesUsage, }; use crate::stream::types::unified::UnifiedTokenUsage; + use bitfun_core_types::ModelResponseReplayItem; use serde_json::json; #[test] @@ -263,4 +348,68 @@ mod tests { assert_eq!(event.output_index, Some(1)); assert_eq!(event.delta.as_deref(), Some("{\"a\":")); } + + #[test] + fn captures_reasoning_message_and_function_call_in_output_order() { + let capture = parse_responses_replay_capture(&[ + json!({ + "id": "rs_1", + "type": "reasoning", + "summary": [{ "type": "summary_text", "text": "summary" }], + "encrypted_content": "opaque" + }), + json!({ + "type": "function_call", + "call_id": "call_1", + "name": "tool", + "arguments": "{}" + }), + json!({ + "type": "message", + "role": "assistant", + "content": [{ "type": "output_text", "text": "done" }] + }), + ]) + .expect("replay capture"); + + assert_eq!(capture.items.len(), 3); + assert!(matches!( + &capture.items[0], + ModelResponseReplayItem::OpaqueReasoning { item_id, summary, opaque_state } + if item_id.as_deref() == Some("rs_1") + && summary.first().map(|part| part.text.as_str()) == Some("summary") + && opaque_state == "opaque" + )); + assert!(matches!( + &capture.items[1], + ModelResponseReplayItem::FunctionCall { call_id } if call_id == "call_1" + )); + assert!(matches!( + &capture.items[2], + ModelResponseReplayItem::AssistantMessage + )); + } + + #[test] + fn does_not_capture_reasoning_without_opaque_state() { + assert!(parse_responses_replay_capture(&[json!({ + "id": "rs_1", + "type": "reasoning", + "summary": [{ "type": "summary_text", "text": "summary" }] + })]) + .is_none()); + } + + #[test] + fn capture_is_atomic_when_output_contains_an_unsupported_item() { + assert!(parse_responses_replay_capture(&[ + json!({ + "type": "reasoning", + "encrypted_content": "opaque", + "summary": [] + }), + json!({ "type": "computer_call", "id": "computer_1" }), + ]) + .is_none()); + } } diff --git a/src/crates/adapters/ai-adapters/tests/ai_stream_contracts/stream_replay_regressions.rs b/src/crates/adapters/ai-adapters/tests/ai_stream_contracts/stream_replay_regressions.rs index c499c7ec69..3e3dba4700 100644 --- a/src/crates/adapters/ai-adapters/tests/ai_stream_contracts/stream_replay_regressions.rs +++ b/src/crates/adapters/ai-adapters/tests/ai_stream_contracts/stream_replay_regressions.rs @@ -41,6 +41,7 @@ fn build_replay_assistant_message(result: &StreamResult) -> AIMessage { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, } } diff --git a/src/crates/assembly/core/src/agentic/core/message.rs b/src/crates/assembly/core/src/agentic/core/message.rs index 1bbe0c5b5c..0e7463f728 100644 --- a/src/crates/assembly/core/src/agentic/core/message.rs +++ b/src/crates/assembly/core/src/agentic/core/message.rs @@ -2,6 +2,7 @@ use crate::agentic::image_analysis::ImageContextData; use crate::util::types::{Message as AIMessage, ToolCall as AIToolCall, ToolImageAttachment}; use crate::util::TokenCounter; use bitfun_agent_runtime::prompt_markup::is_system_reminder_only; +use bitfun_core_types::ModelResponseReplay; pub use bitfun_runtime_ports::{CompressionContract, CompressionContractItem}; use log::warn; use serde::{Deserialize, Serialize}; @@ -74,6 +75,8 @@ pub struct MessageMetadata { /// reminders so activation can be reconstructed from persisted history. #[serde(default, skip_serializing_if = "Vec::is_empty")] pub activated_instruction_sources: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model_response_replay: Option, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] @@ -259,6 +262,7 @@ impl From for AIMessage { MessageRole::System => "system", }; let thinking_signature = msg.metadata.thinking_signature.clone(); + let model_response_replay = msg.metadata.model_response_replay.clone(); match msg.content { MessageContent::Text(text) => { @@ -287,6 +291,7 @@ impl From for AIMessage { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: model_response_replay.clone(), } } MessageContent::Multimodal { text, images } => { @@ -321,6 +326,7 @@ impl From for AIMessage { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: model_response_replay.clone(), } } MessageContent::Mixed { @@ -363,6 +369,7 @@ impl From for AIMessage { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: model_response_replay.clone(), } } MessageContent::ToolResult { @@ -401,6 +408,7 @@ impl From for AIMessage { name: Some(tool_name), is_error: Some(is_error), tool_image_attachments: image_attachments.clone(), + model_response_replay: model_response_replay.clone(), } } } @@ -605,6 +613,14 @@ impl Message { self } + pub fn with_model_response_replay( + mut self, + model_response_replay: Option, + ) -> Self { + self.metadata.model_response_replay = model_response_replay; + self + } + /// Get message's token count pub fn get_tokens(&mut self) -> usize { if let Some(tokens) = self.metadata.tokens { @@ -758,6 +774,7 @@ mod tests { use super::{Message, ToolCall}; use crate::util::types::Message as AIMessage; use bitfun_agent_stream::ToolArgumentRepairKind; + use bitfun_core_types::{ModelResponseReplay, ModelResponseReplayItem}; use serde_json::json; #[test] @@ -771,6 +788,69 @@ mod tests { assert_eq!(ai_msg.thinking_signature.as_deref(), Some("sig_1")); } + #[test] + fn persists_and_restores_model_response_replay() { + let message = Message::assistant("done".to_string()).with_model_response_replay(Some( + ModelResponseReplay { + protocol: "openai_responses".to_string(), + model_binding_fingerprint: "binding-1".to_string(), + items: vec![ModelResponseReplayItem::OpaqueReasoning { + item_id: Some("rs_1".to_string()), + summary: vec![], + opaque_state: "opaque".to_string(), + }], + }, + )); + + let encoded = serde_json::to_string(&message).expect("serialize message"); + let restored: Message = serde_json::from_str(&encoded).expect("deserialize message"); + let replay = restored + .metadata + .model_response_replay + .expect("restored replay"); + + assert_eq!(replay.protocol, "openai_responses"); + assert_eq!(replay.model_binding_fingerprint, "binding-1"); + assert!(matches!( + &replay.items[0], + ModelResponseReplayItem::OpaqueReasoning { opaque_state, .. } + if opaque_state == "opaque" + )); + } + + #[test] + fn legacy_message_without_model_response_replay_still_deserializes() { + let message = Message::assistant("done".to_string()); + let mut encoded = serde_json::to_value(message).expect("serialize message"); + encoded["metadata"] + .as_object_mut() + .expect("metadata object") + .remove("model_response_replay"); + + let restored: Message = serde_json::from_value(encoded).expect("legacy message"); + assert!(restored.metadata.model_response_replay.is_none()); + } + + #[test] + fn carries_model_response_replay_into_adapter_message() { + let message = Message::assistant("done".to_string()).with_model_response_replay(Some( + ModelResponseReplay { + protocol: "openai_responses".to_string(), + model_binding_fingerprint: "binding-1".to_string(), + items: vec![ModelResponseReplayItem::AssistantMessage], + }, + )); + + let ai_message = AIMessage::from(message); + assert_eq!( + ai_message + .model_response_replay + .as_ref() + .map(|replay| replay.model_binding_fingerprint.as_str()), + Some("binding-1") + ); + } + #[test] fn preserves_tool_argument_repair_provenance_from_stream_contract() { let tool_call = ToolCall::from(bitfun_agent_stream::ToolCall { diff --git a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs index c62e155634..323567a606 100644 --- a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs +++ b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs @@ -577,6 +577,7 @@ impl ExecutionEngine { "bitfun-pc-v1-{}", hex::encode(hasher.finalize()) )), + model_binding_fingerprint: Some(model_binding_fingerprint.to_string()), } } @@ -6481,6 +6482,14 @@ mod tests { .prompt_cache_route_key .as_deref() .is_some_and(|key| key.starts_with("bitfun-pc-v1-"))); + assert_eq!( + first.model_binding_fingerprint.as_deref(), + Some("binding-1") + ); + assert_eq!( + changed.model_binding_fingerprint.as_deref(), + Some("binding-2") + ); } fn command_result(tool_name: &str, success: bool, exit_code: Option) -> Message { diff --git a/src/crates/assembly/core/src/agentic/execution/round_executor.rs b/src/crates/assembly/core/src/agentic/execution/round_executor.rs index 9967c0c54a..fa68520747 100644 --- a/src/crates/assembly/core/src/agentic/execution/round_executor.rs +++ b/src/crates/assembly/core/src/agentic/execution/round_executor.rs @@ -42,6 +42,7 @@ use bitfun_ai_adapters::{ ModelExchangeRequestTraceHandle, ModelExchangeResponseTrace, ModelExchangeTraceConfig, }; use bitfun_core_types::errors::{AiProviderError, ErrorCategory}; +use bitfun_core_types::ModelResponseReplay; use bitfun_runtime_ports::PermissionRule; use log::{debug, error, warn}; use std::sync::Arc; @@ -219,6 +220,23 @@ impl RoundExecutor { .map(Into::into) } + fn bound_model_response_replay( + stream_result: &StreamResult, + context: &RoundContext, + ) -> Option { + let capture = stream_result.model_response_replay.as_ref()?; + let model_binding_fingerprint = context + .model_request_context + .model_binding_fingerprint + .as_ref()? + .clone(); + Some(ModelResponseReplay { + protocol: capture.protocol.clone(), + model_binding_fingerprint, + items: capture.items.clone(), + }) + } + fn map_subagent_batch_execution_policy( policy: ConfigSubagentBatchExecutionPolicy, ) -> PipelineSubagentBatchExecutionPolicy { @@ -964,13 +982,15 @@ impl RoundExecutor { }; let parsed_memory_citation = Self::parsed_memory_citation_from_stream_result(&stream_result); + let model_response_replay = Self::bound_model_response_replay(&stream_result, &context); let (clean_text, _) = strip_bitfun_memory_citations(&stream_result.full_text); let assistant_message = Message::assistant_with_reasoning(reasoning, clean_text, vec![]) .with_turn_id(context.dialog_turn_id.clone()) .with_round_id(round_id.clone()) .with_thinking_signature(stream_result.thinking_signature.clone()) - .with_memory_citation(parsed_memory_citation); + .with_memory_citation(parsed_memory_citation) + .with_model_response_replay(model_response_replay); debug!("Returning RoundResult: has_more_rounds=false"); debug!( @@ -1186,13 +1206,15 @@ impl RoundExecutor { }; let parsed_memory_citation = Self::parsed_memory_citation_from_stream_result(&stream_result); + let model_response_replay = Self::bound_model_response_replay(&stream_result, &context); let (clean_text, _) = strip_bitfun_memory_citations(&stream_result.full_text); let assistant_message = Message::assistant_with_reasoning(reasoning, clean_text, tool_calls.clone()) .with_turn_id(context.dialog_turn_id.clone()) .with_round_id(round_id.clone()) .with_thinking_signature(stream_result.thinking_signature.clone()) - .with_memory_citation(parsed_memory_citation); + .with_memory_citation(parsed_memory_citation) + .with_model_response_replay(model_response_replay); debug!( "Tool execution completed, creating message: assistant_msg_len={}, tool_results={}", @@ -2011,6 +2033,7 @@ mod tests { cache_creation_token_count: None, }), provider_metadata: Some(json!({ "finish_reason": "tool_calls" })), + model_response_replay: None, has_effective_output: false, first_chunk_ms: Some(10), first_visible_output_ms: None, diff --git a/src/crates/assembly/core/src/agentic/execution/stream_processor.rs b/src/crates/assembly/core/src/agentic/execution/stream_processor.rs index 039dbcf09e..f407eafe80 100644 --- a/src/crates/assembly/core/src/agentic/execution/stream_processor.rs +++ b/src/crates/assembly/core/src/agentic/execution/stream_processor.rs @@ -12,8 +12,8 @@ use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; pub use bitfun_agent_stream::{ - HiddenTextBlock, HiddenTextTag, StreamProcessOptions, StreamProcessorError, - ToolCall as StreamToolCall, + HiddenTextBlock, HiddenTextTag, ModelResponseReplayCapture, StreamProcessOptions, + StreamProcessorError, ToolCall as StreamToolCall, }; const MEMORY_CITATION_HIDDEN_TEXT_TAG: &str = "memory_citation"; @@ -29,6 +29,7 @@ pub struct StreamResult { pub tool_calls: Vec, pub usage: Option, pub provider_metadata: Option, + pub model_response_replay: Option, pub has_effective_output: bool, pub first_chunk_ms: Option, pub first_visible_output_ms: Option, @@ -46,6 +47,7 @@ impl From for StreamResult { tool_calls: result.tool_calls.into_iter().map(Into::into).collect(), usage: result.usage.map(Into::into), provider_metadata: result.provider_metadata, + model_response_replay: result.model_response_replay, has_effective_output: result.has_effective_output, first_chunk_ms: result.first_chunk_ms, first_visible_output_ms: result.first_visible_output_ms, diff --git a/src/crates/assembly/core/src/agentic/image_analysis/image_processing.rs b/src/crates/assembly/core/src/agentic/image_analysis/image_processing.rs index 1f9c09f48d..45a0774b19 100644 --- a/src/crates/assembly/core/src/agentic/image_analysis/image_processing.rs +++ b/src/crates/assembly/core/src/agentic/image_analysis/image_processing.rs @@ -304,6 +304,7 @@ pub fn build_multimodal_message( name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, } } else if provider_lower.contains("gemini") || provider_lower.contains("google") { Message { @@ -326,6 +327,7 @@ pub fn build_multimodal_message( name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, } } else { // Default to OpenAI-compatible payload shape for OpenAI and most OpenAI-compatible providers. @@ -350,6 +352,7 @@ pub fn build_multimodal_message( name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, } }; @@ -467,6 +470,7 @@ pub fn build_multimodal_message_with_images( name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, }]) } diff --git a/src/crates/assembly/core/src/agentic/session/session_manager.rs b/src/crates/assembly/core/src/agentic/session/session_manager.rs index 011a256682..e43f0a298b 100644 --- a/src/crates/assembly/core/src/agentic/session/session_manager.rs +++ b/src/crates/assembly/core/src/agentic/session/session_manager.rs @@ -8969,6 +8969,7 @@ impl SessionManager { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, }, Message { role: "user".to_string(), @@ -8980,6 +8981,7 @@ impl SessionManager { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, }, ]; diff --git a/src/crates/contracts/core-types/src/ai.rs b/src/crates/contracts/core-types/src/ai.rs index 304215364d..c67294203e 100644 --- a/src/crates/contracts/core-types/src/ai.rs +++ b/src/crates/contracts/core-types/src/ai.rs @@ -653,6 +653,40 @@ pub struct AIConfig { pub struct ModelRequestContext { /// Stable, opaque routing identity for provider-side prompt-prefix caches. pub prompt_cache_route_key: Option, + /// Fingerprint of the resolved provider/model/endpoint/request binding. + /// + /// Provider adapters use this only to decide whether opaque response state + /// from an earlier round is compatible with the current request. + pub model_binding_fingerprint: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ModelResponseReplay { + pub protocol: String, + pub model_binding_fingerprint: String, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub items: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum ModelResponseReplayItem { + OpaqueReasoning { + #[serde(default, skip_serializing_if = "Option::is_none")] + item_id: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + summary: Vec, + opaque_state: String, + }, + AssistantMessage, + FunctionCall { + call_id: String, + }, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ModelReasoningSummaryPart { + pub text: String, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -731,6 +765,8 @@ pub struct Message { pub is_error: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub tool_image_attachments: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model_response_replay: Option, } impl Message { @@ -745,6 +781,7 @@ impl Message { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, } } @@ -759,6 +796,7 @@ impl Message { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, } } @@ -773,6 +811,7 @@ impl Message { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, } } @@ -787,6 +826,7 @@ impl Message { name: None, is_error: None, tool_image_attachments: None, + model_response_replay: None, } } } diff --git a/src/crates/contracts/core-types/src/lib.rs b/src/crates/contracts/core-types/src/lib.rs index 9e0cfe79ce..8766e5b8ea 100644 --- a/src/crates/contracts/core-types/src/lib.rs +++ b/src/crates/contracts/core-types/src/lib.rs @@ -15,17 +15,17 @@ pub mod tool_image_attachment; pub mod worktree; pub use ai::{ - AIConfig, ConnectionTestMessageCode, ConnectionTestResult, Message, ModelRequestContext, - ModelsDevCatalogSource, ModelsDevCatalogStatus, ModelsDevReasoningCatalog, - ModelsDevReasoningModel, ModelsDevReasoningProvider, ModelsDevRefreshResult, - ModelsDevRefreshStatus, ProviderCatalog, ProviderCatalogEndpoint, ProviderCatalogModel, - ProviderCatalogModelCapabilities, ProviderCatalogModelLimits, ProviderCatalogModelPricing, - ProviderCatalogModelSource, ProviderCatalogProvider, ProviderCatalogSource, - ProviderCatalogUpstreamProvider, ProxyConfig, ReasoningCapabilityStatus, - ReasoningCatalogBinding, ReasoningCatalogProjection, ReasoningCatalogProjectionRequest, - ReasoningConfig, ReasoningPreset, ReasoningPresetAction, ReasoningPresetDescriptor, - ReasoningPresetSource, RemoteModelInfo, ToolCall, ToolCallConfirmationDetails, - ToolCallRequestInfo, ToolCallResponseInfo, ToolDefinition, + AIConfig, ConnectionTestMessageCode, ConnectionTestResult, Message, ModelReasoningSummaryPart, + ModelRequestContext, ModelResponseReplay, ModelResponseReplayItem, ModelsDevCatalogSource, + ModelsDevCatalogStatus, ModelsDevReasoningCatalog, ModelsDevReasoningModel, + ModelsDevReasoningProvider, ModelsDevRefreshResult, ModelsDevRefreshStatus, ProviderCatalog, + ProviderCatalogEndpoint, ProviderCatalogModel, ProviderCatalogModelCapabilities, + ProviderCatalogModelLimits, ProviderCatalogModelPricing, ProviderCatalogModelSource, + ProviderCatalogProvider, ProviderCatalogSource, ProviderCatalogUpstreamProvider, ProxyConfig, + ReasoningCapabilityStatus, ReasoningCatalogBinding, ReasoningCatalogProjection, + ReasoningCatalogProjectionRequest, ReasoningConfig, ReasoningPreset, ReasoningPresetAction, + ReasoningPresetDescriptor, ReasoningPresetSource, RemoteModelInfo, ToolCall, + ToolCallConfirmationDetails, ToolCallRequestInfo, ToolCallResponseInfo, ToolDefinition, }; pub use errors::{AiErrorDetail, ErrorCategory}; pub use model::{ diff --git a/src/crates/execution/agent-stream/src/lib.rs b/src/crates/execution/agent-stream/src/lib.rs index fa4f6cd849..9a90fc391f 100644 --- a/src/crates/execution/agent-stream/src/lib.rs +++ b/src/crates/execution/agent-stream/src/lib.rs @@ -24,7 +24,9 @@ use std::fmt; use std::sync::Arc; use std::time::Instant; use tokio::sync::mpsc; -pub use unified::{UnifiedResponse, UnifiedTokenUsage, UnifiedToolCall}; +pub use unified::{ + ModelResponseReplayCapture, UnifiedResponse, UnifiedTokenUsage, UnifiedToolCall, +}; /// Minimal tool-call value emitted by the stream processor. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -259,6 +261,8 @@ pub struct StreamResult { pub usage: Option, /// Provider-specific metadata captured from the stream tail. pub provider_metadata: Option, + /// Complete provider response layout and opaque state eligible for replay. + pub model_response_replay: Option, /// Whether this stream produced any user-visible output (text/thinking/tool events) pub has_effective_output: bool, /// Milliseconds from stream processing start to the first upstream response item. @@ -318,6 +322,7 @@ struct StreamContext { tool_calls: Vec, usage: Option, provider_metadata: Option, + model_response_replay: Option, // Current tool call state pending_tool_calls: PendingToolCalls, @@ -363,6 +368,7 @@ impl StreamContext { tool_calls: Vec::new(), usage: None, provider_metadata: None, + model_response_replay: None, pending_tool_calls: PendingToolCalls::new(), finalized_tool_call_ids: HashSet::new(), stream_started_at: Instant::now(), @@ -389,6 +395,7 @@ impl StreamContext { tool_calls: self.tool_calls, usage: self.usage, provider_metadata: self.provider_metadata, + model_response_replay: self.model_response_replay, has_effective_output: self.has_effective_output, first_chunk_ms: self.first_chunk_ms, first_visible_output_ms: self.first_visible_output_ms, @@ -1122,6 +1129,7 @@ impl StreamProcessor { tool_call_completion, finish_reason, provider_metadata, + model_response_replay, } = response; ctx.mark_first_stream_chunk(); @@ -1176,6 +1184,10 @@ impl StreamProcessor { } } + if let Some(model_response_replay) = model_response_replay { + ctx.model_response_replay = Some(model_response_replay); + } + if let Some(reason) = finish_reason { let completion = tool_call_completion.unwrap_or(ToolCallCompletion::Unknown); let _ = ctx.finalize_all_pending_tool_calls( @@ -1231,12 +1243,14 @@ impl StreamProcessor { #[cfg(test)] mod tests { use super::{ - is_token_limit_finish_reason, GracefulShutdownInput, HiddenTextTag, SseLogCollector, - SseLogConfig, StreamEventSink, StreamProcessOptions, StreamProcessor, StreamProcessorError, - ToolArgumentRepairKind, ToolCall, ToolCallCompletion, + is_token_limit_finish_reason, GracefulShutdownInput, HiddenTextTag, + ModelResponseReplayCapture, SseLogCollector, SseLogConfig, StreamEventSink, + StreamProcessOptions, StreamProcessor, StreamProcessorError, ToolArgumentRepairKind, + ToolCall, ToolCallCompletion, }; use super::{UnifiedResponse, UnifiedTokenUsage, UnifiedToolCall}; use bitfun_core_types::errors::{AiProviderError, ErrorCategory}; + use bitfun_core_types::ModelResponseReplayItem; use bitfun_events::{AgenticEvent, AgenticEventPriority as EventPriority, ToolEventData}; use futures::StreamExt; use serde_json::json; @@ -2165,4 +2179,50 @@ mod tests { assert!(result.full_thinking.is_empty()); assert!(!result.has_effective_output); } + + #[tokio::test] + async fn carries_complete_model_response_replay_to_stream_result() { + let processor = build_processor(); + let stream = iter(vec![ + Ok(UnifiedResponse { + text: Some("done".to_string()), + ..Default::default() + }), + Ok(UnifiedResponse { + model_response_replay: Some(ModelResponseReplayCapture { + protocol: "openai_responses".to_string(), + items: vec![ + ModelResponseReplayItem::OpaqueReasoning { + item_id: Some("rs_1".to_string()), + summary: vec![], + opaque_state: "opaque".to_string(), + }, + ModelResponseReplayItem::AssistantMessage, + ], + }), + finish_reason: Some("stop".to_string()), + ..Default::default() + }), + ]) + .boxed(); + + let result = processor + .process_stream( + stream, + None, + None, + "session_1".to_string(), + "turn_1".to_string(), + "round_1".to_string(), + "round_1:attempt:1".to_string(), + 1, + &CancellationToken::new(), + ) + .await + .expect("stream result"); + + let replay = result.model_response_replay.expect("replay capture"); + assert_eq!(replay.protocol, "openai_responses"); + assert_eq!(replay.items.len(), 2); + } } diff --git a/src/crates/execution/agent-stream/src/unified.rs b/src/crates/execution/agent-stream/src/unified.rs index df63737d86..1aada37dc3 100644 --- a/src/crates/execution/agent-stream/src/unified.rs +++ b/src/crates/execution/agent-stream/src/unified.rs @@ -1,4 +1,5 @@ use crate::tool_call_accumulator::ToolCallCompletion; +use bitfun_core_types::ModelResponseReplayItem; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::borrow::Cow; @@ -15,6 +16,12 @@ pub struct UnifiedToolCall { pub arguments_is_snapshot: bool, } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct ModelResponseReplayCapture { + pub protocol: String, + pub items: Vec, +} + /// Unified AI response format #[derive(Clone, Serialize, Deserialize, Default)] pub struct UnifiedResponse { @@ -33,6 +40,8 @@ pub struct UnifiedResponse { pub finish_reason: Option, #[serde(skip_serializing_if = "Option::is_none")] pub provider_metadata: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub model_response_replay: Option, } impl fmt::Debug for UnifiedResponse { @@ -61,6 +70,12 @@ impl fmt::Debug for UnifiedResponse { .field("tool_call_completion", &self.tool_call_completion) .field("finish_reason", &self.finish_reason) .field("provider_metadata", &"") + .field( + "model_response_replay", + &self.model_response_replay.as_ref().map(|replay| { + format!("protocol={}, items={}", replay.protocol, replay.items.len()) + }), + ) .finish() } } From 4501deea752131d1f82b777cc47a13b8e1a097ab Mon Sep 17 00:00:00 2001 From: wsp Date: Mon, 17 Aug 2026 00:41:59 +0800 Subject: [PATCH 3/5] fix(responses): preserve prompt cache lineage Persist a stable prompt cache lineage independently from the concrete session ID and inherit it when sessions retain a shared prompt prefix. - Route provider caches by lineage instead of model or prompt scope - Keep fresh subagents on independent cache identities - Reuse lineage for forked, BTW, and branched sessions - Pass request context through compaction summary requests - Keep binding fingerprints separate for encrypted reasoning replay --- src/crates/adapters/ai-adapters/src/client.rs | 39 ++++++- .../src/providers/openai/responses.rs | 7 +- .../src/agentic/coordination/coordinator.rs | 6 + .../src/agentic/execution/execution_engine.rs | 103 ++++++------------ .../core/src/agentic/fork_agent/mod.rs | 30 ++++- .../src/agentic/persistence/session_branch.rs | 16 ++- .../execution/agent-runtime/src/session.rs | 49 +++++++++ 7 files changed, 172 insertions(+), 78 deletions(-) diff --git a/src/crates/adapters/ai-adapters/src/client.rs b/src/crates/adapters/ai-adapters/src/client.rs index 0832ce3cdb..54c593a578 100644 --- a/src/crates/adapters/ai-adapters/src/client.rs +++ b/src/crates/adapters/ai-adapters/src/client.rs @@ -230,6 +230,7 @@ impl AIClient { extra_body, SEND_MESSAGE_STREAM_ATTEMPTS, trace, + None, ) .await } @@ -281,6 +282,7 @@ impl AIClient { custom_body, 1, trace, + None, ) .await } @@ -311,11 +313,31 @@ impl AIClient { tools: Option>, trace: Option, ) -> Result { - let custom_body = self.config.custom_request_body.clone(); - self.send_message_with_extra_body_and_trace(messages, tools, custom_body, trace) + self.send_message_with_trace_and_request_context(messages, tools, None, trace) .await } + /// Aggregate one model response while carrying provider-neutral + /// request-scoped facts to adapters that support them. + pub async fn send_message_with_trace_and_request_context( + &self, + messages: Vec, + tools: Option>, + request_context: Option, + trace: Option, + ) -> Result { + let custom_body = self.config.custom_request_body.clone(); + self.send_message_with_extra_body_trace_and_max_attempts( + messages, + tools, + custom_body, + request_context, + trace, + SEND_MESSAGE_STREAM_ATTEMPTS, + ) + .await + } + pub async fn send_message_with_extra_body_and_trace( &self, messages: Vec, @@ -327,6 +349,7 @@ impl AIClient { messages, tools, extra_body, + None, trace, SEND_MESSAGE_STREAM_ATTEMPTS, ) @@ -338,6 +361,7 @@ impl AIClient { messages: Vec, tools: Option>, extra_body: Option, + request_context: Option, trace: Option, max_attempts: usize, ) -> Result { @@ -349,6 +373,7 @@ impl AIClient { extra_body.clone(), 1, trace.clone(), + request_context.clone(), ) .await { @@ -410,6 +435,7 @@ impl AIClient { extra_body: Option, max_tries: usize, trace: Option, + request_context: Option, ) -> Result { match ApiFormat::parse(&self.config.format)? { ApiFormat::OpenAIChat => { @@ -417,7 +443,13 @@ impl AIClient { } ApiFormat::OpenAIResponses => { openai::responses::send_stream( - self, messages, tools, extra_body, max_tries, trace, None, + self, + messages, + tools, + extra_body, + max_tries, + trace, + request_context, ) .await } @@ -458,6 +490,7 @@ impl AIClient { tools, custom_body, None, + None, max_attempts, ) .await diff --git a/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs b/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs index df1c5d8d50..e622a66ca2 100644 --- a/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs +++ b/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs @@ -386,7 +386,7 @@ mod tests { fn attaches_runtime_prompt_cache_key_after_custom_body_merge() { let client = test_client(); let request_context = ModelRequestContext { - prompt_cache_route_key: Some("bitfun-pc-v1-stable".to_string()), + prompt_cache_route_key: Some("lineage-1".to_string()), model_binding_fingerprint: Some("binding-1".to_string()), }; let request_body = build_request_body_with_context( @@ -398,9 +398,6 @@ mod tests { Some(&request_context), ); - assert_eq!( - request_body["prompt_cache_key"], - json!("bitfun-pc-v1-stable") - ); + assert_eq!(request_body["prompt_cache_key"], json!("lineage-1")); } } diff --git a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs index 2b7ee2dcee..8709105d68 100644 --- a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs +++ b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs @@ -18420,6 +18420,7 @@ mod tests { project_workspace_path: Some(project_workspace_path.clone()), execution_target: Some(execution_target.clone()), workspace_id: Some("workspace-1".to_string()), + prompt_cache_lineage_id: Some("parent-lineage".to_string()), ..Default::default() }, ) @@ -18467,6 +18468,7 @@ mod tests { resolved.session_config.workspace_id.as_deref(), Some("workspace-1") ); + assert!(resolved.session_config.prompt_cache_lineage_id.is_none()); } #[tokio::test] @@ -19092,6 +19094,10 @@ mod tests { child_session.config.remote_ssh_host.as_deref(), Some("example.test") ); + assert_eq!( + child_session.config.prompt_cache_lineage_id.as_deref(), + Some(parent_session.session_id.as_str()) + ); assert_eq!( session_manager .cached_system_prompt(&child_session.session_id, &system_prompt_identity) diff --git a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs index 323567a606..15cc6132ef 100644 --- a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs +++ b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs @@ -66,7 +66,6 @@ use crate::util::types::ToolDefinition; use crate::util::{elapsed_ms_u64, truncate_at_char_boundary}; use bitfun_agent_runtime::output_surface::TOOL_CONTEXT_INLINE_MARKDOWN_IMAGE_DISPLAY_KEY; use bitfun_agent_runtime::permission::PERMISSION_MODE_CONTEXT_KEY; -use bitfun_agent_runtime::prompt_cache::prompt_cache_scope_key; use bitfun_agent_runtime::remote_file_delivery::TOOL_CONTEXT_REMOTE_FILE_DELIVERY_KEY; use bitfun_agent_runtime::thread_goal_tools::ensure_thread_goal_tools; use bitfun_ai_adapters::ModelExchangeTraceConfig; @@ -271,6 +270,7 @@ async fn activate_conditional_instructions_after_round( struct CompressionRuntimeScaffold { ai_client: Arc, + model_request_context: ModelRequestContext, tool_definitions: Option>, system_prompt_message: Message, prepended_prompt_reminders: PrependedPromptReminders, @@ -506,6 +506,7 @@ struct FinalizeRoundInput<'a> { struct CompressionModelSummaryInput<'a> { trace_config: Option, + model_request_context: &'a ModelRequestContext, primary_supports_image_understanding: bool, prepended_prompt_reminders: &'a PrependedPromptReminders, tool_definitions: &'a Option>, @@ -526,7 +527,6 @@ pub struct ExecutionEngine { } impl ExecutionEngine { - const PROVIDER_PROMPT_CACHE_ROUTE_SCHEMA: &'static str = "bitfun-provider-prompt-cache-v1"; const AUTO_COMPRESSION_SAFETY_RESERVE_TOKENS: usize = 10_000; const MAX_COMPRESSION_OVERFLOW_ATTEMPTS: usize = 4; const MAX_MAIN_CONTEXT_OVERFLOW_RECOVERIES: usize = 2; @@ -538,45 +538,11 @@ impl ExecutionEngine { "Provide a final answer. You MUST not call any tools."; fn model_request_context( - session_id: &str, + prompt_cache_lineage_id: &str, model_binding_fingerprint: &str, - effective_model_name: &str, - current_agent: &dyn crate::agentic::agents::Agent, ) -> ModelRequestContext { - let prompt_scope = prompt_cache_scope_key( - ¤t_agent.system_prompt_cache_identity(Some(effective_model_name)), - ¤t_agent.user_context_cache_identity(), - ); - Self::model_request_context_from_scope_key( - session_id, - model_binding_fingerprint, - effective_model_name, - &prompt_scope, - ) - } - - fn model_request_context_from_scope_key( - session_id: &str, - model_binding_fingerprint: &str, - effective_model_name: &str, - prompt_scope: &str, - ) -> ModelRequestContext { - let mut hasher = Sha256::new(); - for component in [ - Self::PROVIDER_PROMPT_CACHE_ROUTE_SCHEMA, - session_id, - model_binding_fingerprint, - effective_model_name, - prompt_scope, - ] { - hasher.update(component.as_bytes()); - hasher.update([0]); - } ModelRequestContext { - prompt_cache_route_key: Some(format!( - "bitfun-pc-v1-{}", - hex::encode(hasher.finalize()) - )), + prompt_cache_route_key: Some(prompt_cache_lineage_id.to_string()), model_binding_fingerprint: Some(model_binding_fingerprint.to_string()), } } @@ -2260,6 +2226,7 @@ impl ExecutionEngine { ai_client: Arc, request_messages: Vec, tool_definitions: Option>, + model_request_context: &ModelRequestContext, trace_config: Option, max_tries: usize, ) -> BitFunResult { @@ -2268,9 +2235,10 @@ impl ExecutionEngine { for attempt in 0..max_tries { let result = ai_client - .send_message_with_trace( + .send_message_with_trace_and_request_context( request_messages.clone(), tool_definitions.clone(), + Some(model_request_context.clone()), trace_config.clone(), ) .await; @@ -2363,6 +2331,7 @@ impl ExecutionEngine { input.ai_client, request_messages, input.tool_definitions.clone(), + input.model_request_context, input.trace_config, 2, ) @@ -2385,6 +2354,7 @@ impl ExecutionEngine { context_window: usize, compression_contract: Option, ai_client: Arc, + model_request_context: &ModelRequestContext, tool_definitions: &Option>, prepended_prompt_reminders: &PrependedPromptReminders, primary_supports_image_understanding: bool, @@ -2426,6 +2396,7 @@ impl ExecutionEngine { let summary_result = self .generate_compression_model_summary(CompressionModelSummaryInput { ai_client: ai_client.clone(), + model_request_context, runtime_messages: &plan.summary_request_messages, dialog_turn_id, workspace, @@ -2523,7 +2494,7 @@ impl ExecutionEngine { .get("original_user_input") .cloned() .unwrap_or_default(); - let (model_id, _) = self + let (model_id, model_binding_fingerprint) = self .resolve_model_id_for_turn( session, &context.agent_type, @@ -2598,6 +2569,10 @@ impl ExecutionEngine { }; Self::validate_frozen_model_contract(context).await?; Self::validate_frozen_reasoning_contract(context, ai_client.as_ref())?; + let model_request_context = Self::model_request_context( + session.effective_prompt_cache_lineage_id(), + &model_binding_fingerprint, + ); let primary_model_facts = Self::resolve_primary_model_context( &model_id, @@ -2696,6 +2671,7 @@ impl ExecutionEngine { Ok(CompressionRuntimeScaffold { ai_client, + model_request_context, tool_definitions, system_prompt_message: turn_prompt_scaffold.system_prompt_message, prepended_prompt_reminders: turn_prompt_scaffold.prepended_prompt_reminders, @@ -2743,6 +2719,7 @@ impl ExecutionEngine { before_pressure: TokenPressureSnapshot, context_window: usize, ai_client: Arc, + model_request_context: &ModelRequestContext, tool_definitions: &Option>, system_prompt_message: Message, prepended_prompt_reminders: &PrependedPromptReminders, @@ -2819,6 +2796,7 @@ impl ExecutionEngine { context_window, compression_contract, ai_client, + model_request_context, tool_definitions, prepended_prompt_reminders, primary_supports_image_understanding, @@ -3087,6 +3065,7 @@ impl ExecutionEngine { context_window, compression_contract, scaffold.ai_client.clone(), + &scaffold.model_request_context, &scaffold.tool_definitions, &scaffold.prepended_prompt_reminders, scaffold.primary_supports_image_understanding, @@ -3517,10 +3496,8 @@ impl ExecutionEngine { Self::validate_frozen_model_contract(&context).await?; Self::validate_frozen_reasoning_contract(&context, ai_client.as_ref())?; let model_request_context = Self::model_request_context( - &context.session_id, + session.effective_prompt_cache_lineage_id(), &model_binding_fingerprint, - &ai_client.config.model, - current_agent.as_ref(), ); // Primary model vision capability (tools + system prompt appendix; also used below for API message stripping). @@ -3934,6 +3911,7 @@ impl ExecutionEngine { token_pressure, context_window, ai_client.clone(), + &model_request_context, &tool_definitions, turn_prompt_scaffold.system_prompt_message.clone(), &turn_prompt_scaffold.prepended_prompt_reminders, @@ -4165,6 +4143,7 @@ impl ExecutionEngine { send_pressure, context_window, ai_client.clone(), + &model_request_context, &tool_definitions, turn_prompt_scaffold.system_prompt_message.clone(), &turn_prompt_scaffold.prepended_prompt_reminders, @@ -6456,38 +6435,26 @@ mod tests { } #[test] - fn provider_prompt_cache_route_key_is_stable_and_changes_with_scope() { - let first = ExecutionEngine::model_request_context_from_scope_key( - "session-1", - "binding-1", - "gpt-5.6-terra", - "scope-1", - ); - let retry = ExecutionEngine::model_request_context_from_scope_key( - "session-1", - "binding-1", - "gpt-5.6-terra", - "scope-1", + fn provider_prompt_cache_route_key_depends_only_on_lineage() { + let first = ExecutionEngine::model_request_context("session-1", "binding-1"); + let changed_binding = ExecutionEngine::model_request_context("session-1", "binding-2"); + let changed_lineage = ExecutionEngine::model_request_context("session-2", "binding-1"); + + assert_eq!(first.prompt_cache_route_key.as_deref(), Some("session-1")); + assert_eq!( + first.prompt_cache_route_key, + changed_binding.prompt_cache_route_key ); - let changed = ExecutionEngine::model_request_context_from_scope_key( - "session-1", - "binding-2", - "gpt-5.6-terra", - "scope-1", + assert_ne!( + first.prompt_cache_route_key, + changed_lineage.prompt_cache_route_key ); - - assert_eq!(first, retry); - assert_ne!(first, changed); - assert!(first - .prompt_cache_route_key - .as_deref() - .is_some_and(|key| key.starts_with("bitfun-pc-v1-"))); assert_eq!( first.model_binding_fingerprint.as_deref(), Some("binding-1") ); assert_eq!( - changed.model_binding_fingerprint.as_deref(), + changed_binding.model_binding_fingerprint.as_deref(), Some("binding-2") ); } diff --git a/src/crates/assembly/core/src/agentic/fork_agent/mod.rs b/src/crates/assembly/core/src/agentic/fork_agent/mod.rs index 4ff56441b6..e27c49e4ce 100644 --- a/src/crates/assembly/core/src/agentic/fork_agent/mod.rs +++ b/src/crates/assembly/core/src/agentic/fork_agent/mod.rs @@ -38,6 +38,13 @@ impl ForkAgentContextSnapshot { )) })?; + let mut session_config = parent_session.config.clone(); + session_config.prompt_cache_lineage_id = Some( + parent_session + .effective_prompt_cache_lineage_id() + .to_string(), + ); + Ok(Self { parent_session_id: parent_session.session_id.clone(), parent_agent_type: parent_session.agent_type.clone(), @@ -45,7 +52,7 @@ impl ForkAgentContextSnapshot { remote_connection_id: parent_session.config.remote_connection_id.clone(), remote_ssh_host: parent_session.config.remote_ssh_host.clone(), session_model_id: parent_session.config.model_id.clone(), - session_config: parent_session.config.clone(), + session_config, last_user_dialog_agent_type: parent_session.last_user_dialog_agent_type.clone(), last_submitted_agent_type: parent_session.last_submitted_agent_type.clone(), messages, @@ -116,6 +123,27 @@ mod tests { assert_eq!(child_config.remote_ssh_host.as_deref(), Some("prod-box")); assert_eq!(child_config.model_id.as_deref(), Some("primary")); assert_eq!(child_config.max_turns, 7); + assert_eq!( + child_config.prompt_cache_lineage_id.as_deref(), + Some(parent.session_id.as_str()) + ); + } + + #[test] + fn snapshot_preserves_an_existing_prompt_cache_lineage() { + let mut parent = parent_session(); + parent.config.prompt_cache_lineage_id = Some("root-lineage".to_string()); + + let snapshot = + ForkAgentContextSnapshot::from_parent_session(&parent, Vec::new()).expect("snapshot"); + + assert_eq!( + snapshot + .build_child_session_config(None) + .prompt_cache_lineage_id + .as_deref(), + Some("root-lineage") + ); } #[test] diff --git a/src/crates/assembly/core/src/agentic/persistence/session_branch.rs b/src/crates/assembly/core/src/agentic/persistence/session_branch.rs index 0b0d211997..97e6251950 100644 --- a/src/crates/assembly/core/src/agentic/persistence/session_branch.rs +++ b/src/crates/assembly/core/src/agentic/persistence/session_branch.rs @@ -71,10 +71,16 @@ impl PersistenceManager { format_branch_session_name(&branch_lineage.base_session_name, branch_lineage.ordinal); let target_agent_type = source_session.agent_type.clone(); + let mut target_config = source_session.config.clone(); + target_config.prompt_cache_lineage_id = Some( + source_session + .effective_prompt_cache_lineage_id() + .to_string(), + ); let mut target_session = Session::new( target_session_name.clone(), target_agent_type.clone(), - source_session.config.clone(), + target_config, ); target_session.created_by = None; target_session.kind = SessionKind::Standard; @@ -439,6 +445,14 @@ mod tests { assert_ne!(result.session_id, source_session.session_id); assert_eq!(result.session_name, "Source Title (1)"); assert_eq!(result.agent_type, "agentic"); + let branched_session = manager + .load_session(workspace.path(), &result.session_id) + .await + .expect("branched session should load"); + assert_eq!( + branched_session.config.prompt_cache_lineage_id.as_deref(), + Some(source_session.session_id.as_str()) + ); let branched_turns = manager .load_session_turns(workspace.path(), &result.session_id) diff --git a/src/crates/execution/agent-runtime/src/session.rs b/src/crates/execution/agent-runtime/src/session.rs index c8357eca73..11ccd1acfd 100644 --- a/src/crates/execution/agent-runtime/src/session.rs +++ b/src/crates/execution/agent-runtime/src/session.rs @@ -135,6 +135,18 @@ impl Session { last_activity_at: now, } } + + /// Stable routing identity for provider-side prompt-prefix caches. + /// + /// Legacy and independent sessions use their own session ID. Derived + /// sessions that preserve a parent prompt prefix persist the parent's + /// effective lineage in `SessionConfig`. + pub fn effective_prompt_cache_lineage_id(&self) -> &str { + self.config + .prompt_cache_lineage_id + .as_deref() + .unwrap_or(&self.session_id) + } } impl From for bitfun_runtime_ports::AgentSessionCreateResult { @@ -212,6 +224,11 @@ pub struct SessionConfig { /// Mutable sessions leave this unset and continue to resolve selectors. #[serde(default, skip_serializing_if = "Option::is_none")] pub model_binding_fingerprint: Option, + /// Stable provider-cache lineage shared only by sessions that preserve an + /// exact prompt prefix. `None` keeps legacy and independent sessions scoped + /// to their own session ID. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub prompt_cache_lineage_id: Option, /// Durable owner of the logical main-agent route. External ownership is /// revalidated for every turn and never falls back by name alone. #[serde(default, skip_serializing_if = "is_local_agent_route_owner")] @@ -251,6 +268,7 @@ impl Default for SessionConfig { continuation_policy: SessionContinuationPolicy::default(), model_binding_policy: SessionModelBindingPolicy::default(), model_binding_fingerprint: None, + prompt_cache_lineage_id: None, agent_route_owner: SessionAgentRouteOwner::Local, } } @@ -372,6 +390,37 @@ mod tests { let serialized = serde_json::to_value(SessionConfig::default()).expect("serialize"); assert!(serialized.get("permission_mode").is_none()); } + + #[test] + fn prompt_cache_lineage_defaults_to_the_session_id_and_round_trips() { + let mut legacy_value = serde_json::to_value(SessionConfig::default()).expect("serialize"); + legacy_value + .as_object_mut() + .expect("config should serialize as an object") + .remove("prompt_cache_lineage_id"); + let legacy_config: SessionConfig = + serde_json::from_value(legacy_value).expect("deserialize legacy config"); + let mut session = Session::new("Session".to_string(), "agentic".to_string(), legacy_config); + assert!(session.config.prompt_cache_lineage_id.is_none()); + assert_eq!( + session.effective_prompt_cache_lineage_id(), + session.session_id + ); + + session.config.prompt_cache_lineage_id = Some("root-session".to_string()); + assert_eq!(session.effective_prompt_cache_lineage_id(), "root-session"); + + let serialized = serde_json::to_value(&session.config).expect("serialize"); + assert_eq!( + serialized["prompt_cache_lineage_id"], + serde_json::json!("root-session") + ); + let restored: SessionConfig = serde_json::from_value(serialized).expect("deserialize"); + assert_eq!( + restored.prompt_cache_lineage_id.as_deref(), + Some("root-session") + ); + } use bitfun_core_types::{ SessionExecutionTarget, SessionExecutionTargetKind, WorktreeLifecycle, }; From 29da0207e79b09e32fcac72b42e7e1741c344c4c Mon Sep 17 00:00:00 2001 From: wsp Date: Mon, 17 Aug 2026 01:37:21 +0800 Subject: [PATCH 4/5] fix(responses): preserve BTW prompt prefix Route BTW turns through the Desktop UI submission policy so the runtime prompt retains the Chat Image Display context. Pass the policy from the desktop adapter into the coordinator and add regression coverage for the BTW output surface classification. --- src/apps/desktop/src/api/btw_api.rs | 22 ++++++++++++++++++- .../src/agentic/coordination/coordinator.rs | 3 ++- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/src/apps/desktop/src/api/btw_api.rs b/src/apps/desktop/src/api/btw_api.rs index c636011db6..729311c996 100644 --- a/src/apps/desktop/src/api/btw_api.rs +++ b/src/apps/desktop/src/api/btw_api.rs @@ -12,9 +12,15 @@ use tauri::State; use crate::api::app_state::AppState; -use bitfun_core::agentic::coordination::ConversationCoordinator; +use bitfun_core::agentic::coordination::{ + ConversationCoordinator, DialogSubmissionPolicy, DialogTriggerSource, +}; use bitfun_core::agentic::image_analysis::ImageContextData; +fn desktop_btw_submission_policy() -> DialogSubmissionPolicy { + DialogSubmissionPolicy::for_source(DialogTriggerSource::DesktopUi) +} + #[derive(Debug, Clone, Deserialize)] #[serde(rename_all = "camelCase")] pub struct BtwAskStreamRequest { @@ -105,6 +111,7 @@ pub async fn btw_ask_stream( &child_session_id, child_session_name.as_deref(), &request.question, + desktop_btw_submission_policy(), model_id.as_deref(), image_contexts, request.parent_dialog_turn_id.as_deref(), @@ -150,3 +157,16 @@ pub async fn btw_ask_stream( Ok(BtwAskStreamResponse { ok: true }) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn btw_turns_use_the_desktop_chat_output_surface() { + assert_eq!( + desktop_btw_submission_policy(), + DialogSubmissionPolicy::for_source(DialogTriggerSource::DesktopUi) + ); + } +} diff --git a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs index 8709105d68..99d1154668 100644 --- a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs +++ b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs @@ -10393,6 +10393,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet child_session_id: &str, child_session_name: Option<&str>, question: &str, + submission_policy: DialogSubmissionPolicy, model_id: Option<&str>, image_contexts: Option>, parent_dialog_turn_id: Option<&str>, @@ -10478,7 +10479,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet child_session.config.workspace_path.clone(), child_session.config.remote_connection_id.clone(), child_session.config.remote_ssh_host.clone(), - DialogSubmissionPolicy::for_source(DialogTriggerSource::DesktopApi), + submission_policy, Some(user_message_metadata), prepended_messages, true, From 31efcae9a37cf7c8ad8b34416d760df86185e60e Mon Sep 17 00:00:00 2001 From: wsp Date: Mon, 17 Aug 2026 01:51:21 +0800 Subject: [PATCH 5/5] fix(responses): preserve reasoning across model switches Remove model-binding fingerprints from persisted response replay state and request context. Forward valid encrypted reasoning items without local model filtering, allowing the Responses service to determine compatibility while retaining atomic fallback for malformed replay layouts. --- .../src/providers/openai/codex_chatgpt.rs | 11 +--- .../src/providers/openai/message_converter.rs | 66 ++----------------- .../src/providers/openai/responses.rs | 17 +---- .../assembly/core/src/agentic/core/message.rs | 11 +--- .../src/agentic/execution/execution_engine.rs | 38 ++++------- .../src/agentic/execution/round_executor.rs | 15 +---- src/crates/contracts/core-types/src/ai.rs | 6 -- 7 files changed, 26 insertions(+), 138 deletions(-) diff --git a/src/crates/adapters/ai-adapters/src/providers/openai/codex_chatgpt.rs b/src/crates/adapters/ai-adapters/src/providers/openai/codex_chatgpt.rs index bc6c6e85b5..5d27afaca4 100644 --- a/src/crates/adapters/ai-adapters/src/providers/openai/codex_chatgpt.rs +++ b/src/crates/adapters/ai-adapters/src/providers/openai/codex_chatgpt.rs @@ -24,7 +24,7 @@ use crate::client::{AIClient, StreamResponse}; use crate::providers::shared; use crate::stream::handle_responses_stream; use crate::trace::ModelExchangeTraceConfig; -use crate::types::{Message, ModelRequestContext, ReasoningPresetAction, ToolDefinition}; +use crate::types::{Message, ReasoningPresetAction, ToolDefinition}; use anyhow::Result; use log::debug; use serde_json::{json, Value}; @@ -184,7 +184,6 @@ pub(crate) async fn send_stream( extra_body: Option, max_tries: usize, trace: Option, - request_context: Option, ) -> Result { let url = client.config.request_url.clone(); debug!( @@ -192,14 +191,8 @@ pub(crate) async fn send_stream( client.config.model, url, max_tries ); - let model_binding_fingerprint = request_context - .as_ref() - .and_then(|context| context.model_binding_fingerprint.as_deref()); let (instructions, response_input) = - OpenAIMessageConverter::convert_messages_to_responses_input_with_context( - messages, - model_binding_fingerprint, - ); + OpenAIMessageConverter::convert_messages_to_responses_input(messages); let tools_flat = common::convert_tools_flat(tools); let request_body = try_build_request_body(client, instructions, response_input, tools_flat, extra_body)?; diff --git a/src/crates/adapters/ai-adapters/src/providers/openai/message_converter.rs b/src/crates/adapters/ai-adapters/src/providers/openai/message_converter.rs index ccf6fff9e5..b5ea06ebae 100644 --- a/src/crates/adapters/ai-adapters/src/providers/openai/message_converter.rs +++ b/src/crates/adapters/ai-adapters/src/providers/openai/message_converter.rs @@ -12,13 +12,6 @@ pub struct OpenAIMessageConverter; impl OpenAIMessageConverter { pub fn convert_messages_to_responses_input( messages: Vec, - ) -> (Option, Vec) { - Self::convert_messages_to_responses_input_with_context(messages, None) - } - - pub fn convert_messages_to_responses_input_with_context( - messages: Vec, - model_binding_fingerprint: Option<&str>, ) -> (Option, Vec) { let mut instructions = Vec::new(); let mut input = Vec::new(); @@ -37,9 +30,7 @@ impl OpenAIMessageConverter { } } "assistant" => { - if let Some(replay_items) = - Self::convert_assistant_replay_items(&msg, model_binding_fingerprint) - { + if let Some(replay_items) = Self::convert_assistant_replay_items(&msg) { input.extend(replay_items); continue; } @@ -91,16 +82,9 @@ impl OpenAIMessageConverter { (instructions, input) } - fn convert_assistant_replay_items( - msg: &Message, - model_binding_fingerprint: Option<&str>, - ) -> Option> { + fn convert_assistant_replay_items(msg: &Message) -> Option> { let replay = msg.model_response_replay.as_ref()?; - let current_fingerprint = model_binding_fingerprint.filter(|value| !value.is_empty())?; - if replay.protocol != OPENAI_RESPONSES_REPLAY_PROTOCOL - || replay.model_binding_fingerprint != current_fingerprint - || replay.items.is_empty() - { + if replay.protocol != OPENAI_RESPONSES_REPLAY_PROTOCOL || replay.items.is_empty() { return None; } @@ -554,7 +538,6 @@ mod tests { content: Option<&str>, tool_calls: Vec, items: Vec, - fingerprint: &str, ) -> Message { Message { role: "assistant".to_string(), @@ -568,7 +551,6 @@ mod tests { tool_image_attachments: None, model_response_replay: Some(ModelResponseReplay { protocol: "openai_responses".to_string(), - model_binding_fingerprint: fingerprint.to_string(), items, }), } @@ -585,7 +567,7 @@ mod tests { } #[test] - fn replays_reasoning_before_final_assistant_message_when_fingerprint_matches() { + fn replays_reasoning_before_final_assistant_message() { let message = assistant_with_replay( Some("done"), vec![], @@ -593,13 +575,9 @@ mod tests { opaque_reasoning("rs_1", "opaque_1"), ModelResponseReplayItem::AssistantMessage, ], - "binding-1", ); - let (_, input) = OpenAIMessageConverter::convert_messages_to_responses_input_with_context( - vec![message], - Some("binding-1"), - ); + let (_, input) = OpenAIMessageConverter::convert_messages_to_responses_input(vec![message]); assert_eq!(input.len(), 2); assert_eq!(input[0]["type"], json!("reasoning")); @@ -637,13 +615,9 @@ mod tests { call_id: "call_1".to_string(), }, ], - "binding-1", ); - let (_, input) = OpenAIMessageConverter::convert_messages_to_responses_input_with_context( - vec![message], - Some("binding-1"), - ); + let (_, input) = OpenAIMessageConverter::convert_messages_to_responses_input(vec![message]); assert_eq!( input @@ -657,28 +631,6 @@ mod tests { assert_eq!(input[3]["call_id"], json!("call_1")); } - #[test] - fn fingerprint_mismatch_uses_ordinary_responses_conversion() { - let message = assistant_with_replay( - Some("done"), - vec![], - vec![ - opaque_reasoning("rs_1", "opaque_1"), - ModelResponseReplayItem::AssistantMessage, - ], - "binding-old", - ); - - let (_, input) = OpenAIMessageConverter::convert_messages_to_responses_input_with_context( - vec![message], - Some("binding-new"), - ); - - assert_eq!(input.len(), 1); - assert_eq!(input[0]["type"], json!("message")); - assert!(input[0].get("encrypted_content").is_none()); - } - #[test] fn invalid_replay_layout_falls_back_atomically() { let message = assistant_with_replay( @@ -703,13 +655,9 @@ mod tests { call_id: "call_1".to_string(), }, ], - "binding-1", ); - let (_, input) = OpenAIMessageConverter::convert_messages_to_responses_input_with_context( - vec![message], - Some("binding-1"), - ); + let (_, input) = OpenAIMessageConverter::convert_messages_to_responses_input(vec![message]); assert_eq!(input.len(), 2); assert!(input.iter().all(|item| item["type"] == "function_call")); diff --git a/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs b/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs index e622a66ca2..dca43fb662 100644 --- a/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs +++ b/src/crates/adapters/ai-adapters/src/providers/openai/responses.rs @@ -248,13 +248,7 @@ pub(crate) async fn send_stream( // self-contained so the standard Responses path stays untouched. if super::codex_chatgpt::is_codex_chatgpt_endpoint(&client.config.request_url) { return super::codex_chatgpt::send_stream( - client, - messages, - tools, - extra_body, - max_tries, - trace, - request_context, + client, messages, tools, extra_body, max_tries, trace, ) .await; } @@ -265,14 +259,8 @@ pub(crate) async fn send_stream( client.config.model, client.config.request_url, max_tries ); - let model_binding_fingerprint = request_context - .as_ref() - .and_then(|context| context.model_binding_fingerprint.as_deref()); let (instructions, response_input) = - OpenAIMessageConverter::convert_messages_to_responses_input_with_context( - messages, - model_binding_fingerprint, - ); + OpenAIMessageConverter::convert_messages_to_responses_input(messages); let openai_tools = common::convert_tools_flat(tools); let request_body = try_build_request_body_with_context( client, @@ -387,7 +375,6 @@ mod tests { let client = test_client(); let request_context = ModelRequestContext { prompt_cache_route_key: Some("lineage-1".to_string()), - model_binding_fingerprint: Some("binding-1".to_string()), }; let request_body = build_request_body_with_context( &client, diff --git a/src/crates/assembly/core/src/agentic/core/message.rs b/src/crates/assembly/core/src/agentic/core/message.rs index 0e7463f728..bbdff6cc5e 100644 --- a/src/crates/assembly/core/src/agentic/core/message.rs +++ b/src/crates/assembly/core/src/agentic/core/message.rs @@ -793,7 +793,6 @@ mod tests { let message = Message::assistant("done".to_string()).with_model_response_replay(Some( ModelResponseReplay { protocol: "openai_responses".to_string(), - model_binding_fingerprint: "binding-1".to_string(), items: vec![ModelResponseReplayItem::OpaqueReasoning { item_id: Some("rs_1".to_string()), summary: vec![], @@ -810,7 +809,6 @@ mod tests { .expect("restored replay"); assert_eq!(replay.protocol, "openai_responses"); - assert_eq!(replay.model_binding_fingerprint, "binding-1"); assert!(matches!( &replay.items[0], ModelResponseReplayItem::OpaqueReasoning { opaque_state, .. } @@ -836,19 +834,12 @@ mod tests { let message = Message::assistant("done".to_string()).with_model_response_replay(Some( ModelResponseReplay { protocol: "openai_responses".to_string(), - model_binding_fingerprint: "binding-1".to_string(), items: vec![ModelResponseReplayItem::AssistantMessage], }, )); let ai_message = AIMessage::from(message); - assert_eq!( - ai_message - .model_response_replay - .as_ref() - .map(|replay| replay.model_binding_fingerprint.as_str()), - Some("binding-1") - ); + assert!(ai_message.model_response_replay.is_some()); } #[test] diff --git a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs index 15cc6132ef..f4fd8fd7cb 100644 --- a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs +++ b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs @@ -537,13 +537,9 @@ impl ExecutionEngine { const FINALIZE_USER_FOLLOWUP: &'static str = "Provide a final answer. You MUST not call any tools."; - fn model_request_context( - prompt_cache_lineage_id: &str, - model_binding_fingerprint: &str, - ) -> ModelRequestContext { + fn model_request_context(prompt_cache_lineage_id: &str) -> ModelRequestContext { ModelRequestContext { prompt_cache_route_key: Some(prompt_cache_lineage_id.to_string()), - model_binding_fingerprint: Some(model_binding_fingerprint.to_string()), } } @@ -2494,7 +2490,7 @@ impl ExecutionEngine { .get("original_user_input") .cloned() .unwrap_or_default(); - let (model_id, model_binding_fingerprint) = self + let (model_id, _) = self .resolve_model_id_for_turn( session, &context.agent_type, @@ -2569,10 +2565,8 @@ impl ExecutionEngine { }; Self::validate_frozen_model_contract(context).await?; Self::validate_frozen_reasoning_contract(context, ai_client.as_ref())?; - let model_request_context = Self::model_request_context( - session.effective_prompt_cache_lineage_id(), - &model_binding_fingerprint, - ); + let model_request_context = + Self::model_request_context(session.effective_prompt_cache_lineage_id()); let primary_model_facts = Self::resolve_primary_model_context( &model_id, @@ -3417,7 +3411,7 @@ impl ExecutionEngine { } } - let (model_id, model_binding_fingerprint) = self + let (model_id, _) = self .resolve_model_id_for_turn( &session, &agent_type, @@ -3495,10 +3489,8 @@ impl ExecutionEngine { }; Self::validate_frozen_model_contract(&context).await?; Self::validate_frozen_reasoning_contract(&context, ai_client.as_ref())?; - let model_request_context = Self::model_request_context( - session.effective_prompt_cache_lineage_id(), - &model_binding_fingerprint, - ); + let model_request_context = + Self::model_request_context(session.effective_prompt_cache_lineage_id()); // Primary model vision capability (tools + system prompt appendix; also used below for API message stripping). let primary_model_facts = Self::resolve_primary_model_context( @@ -6436,27 +6428,19 @@ mod tests { #[test] fn provider_prompt_cache_route_key_depends_only_on_lineage() { - let first = ExecutionEngine::model_request_context("session-1", "binding-1"); - let changed_binding = ExecutionEngine::model_request_context("session-1", "binding-2"); - let changed_lineage = ExecutionEngine::model_request_context("session-2", "binding-1"); + let first = ExecutionEngine::model_request_context("session-1"); + let same_lineage = ExecutionEngine::model_request_context("session-1"); + let changed_lineage = ExecutionEngine::model_request_context("session-2"); assert_eq!(first.prompt_cache_route_key.as_deref(), Some("session-1")); assert_eq!( first.prompt_cache_route_key, - changed_binding.prompt_cache_route_key + same_lineage.prompt_cache_route_key ); assert_ne!( first.prompt_cache_route_key, changed_lineage.prompt_cache_route_key ); - assert_eq!( - first.model_binding_fingerprint.as_deref(), - Some("binding-1") - ); - assert_eq!( - changed_binding.model_binding_fingerprint.as_deref(), - Some("binding-2") - ); } fn command_result(tool_name: &str, success: bool, exit_code: Option) -> Message { diff --git a/src/crates/assembly/core/src/agentic/execution/round_executor.rs b/src/crates/assembly/core/src/agentic/execution/round_executor.rs index fa68520747..6172f06a7b 100644 --- a/src/crates/assembly/core/src/agentic/execution/round_executor.rs +++ b/src/crates/assembly/core/src/agentic/execution/round_executor.rs @@ -220,19 +220,10 @@ impl RoundExecutor { .map(Into::into) } - fn bound_model_response_replay( - stream_result: &StreamResult, - context: &RoundContext, - ) -> Option { + fn model_response_replay(stream_result: &StreamResult) -> Option { let capture = stream_result.model_response_replay.as_ref()?; - let model_binding_fingerprint = context - .model_request_context - .model_binding_fingerprint - .as_ref()? - .clone(); Some(ModelResponseReplay { protocol: capture.protocol.clone(), - model_binding_fingerprint, items: capture.items.clone(), }) } @@ -982,7 +973,7 @@ impl RoundExecutor { }; let parsed_memory_citation = Self::parsed_memory_citation_from_stream_result(&stream_result); - let model_response_replay = Self::bound_model_response_replay(&stream_result, &context); + let model_response_replay = Self::model_response_replay(&stream_result); let (clean_text, _) = strip_bitfun_memory_citations(&stream_result.full_text); let assistant_message = Message::assistant_with_reasoning(reasoning, clean_text, vec![]) @@ -1206,7 +1197,7 @@ impl RoundExecutor { }; let parsed_memory_citation = Self::parsed_memory_citation_from_stream_result(&stream_result); - let model_response_replay = Self::bound_model_response_replay(&stream_result, &context); + let model_response_replay = Self::model_response_replay(&stream_result); let (clean_text, _) = strip_bitfun_memory_citations(&stream_result.full_text); let assistant_message = Message::assistant_with_reasoning(reasoning, clean_text, tool_calls.clone()) diff --git a/src/crates/contracts/core-types/src/ai.rs b/src/crates/contracts/core-types/src/ai.rs index c67294203e..dd5bc22e14 100644 --- a/src/crates/contracts/core-types/src/ai.rs +++ b/src/crates/contracts/core-types/src/ai.rs @@ -653,17 +653,11 @@ pub struct AIConfig { pub struct ModelRequestContext { /// Stable, opaque routing identity for provider-side prompt-prefix caches. pub prompt_cache_route_key: Option, - /// Fingerprint of the resolved provider/model/endpoint/request binding. - /// - /// Provider adapters use this only to decide whether opaque response state - /// from an earlier round is compatible with the current request. - pub model_binding_fingerprint: Option, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct ModelResponseReplay { pub protocol: String, - pub model_binding_fingerprint: String, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub items: Vec, }