From 08c48e37007683d366f648d4c40ec17ba5009e98 Mon Sep 17 00:00:00 2001 From: felixxia-oai Date: Fri, 14 Aug 2026 11:48:31 +0000 Subject: [PATCH] Reuse compatible parent compactions in Guardian V2 (#38564) ## Why Guardian V2 classification requests should retain compacted conversation context only when the parent and classifier model configurations can share the encrypted compaction. ## What changed - Resolve the classifier model's compaction compatibility hash when initializing its sampler. - Add the latest encrypted parent compaction to classification input when both models have the same non-empty compatibility hash. - Ignore missing, empty, or unidentified compactions, and omit parent context when compatibility cannot be established. ## Testing - Cover latest-compaction selection and invalid compaction data. - Cover matching, mismatched, missing, empty, and changed compatibility hashes. GitOrigin-RevId: f0fcb09fd028cd2f397fdbc3d541b4bc464b0beb --- codex-rs/ext/guardian-v2/src/extension.rs | 47 +++++ .../ext/guardian-v2/src/extension_tests.rs | 176 ++++++++++++++++++ codex-rs/ext/guardian-v2/src/sampler.rs | 74 +++++--- codex-rs/ext/guardian-v2/src/sampler_tests.rs | 77 ++++++++ 4 files changed, 346 insertions(+), 28 deletions(-) diff --git a/codex-rs/ext/guardian-v2/src/extension.rs b/codex-rs/ext/guardian-v2/src/extension.rs index 944a54336a..66799ea388 100644 --- a/codex-rs/ext/guardian-v2/src/extension.rs +++ b/codex-rs/ext/guardian-v2/src/extension.rs @@ -7,6 +7,7 @@ use codex_extension_api::ExtensionEventSink; use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ExtensionWarning; +use codex_extension_api::ResponseItem; use codex_extension_api::ThreadLifecycleContributor; use codex_extension_api::ThreadOriginator; use codex_extension_api::ThreadStartInput; @@ -20,6 +21,7 @@ use codex_login::AgentIdentityAuthPolicy; use codex_login::AuthManager; use codex_model_provider::create_model_provider; use codex_protocol::ThreadId; +use codex_protocol::openai_models::ModelInfo; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::security_risk::SecurityRiskScore; use serde_json::json; @@ -27,6 +29,7 @@ use serde_json::json; use crate::LunaSampler; use crate::LunaSamplerConfig; use crate::LunaSamplingRequest; +use crate::sampler::MODEL; use crate::transcript::TranscriptConfig; const CLASSIFIER_INSTRUCTIONS: &str = "Classify the risk of the assistant tool action. Consider \ @@ -52,6 +55,15 @@ impl ThreadLifecycleContributor for GuardianV2Extension { } let thread_id = input.thread_store.level_id().to_string(); + let luna_compaction_hash = if let Some(thread_manager) = self.thread_manager.upgrade() { + thread_manager + .get_models_manager() + .get_model_info(MODEL, &input.config.to_models_manager_config()) + .await + .comp_hash + } else { + None + }; let sampler = LunaSampler::connect(LunaSamplerConfig { provider: create_model_provider( input.config.model_provider.clone(), @@ -71,6 +83,7 @@ impl ThreadLifecycleContributor for GuardianV2Extension { .get::() .map(|originator| originator.0.clone()), service_tier: input.config.service_tier.clone(), + luna_compaction_hash, }) .await; @@ -99,9 +112,14 @@ impl ToolLifecycleContributor for GuardianV2Extension { let turn_id = input.turn_id.to_owned(); let tool_name = input.tool_name.to_string(); let payload = input.payload.clone(); + let parent_compaction_hash = input + .thread_store + .get::() + .and_then(|model_info| model_info.comp_hash.clone()); let conversation_history = Arc::clone(&input.conversation_history); tokio::spawn(async move { + let parent_compaction = encrypted_parent_compaction(conversation_history.items()); let transcript = TranscriptConfig::default().build(conversation_history.items()); drop(conversation_history); let arguments = match payload { @@ -142,6 +160,8 @@ impl ToolLifecycleContributor for GuardianV2Extension { .sample(LunaSamplingRequest { instructions: CLASSIFIER_INSTRUCTIONS.to_owned(), input: classification_input, + parent_compaction, + parent_compaction_hash, output_schema: json!({ "type": "object", "properties": { @@ -215,6 +235,33 @@ impl ToolLifecycleContributor for GuardianV2Extension { } } +fn encrypted_parent_compaction<'a>( + items: impl Iterator, +) -> Option { + let item = items + .filter(|item| { + matches!( + item, + ResponseItem::Compaction { .. } | ResponseItem::ContextCompaction { .. } + ) + }) + .last()?; + + match item { + ResponseItem::Compaction { + id: Some(_), + encrypted_content, + .. + } if !encrypted_content.is_empty() => Some(item.clone()), + ResponseItem::ContextCompaction { + id: Some(_), + encrypted_content: Some(encrypted_content), + .. + } if !encrypted_content.is_empty() => Some(item.clone()), + _ => None, + } +} + /// Installs feature-gated Guardian V2 tool classification for each thread. pub fn install( registry: &mut ExtensionRegistryBuilder, diff --git a/codex-rs/ext/guardian-v2/src/extension_tests.rs b/codex-rs/ext/guardian-v2/src/extension_tests.rs index 352a67ce69..fdcadbbd98 100644 --- a/codex-rs/ext/guardian-v2/src/extension_tests.rs +++ b/codex-rs/ext/guardian-v2/src/extension_tests.rs @@ -16,6 +16,7 @@ use codex_history::RolloutItem; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderInfo; +use codex_protocol::ResponseItemId; use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ReasoningItemReasoningSummary; @@ -30,6 +31,9 @@ use core_test_support::test_codex::test_codex; use pretty_assertions::assert_eq; use serde_json::json; +use super::encrypted_parent_compaction; +use crate::sampler::MODEL; + struct TestConversationHistory(Vec); impl ConversationHistorySnapshot for TestConversationHistory { @@ -38,6 +42,73 @@ impl ConversationHistorySnapshot for TestConversationHistory { } } +#[test] +fn encrypted_parent_compaction_preserves_the_latest_valid_item() { + let older = ResponseItem::Compaction { + id: Some(ResponseItemId::from_server("cmp_older".to_owned())), + encrypted_content: "older encrypted summary".to_owned(), + internal_chat_message_metadata_passthrough: None, + }; + let latest = ResponseItem::ContextCompaction { + id: Some(ResponseItemId::from_server("cmp_latest".to_owned())), + encrypted_content: Some("latest encrypted summary".to_owned()), + internal_chat_message_metadata_passthrough: None, + }; + + assert_eq!( + encrypted_parent_compaction([&older, &latest].into_iter()), + Some(latest.clone()) + ); + assert_eq!( + encrypted_parent_compaction([&latest, &older].into_iter()), + Some(older) + ); +} + +#[test] +fn encrypted_parent_compaction_rejects_invalid_latest_item() { + let older = ResponseItem::Compaction { + id: Some(ResponseItemId::from_server("cmp_older".to_owned())), + encrypted_content: "older encrypted summary".to_owned(), + internal_chat_message_metadata_passthrough: None, + }; + let invalid = [ + ResponseItem::Compaction { + id: None, + encrypted_content: "encrypted summary without an ID".to_owned(), + internal_chat_message_metadata_passthrough: None, + }, + ResponseItem::Compaction { + id: Some(ResponseItemId::from_server("cmp_empty".to_owned())), + encrypted_content: String::new(), + internal_chat_message_metadata_passthrough: None, + }, + ResponseItem::ContextCompaction { + id: None, + encrypted_content: Some("encrypted context without an ID".to_owned()), + internal_chat_message_metadata_passthrough: None, + }, + ResponseItem::ContextCompaction { + id: Some(ResponseItemId::from_server("cmp_missing".to_owned())), + encrypted_content: None, + internal_chat_message_metadata_passthrough: None, + }, + ResponseItem::ContextCompaction { + id: Some(ResponseItemId::from_server("cmp_empty".to_owned())), + encrypted_content: Some(String::new()), + internal_chat_message_metadata_passthrough: None, + }, + ]; + + for latest in &invalid { + assert_eq!( + encrypted_parent_compaction([&older, latest].into_iter()), + None, + "an unusable latest summary must not resurrect older context" + ); + } +} + async fn sample_conversation_history( conversation_history: Vec, ) -> Result<(serde_json::Value, TestCodex)> { @@ -301,3 +372,108 @@ async fn contributor_sends_compacted_conversation_history_to_luna() -> Result<() Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn contributor_reuses_the_latest_compatible_parent_compaction() -> Result<()> { + skip_if_no_network!(Ok(())); + + let thread_server = responses::start_mock_server().await; + let test = test_codex().build_with_auto_env(&thread_server).await?; + let events = vec![ + ev_assistant_message("sample", r#"{"scores":{"action_risk":0.25}}"#), + ev_completed("response-1"), + ]; + let server = responses::start_websocket_server(vec![Vec::new(), vec![events]]).await; + let provider_info = ModelProviderInfo::create_openai_provider(Some(format!( + "http://{}/v1", + server.uri().trim_start_matches("ws://") + ))); + let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test-api-key")); + let mut config = test.config.clone(); + config.model_provider = provider_info; + config.features.enable(Feature::GuardianV2)?; + let parent_model = test + .thread_manager + .get_models_manager() + .get_model_info(MODEL, &config.to_models_manager_config()) + .await; + let mut builder = ExtensionRegistryBuilder::new(); + crate::install( + &mut builder, + auth_manager, + Arc::downgrade(&test.thread_manager), + ); + let registry = builder.build(); + let session_store = ExtensionData::new("session-1"); + let thread_store = test.codex.thread_extension_data(); + thread_store.insert(parent_model); + registry.thread_lifecycle_contributors()[0] + .on_thread_start(ThreadStartInput { + config: &config, + session_source: &SessionSource::Exec, + persistent_thread_state_available: false, + environments: &[], + mcp_resource_client: None, + extension_metrics: None, + session_store: &session_store, + thread_store, + }) + .await; + let turn_store = ExtensionData::new("turn-1"); + let tool_name = ToolName::plain("read_file"); + let tool_payload = ToolPayload::Function { + arguments: r#"{"path":"README.md"}"#.to_owned(), + }; + let latest_compaction = ResponseItem::ContextCompaction { + id: Some(ResponseItemId::from_server("cmp_latest".to_owned())), + encrypted_content: Some("latest encrypted parent summary".to_owned()), + internal_chat_message_metadata_passthrough: None, + }; + let conversation_history = TestConversationHistory(vec![ + ResponseItem::Compaction { + id: Some(ResponseItemId::from_server("cmp_old".to_owned())), + encrypted_content: "old encrypted parent summary".to_owned(), + internal_chat_message_metadata_passthrough: None, + }, + latest_compaction.clone(), + ResponseItem::Message { + id: None, + role: "user".to_owned(), + content: vec![ContentItem::InputText { + text: "Inspect the repository guidelines.".to_owned(), + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + }, + ]); + + registry.tool_lifecycle_contributors()[0] + .on_tool_start(ToolStartInput { + session_store: &session_store, + thread_store, + turn_store: &turn_store, + turn_id: "turn-1", + call_id: "call-1", + tool_name: &tool_name, + payload: &tool_payload, + conversation_history: Arc::new(conversation_history), + source: ToolCallSource::Direct, + }) + .await; + + let request = tokio::time::timeout( + Duration::from_secs(5), + server.wait_for_request(/*connection_index*/ 1, /*request_index*/ 0), + ) + .await? + .body_json(); + assert_eq!(request["input"][0]["type"], "additional_tools"); + assert_eq!(request["input"][1]["role"], "developer"); + assert_eq!( + request["input"][2], + serde_json::to_value(latest_compaction)? + ); + assert_eq!(request["input"][3]["role"], "user"); + + Ok(()) +} diff --git a/codex-rs/ext/guardian-v2/src/sampler.rs b/codex-rs/ext/guardian-v2/src/sampler.rs index 3978019eb5..f806b6aa3b 100644 --- a/codex-rs/ext/guardian-v2/src/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/sampler.rs @@ -33,7 +33,7 @@ use thiserror::Error; use tokio::sync::OwnedSemaphorePermit; use tokio::sync::Semaphore; -const MODEL: &str = "gpt-5.6-luna"; +pub(crate) const MODEL: &str = "gpt-5.6-luna"; const MAX_OUTPUT_BYTES: usize = 8 * 1024; const INITIAL_WEBSOCKET_CONNECTIONS: usize = 2; const MAX_WEBSOCKET_CONNECTIONS: usize = 8; @@ -60,6 +60,8 @@ pub struct LunaSamplerConfig { pub originator: Option, /// Optional inference service tier. pub service_tier: Option, + /// Luna model's host-resolved encrypted-compaction compatibility hash. + pub luna_compaction_hash: Option, } /// One tool-less structured Luna request over an already-open connection. @@ -68,6 +70,10 @@ pub struct LunaSamplingRequest { pub instructions: String, /// Ordered untrusted input entries that the model should classify. pub input: Vec, + /// Opaque parent compaction to reuse only for compatible model configurations. + pub parent_compaction: Option, + /// Current parent model's encrypted-compaction compatibility hash. + pub parent_compaction_hash: Option, /// Strict JSON schema constraining the model response. pub output_schema: Value, /// Reasoning budget explicitly selected for this request. @@ -263,36 +269,48 @@ impl LunaSampler { ("turn_id".to_owned(), request.turn_id), (RESPONSES_LITE_METADATA_KEY.to_owned(), "true".to_owned()), ]); + let mut input = vec![ + ResponseItem::AdditionalTools { + id: None, + role: "developer".to_owned(), + tools: Vec::new(), + }, + ResponseItem::Message { + id: None, + role: "developer".to_owned(), + content: vec![ContentItem::InputText { + text: request.instructions, + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + }, + ]; + if request + .parent_compaction_hash + .as_deref() + .zip(self.config.luna_compaction_hash.as_deref()) + .is_some_and(|(parent_hash, luna_hash)| { + !parent_hash.is_empty() && parent_hash == luna_hash + }) + && let Some(parent_compaction) = request.parent_compaction + { + input.push(parent_compaction); + } + input.push(ResponseItem::Message { + id: None, + role: "user".to_owned(), + content: request + .input + .into_iter() + .map(|text| ContentItem::InputText { text }) + .collect(), + phase: None, + internal_chat_message_metadata_passthrough: None, + }); let request = ResponsesApiRequest { model: MODEL.to_owned(), instructions: String::new(), - input: vec![ - ResponseItem::AdditionalTools { - id: None, - role: "developer".to_owned(), - tools: Vec::new(), - }, - ResponseItem::Message { - id: None, - role: "developer".to_owned(), - content: vec![ContentItem::InputText { - text: request.instructions, - }], - phase: None, - internal_chat_message_metadata_passthrough: None, - }, - ResponseItem::Message { - id: None, - role: "user".to_owned(), - content: request - .input - .into_iter() - .map(|text| ContentItem::InputText { text }) - .collect(), - phase: None, - internal_chat_message_metadata_passthrough: None, - }, - ], + input, tools: None, tool_choice: "none".to_owned(), parallel_tool_calls: false, diff --git a/codex-rs/ext/guardian-v2/src/sampler_tests.rs b/codex-rs/ext/guardian-v2/src/sampler_tests.rs index 2d1c30a427..5c7848a599 100644 --- a/codex-rs/ext/guardian-v2/src/sampler_tests.rs +++ b/codex-rs/ext/guardian-v2/src/sampler_tests.rs @@ -6,6 +6,8 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; +use codex_protocol::ResponseItemId; +use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::SessionSource; use core_test_support::responses; @@ -63,6 +65,7 @@ fn sampler_config(base_url: String) -> LunaSamplerConfig { thread_id: "thread-1".to_owned(), originator: Some("guardian-v2-test".to_owned()), service_tier: None, + luna_compaction_hash: None, } } @@ -70,6 +73,8 @@ fn sample_request(turn_id: &str) -> LunaSamplingRequest { LunaSamplingRequest { instructions: "Return a risk score.".to_owned(), input: vec!["The user requested a README summary.".to_owned()], + parent_compaction: None, + parent_compaction_hash: None, output_schema: json!({ "type": "object", "properties": { "score": { "type": "number" } }, @@ -118,6 +123,7 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ thread_id: "thread-1".to_owned(), originator: Some("guardian-v2-test".to_owned()), service_tier: None, + luna_compaction_hash: None, }) .await?; @@ -161,6 +167,8 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ "The user requested a README summary.".to_owned(), "The assistant inspected README.md.".to_owned(), ], + parent_compaction: None, + parent_compaction_hash: None, output_schema: schema.clone(), reasoning_effort: ReasoningEffort::None, turn_id: "turn-1".to_owned(), @@ -182,6 +190,8 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ .sample(LunaSamplingRequest { instructions: "Return a risk score.".to_owned(), input: vec!["The user requested a source review.".to_owned()], + parent_compaction: None, + parent_compaction_hash: None, output_schema: schema, reasoning_effort: ReasoningEffort::Medium, turn_id: "turn-2".to_owned(), @@ -226,6 +236,70 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn sampler_reuses_parent_compaction_only_for_matching_model_hashes() -> Result<()> { + skip_if_no_network!(Ok(())); + + for (parent_hash, luna_hash, should_reuse) in [ + (Some("compatible"), Some("compatible"), true), + (Some("parent"), Some("luna"), false), + (None, Some("compatible"), false), + (Some("compatible"), None, false), + (Some(""), Some(""), false), + ] { + let events = vec![ + ev_assistant_message("sample", r#"{"score":0.25}"#), + ev_completed("response-1"), + ]; + let server = + responses::start_websocket_server(vec![Vec::new(), vec![events.clone(), events]]).await; + let mut config = sampler_config(format!( + "http://{}/v1", + server.uri().trim_start_matches("ws://") + )); + config.luna_compaction_hash = luna_hash.map(str::to_owned); + let sampler = LunaSampler::connect(config).await?; + let parent_compaction = ResponseItem::Compaction { + id: Some(ResponseItemId::from_server("cmp_parent".to_owned())), + encrypted_content: "opaque encrypted summary".to_owned(), + internal_chat_message_metadata_passthrough: None, + }; + let mut request = sample_request("turn-1"); + request.parent_compaction = Some(parent_compaction.clone()); + request.parent_compaction_hash = parent_hash.map(str::to_owned); + + assert_eq!(sampler.sample(request).await?, r#"{"score":0.25}"#); + + let request = server + .wait_for_request(/*connection_index*/ 1, /*request_index*/ 0) + .await + .body_json(); + let input = request["input"].as_array().expect("input items"); + assert_eq!(input[0]["type"], "additional_tools"); + assert_eq!(input[1]["role"], "developer"); + if should_reuse { + assert_eq!(input.len(), 4); + assert_eq!(input[2], serde_json::to_value(&parent_compaction)?); + assert_eq!(input[3]["role"], "user"); + + let mut switched_request = sample_request("turn-2"); + switched_request.parent_compaction = Some(parent_compaction); + switched_request.parent_compaction_hash = Some("incompatible".to_owned()); + assert_eq!(sampler.sample(switched_request).await?, r#"{"score":0.25}"#); + let switched_request = server + .wait_for_request(/*connection_index*/ 1, /*request_index*/ 1) + .await + .body_json(); + assert_eq!(switched_request["input"][2]["role"], "user"); + } else { + assert_eq!(input.len(), 3); + assert_eq!(input[2]["role"], "user"); + } + } + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn sampler_returns_complete_json_before_terminal_response_events() -> Result<()> { skip_if_no_network!(Ok(())); @@ -257,6 +331,7 @@ async fn sampler_returns_complete_json_before_terminal_response_events() -> Resu thread_id: "thread-1".to_owned(), originator: None, service_tier: None, + luna_compaction_hash: None, }) .await?; @@ -265,6 +340,8 @@ async fn sampler_returns_complete_json_before_terminal_response_events() -> Resu sampler.sample(LunaSamplingRequest { instructions: "Return a risk score.".to_owned(), input: vec!["The user requested a README summary.".to_owned()], + parent_compaction: None, + parent_compaction_hash: None, output_schema: json!({ "type": "object", "properties": { "score": { "type": "number" } },