From 726f1492db90501746639e8c61cd3c674735a1a7 Mon Sep 17 00:00:00 2001 From: felixxia-oai Date: Wed, 30 Sep 2026 21:03:33 +0000 Subject: [PATCH] Add conversation mode to Guardian v2 async classification (#49793) ## What changed - Add `features.guardianv2.async_classifier_mode` with `snapshot` and `conversation` modes. Local configuration overrides model defaults; snapshots remain the fallback. - Retain completed classifier requests and responses in conversation mode, appending transcript deltas without repeating previously delivered user instructions. - Rebuild retained history when authorization, classifier instructions, or parent compaction changes, after an incomplete response, or when the next request exceeds the configured token limit. Add `async_classifier_conversation_token_limit`, defaulting to 100,000 tokens, while preserving the model's input limit for fresh requests. - Serialize conversation requests in observation order, bound outstanding work, and reject stale work after authorization changes or thread shutdown. Record retained and new input token costs. ## Testing Add tests for configuration precedence, history reuse and resets, ordered admission, queue overflow, and stale or stopped requests. Exercise the approval matrix in both modes, including resume, fork, mandatory review, authorization changes, and stream failure after an early score. GitOrigin-RevId: 35b9800104648cde4f467086e678dee12458c3e5 --- .../suite/v2/guardian_stateful_async_tests.rs | 38 +++ .../app-server/tests/suite/v2/guardian_v2.rs | 168 +++++++++----- codex-rs/core/config.schema.json | 22 ++ .../src/async_scorer/authorization.rs | 4 + .../src/async_scorer/classification.rs | 107 ++++++--- .../guardian-v2/src/async_scorer/config.rs | 24 ++ .../src/async_scorer/config_tests.rs | 36 ++- .../src/async_scorer/conversation.rs | 194 ++++++++++++++++ .../src/async_scorer/conversation_tests.rs | 218 ++++++++++++++++++ .../guardian-v2/src/async_scorer/extension.rs | 20 ++ .../src/async_scorer/extension_tests.rs | 5 + .../guardian-v2/src/async_scorer/metrics.rs | 46 ++++ .../ext/guardian-v2/src/async_scorer/mod.rs | 3 +- .../src/async_scorer/observation.rs | 25 +- .../guardian-v2/src/async_scorer/request.rs | 5 +- .../async_scorer/retained_sampling_tests.rs | 2 + .../guardian-v2/src/async_scorer/sampler.rs | 27 +-- .../src/async_scorer/sampler/execution.rs | 1 + .../src/async_scorer/sampler_tests.rs | 4 +- .../src/async_scorer/transcript.rs | 14 +- codex-rs/features/src/feature_configs.rs | 13 ++ codex-rs/features/src/tests.rs | 4 + .../src/retained_instructions.rs | 6 +- codex-rs/protocol/src/openai_models.rs | 1 + .../protocol/src/openai_models/guardian_v2.rs | 15 ++ 25 files changed, 872 insertions(+), 130 deletions(-) create mode 100644 codex-rs/app-server/tests/suite/v2/guardian_stateful_async_tests.rs create mode 100644 codex-rs/ext/guardian-v2/src/async_scorer/conversation.rs create mode 100644 codex-rs/ext/guardian-v2/src/async_scorer/conversation_tests.rs diff --git a/codex-rs/app-server/tests/suite/v2/guardian_stateful_async_tests.rs b/codex-rs/app-server/tests/suite/v2/guardian_stateful_async_tests.rs new file mode 100644 index 0000000000..d430225bdd --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/guardian_stateful_async_tests.rs @@ -0,0 +1,38 @@ +use super::*; +use codex_protocol::openai_models::AsyncClassifierMode; +use test_case::test_case; + +#[test_case(GuardianRisk::Low, ThreadLifecycle::New, ModelReviewRequirement::Optional; "low")] +#[test_case(GuardianRisk::Low, ThreadLifecycle::RootUserRestriction, ModelReviewRequirement::Optional; "new_user_instruction_resets_conversation")] +#[test_case(GuardianRisk::LowWithStreamFailure, ThreadLifecycle::New, ModelReviewRequirement::Optional; "failure_after_early_score")] +#[test_case(GuardianRisk::Low, ThreadLifecycle::RootRestrictionDuringClassification, ModelReviewRequirement::Optional; "authorization_changed_during_classification")] +#[test_case(GuardianRisk::High, ThreadLifecycle::New, ModelReviewRequirement::Optional; "high")] +#[test_case(GuardianRisk::InvalidResponse, ThreadLifecycle::New, ModelReviewRequirement::Optional; "failed")] +#[test_case(GuardianRisk::Low, ThreadLifecycle::Resume, ModelReviewRequirement::Optional; "resume")] +#[test_case(GuardianRisk::Low, ThreadLifecycle::Fork, ModelReviewRequirement::Optional; "fork")] +#[test_case(GuardianRisk::Low, ThreadLifecycle::New, ModelReviewRequirement::Required; "mandatory_review")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn classifier_modes_preserve_approval_matrix( + risk: GuardianRisk, + lifecycle: ThreadLifecycle, + requirement: ModelReviewRequirement, +) -> Result<()> { + skip_if_no_network!(Ok(())); + for mode in [ + AsyncClassifierMode::Snapshot, + AsyncClassifierMode::Conversation, + ] { + guardian_v2_routes_scoped_tool_approvals( + risk, + lifecycle, + requirement, + ReviewOutcome::Allow, + TranscriptContent::Normal, + GuardianToolScope::AllTools, + /*sensitive_action*/ None, + mode, + ) + .await?; + } + Ok(()) +} diff --git a/codex-rs/app-server/tests/suite/v2/guardian_v2.rs b/codex-rs/app-server/tests/suite/v2/guardian_v2.rs index b084707a4e..a8935aac26 100644 --- a/codex-rs/app-server/tests/suite/v2/guardian_v2.rs +++ b/codex-rs/app-server/tests/suite/v2/guardian_v2.rs @@ -188,6 +188,7 @@ struct MockResponsesState { context_metrics_recorded: Notify, luna_score: f64, invalid_classification: bool, + fail_after_classification: bool, review_outcome: ReviewOutcome, transcript_content: TranscriptContent, mcp_server_name: Option<&'static str>, @@ -220,6 +221,7 @@ enum TranscriptContent { #[derive(Clone, Copy)] enum GuardianRisk { Low, + LowWithStreamFailure, Threshold, High, InvalidResponse, @@ -524,10 +526,10 @@ async fn parent_response( item["call_id"] == "guardian-action-0" && item["type"] == "function_call_output" }) .expect("first tool result after root revocation"); + let output = output.to_string(); assert!( - output - .to_string() - .contains("Tool execution was declined by Guardian."), + output.contains("Tool execution was declined by Guardian.") + || output.contains("This action was rejected due to unacceptable risk."), "the refreshed review must deny the first tool after revocation: {output}" ); } @@ -623,7 +625,11 @@ async fn luna_response(state: &MockResponsesState, request: Value) -> Vec responses::ev_response_created("luna-score"), responses::ev_output_text_delta(classification), responses::ev_assistant_message("luna-score-message", classification), - responses::ev_completed("luna-score"), + if state.fail_after_classification { + json!({"type":"response.failed", "response":{"id":"luna-score", "error":{"code":"server_error", "message":"stream failed after score"}}}) + } else { + responses::ev_completed("luna-score") + }, ] } @@ -666,10 +672,13 @@ async fn guardian_v2_routes_tool_approvals( transcript_content, GuardianToolScope::AllTools, /*sensitive_action*/ None, + codex_protocol::openai_models::AsyncClassifierMode::Snapshot, ) .await } +// Keep the independent approval scenario inputs explicit at each call site. +#[allow(clippy::too_many_arguments)] async fn guardian_v2_routes_scoped_tool_approvals( risk: GuardianRisk, lifecycle: ThreadLifecycle, @@ -678,7 +687,14 @@ async fn guardian_v2_routes_scoped_tool_approvals( transcript_content: TranscriptContent, scope: GuardianToolScope, sensitive_action: Option, + classifier_mode: codex_protocol::openai_models::AsyncClassifierMode, ) -> Result<()> { + let fail_after_classification = matches!(risk, GuardianRisk::LowWithStreamFailure); + let risk = if fail_after_classification { + GuardianRisk::Low + } else { + risk + }; let server_name = match scope { GuardianToolScope::AllTools => TEST_SERVER_NAME, GuardianToolScope::ComputerUseOnly { server_name } => server_name, @@ -705,7 +721,9 @@ async fn guardian_v2_routes_scoped_tool_approvals( { (0.25, 1) } - GuardianRisk::Low | GuardianRisk::InvalidResponse => (0.25, 2), + GuardianRisk::Low | GuardianRisk::LowWithStreamFailure | GuardianRisk::InvalidResponse => { + (0.25, 2) + } GuardianRisk::Threshold => (0.5, 2), GuardianRisk::High => (0.95, 2), }; @@ -719,6 +737,7 @@ async fn guardian_v2_routes_scoped_tool_approvals( let responses_state = Arc::new(MockResponsesState { luna_score, invalid_classification: matches!(risk, GuardianRisk::InvalidResponse), + fail_after_classification, review_outcome, transcript_content, mcp_server_name: Some(server_name), @@ -874,6 +893,10 @@ async fn guardian_v2_routes_scoped_tool_approvals( } else { "prompt" }; + let classifier_mode = match classifier_mode { + codex_protocol::openai_models::AsyncClassifierMode::Snapshot => "snapshot", + codex_protocol::openai_models::AsyncClassifierMode::Conversation => "conversation", + }; let mut mock_config = MockResponsesConfig::new(&responses_url) .with_model(MODEL) .with_provider_config("supports_websockets = false") @@ -883,7 +906,7 @@ async fn guardian_v2_routes_scoped_tool_approvals( analytics_server.uri(), )) .with_extra_config(&format!( - "[mcp_servers.{server_name}]\nurl = \"{mcp_server_url}/mcp\"\ndefault_tools_approval_mode = \"{tool_approval_mode}\"\n\n[analytics]\nenabled = true\n\n[otel]\nmetrics_exporter = {{ otlp-http = {{ endpoint = \"{responses_url}/metrics\", protocol = \"json\" }} }}\n\n[features.guardianv2]\nenabled = true{guardian_scope_config}" + "[mcp_servers.{server_name}]\nurl = \"{mcp_server_url}/mcp\"\ndefault_tools_approval_mode = \"{tool_approval_mode}\"\n\n[analytics]\nenabled = true\n\n[otel]\nmetrics_exporter = {{ otlp-http = {{ endpoint = \"{responses_url}/metrics\", protocol = \"json\" }} }}\n\n[features.guardianv2]\nenabled = true\nasync_classifier_mode = \"{classifier_mode}\"{guardian_scope_config}" )) .enable_feature(Feature::GuardianApproval); if lifecycle.has_user_input() || lifecycle.has_root_user_input() { @@ -1177,6 +1200,57 @@ async fn guardian_v2_routes_scoped_tool_approvals( submit_user_input_response(&mut app_server, answers).await?; } let second_sample = wait_for_luna_request(responses_state.as_ref(), /*index*/ 1).await?; + if matches!( + lifecycle, + ThreadLifecycle::New + | ThreadLifecycle::Resume + | ThreadLifecycle::Fork + | ThreadLifecycle::RootUserRestriction + | ThreadLifecycle::RootRestrictionDuringClassification + ) { + let retained = classifier_mode == "conversation" + && !fail_after_classification + && !matches!(risk, GuardianRisk::InvalidResponse) + && !late_root_restriction; + let input = second_sample["input"] + .as_array() + .expect("second Luna request input should be an array"); + assert_eq!( + input.iter().any(|item| item["id"] == "luna-score-message"), + retained + ); + if retained { + let first = luna_request["input"] + .as_array() + .expect("first Luna request input should be an array"); + assert_eq!(&input[..first.len()], first); + let delta = serde_json::to_string(&input[first.len()..])?; + assert!(delta.contains("TRANSCRIPT DELTA START")); + if matches!( + lifecycle, + ThreadLifecycle::New | ThreadLifecycle::RootUserRestriction + ) { + assert!( + delta.contains("guardian-0"), + "the new tool evidence must be included" + ); + assert!( + !delta.contains(USER_CONTEXT), + "unchanged user instructions must stay in the retained prefix instead of being appended again" + ); + assert!( + !delta.contains(">>> RETAINED USER INSTRUCTIONS START"), + "a tool-only delta must not append the old retained-instruction section" + ); + } + } else { + assert!(!second_sample.to_string().contains("TRANSCRIPT DELTA START")); + assert!( + second_sample.to_string().contains(USER_CONTEXT), + "snapshot requests and rebuilt conversations must include the user instructions" + ); + } + } if mixed_evidence { let input = second_sample["input"].as_array().expect("Luna input"); assert_eq!( @@ -1675,6 +1749,23 @@ async fn guardian_v2_routes_scoped_tool_approvals( } let post_authorization_change_sample = wait_for_luna_request(responses_state.as_ref(), /*index*/ 2).await?; + if classifier_mode == "conversation" + && matches!(lifecycle, ThreadLifecycle::RootUserRestriction) + { + let input = post_authorization_change_sample["input"] + .as_array() + .expect("fresh classifier input"); + assert!( + !input.iter().any(|item| item["id"] == "luna-score-message"), + "new user instructions must invalidate the previous classifier conversation" + ); + let text = serde_json::to_string(input)?; + assert!(!text.contains("TRANSCRIPT DELTA START")); + assert!( + text.contains(USER_CONTEXT), + "rebuilding must restore the earlier user instructions" + ); + } assert_eq!( post_authorization_change_sample["prompt_cache_key"], format!("guardian-v2:{reviewed_thread_id}") @@ -1833,7 +1924,8 @@ async fn guardian_v2_routes_scoped_tool_approvals( event["event_params"]["risk_level"] ]), match risk { - GuardianRisk::Low => json!(["success", "low"]), + GuardianRisk::Low | GuardianRisk::LowWithStreamFailure => + json!(["success", "low"]), GuardianRisk::Threshold | GuardianRisk::High => json!(["success", "high"]), GuardianRisk::InvalidResponse => json!(["failure", null]), } @@ -1876,19 +1968,6 @@ async fn guardian_v2_routes_scoped_tool_approvals( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn guardian_v2_low_risk_actions_skip_subsequent_reviews() -> Result<()> { - skip_if_no_network!(Ok(())); - guardian_v2_routes_tool_approvals( - GuardianRisk::Low, - ThreadLifecycle::New, - ModelReviewRequirement::Optional, - ReviewOutcome::Allow, - TranscriptContent::Normal, - ) - .await -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn guardian_v2_trusts_invoked_user_skills_but_rejects_repository_forgery() -> Result<()> { skip_if_no_network!(Ok(())); @@ -2120,6 +2199,7 @@ async fn guardian_v2_inherits_root_user_skills_for_delegated_workers() -> Result server_name: "node_repl", }, /*sensitive_action*/ None, + codex_protocol::openai_models::AsyncClassifierMode::Snapshot, ) .await } @@ -2142,6 +2222,7 @@ async fn guardian_v2_computer_use_only_scopes_classification_and_fast_reviews( TranscriptContent::Normal, GuardianToolScope::ComputerUseOnly { server_name }, /*sensitive_action*/ None, + codex_protocol::openai_models::AsyncClassifierMode::Snapshot, ) .await } @@ -2439,6 +2520,7 @@ async fn guardian_v2_required_model_computer_use_preserves_strict_approval( TranscriptContent::Normal, GuardianToolScope::ComputerUseOnly { server_name }, sensitive_action, + codex_protocol::openai_models::AsyncClassifierMode::Snapshot, ) .await } @@ -2460,6 +2542,7 @@ async fn guardian_v2_discards_sync_reviews_after_user_input_answer( server_name: "node_repl", }, /*sensitive_action*/ None, + codex_protocol::openai_models::AsyncClassifierMode::Snapshot, ) .await } @@ -2482,6 +2565,7 @@ async fn guardian_v2_validates_user_input_before_history_truncation( server_name: "node_repl", }, /*sensitive_action*/ None, + codex_protocol::openai_models::AsyncClassifierMode::Snapshot, ) .await } @@ -2504,6 +2588,7 @@ async fn guardian_v2_propagates_root_user_input_to_worker_reviews( server_name: "node_repl", }, /*sensitive_action*/ None, + codex_protocol::openai_models::AsyncClassifierMode::Snapshot, ) .await } @@ -2541,19 +2626,6 @@ async fn guardian_v2_threshold_score_requires_full_reviews() -> Result<()> { .await } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn guardian_v2_required_model_bypasses_scoring_and_runs_full_reviews() -> Result<()> { - skip_if_no_network!(Ok(())); - guardian_v2_routes_tool_approvals( - GuardianRisk::Low, - ThreadLifecycle::New, - ModelReviewRequirement::Required, - ReviewOutcome::Allow, - TranscriptContent::Normal, - ) - .await -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn guardian_v2_required_model_cannot_reuse_a_cached_score_for_skipped_exec() -> Result<()> { skip_if_no_network!(Ok(())); @@ -2571,32 +2643,6 @@ async fn guardian_v2_required_model_cannot_reuse_a_cached_score_for_skipped_exec .await } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn resumed_thread_ignores_persisted_guardian_score() -> Result<()> { - skip_if_no_network!(Ok(())); - guardian_v2_routes_tool_approvals( - GuardianRisk::Low, - ThreadLifecycle::Resume, - ModelReviewRequirement::Optional, - ReviewOutcome::Allow, - TranscriptContent::Normal, - ) - .await -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn forked_thread_ignores_persisted_guardian_score() -> Result<()> { - skip_if_no_network!(Ok(())); - guardian_v2_routes_tool_approvals( - GuardianRisk::Low, - ThreadLifecycle::Fork, - ModelReviewRequirement::Optional, - ReviewOutcome::Allow, - TranscriptContent::Normal, - ) - .await -} - #[test_case(ThreadLifecycle::RootRestriction; "worker_root_restriction")] #[test_case(ThreadLifecycle::RootUserRestriction; "root_user_restriction")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -2635,6 +2681,10 @@ async fn guardian_v2_low_scores_require_current_authorization( server_name: "node_repl", }, /*sensitive_action*/ None, + codex_protocol::openai_models::AsyncClassifierMode::Snapshot, ) .await } + +#[path = "guardian_stateful_async_tests.rs"] +mod stateful_async; diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index de866e4609..f3a1811f04 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -336,6 +336,14 @@ } ] }, + "AsyncClassifierMode": { + "description": "Async classifier experiment, resolved once for each parent task.", + "enum": [ + "snapshot", + "conversation" + ], + "type": "string" + }, "AuthCredentialsStoreMode": { "description": "Determine where Codex should store CLI auth credentials.", "oneOf": [ @@ -1853,6 +1861,20 @@ "additionalProperties": false, "description": "User-configurable prompt, approval, and context settings for Guardian v2.", "properties": { + "async_classifier_conversation_token_limit": { + "description": "Reset retained async history when the next request exceeds this token estimate. Defaults to 100,000. Fresh requests remain subject to the model's input limit.", + "format": "uint", + "minimum": 1.0, + "type": "integer" + }, + "async_classifier_mode": { + "allOf": [ + { + "$ref": "#/definitions/AsyncClassifierMode" + } + ], + "description": "Classifier experiment override. Otherwise use the model default, then snapshots." + }, "classifier_instructions": { "type": "string" }, diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/authorization.rs b/codex-rs/ext/guardian-v2/src/async_scorer/authorization.rs index 5150c6f84f..bdcb08384b 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/authorization.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/authorization.rs @@ -17,6 +17,10 @@ pub(super) struct ScoreAuthorization { } impl ScoreAuthorization { + pub(super) async fn is_current(&self, thread: &CodexThread) -> bool { + self == &Self::current(thread, &self.permissions).await + } + pub(super) async fn current( thread: &CodexThread, permissions: &codex_guardian_context::PermissionContext, diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/classification.rs b/codex-rs/ext/guardian-v2/src/async_scorer/classification.rs index fe78b54f51..23118ea7db 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/classification.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/classification.rs @@ -39,6 +39,8 @@ use codex_protocol::security_risk::SecurityRiskScore; use super::authorization::ScoreAuthorization; use super::config::GuardianV2Config; +use super::conversation::ConversationRequest; +use super::conversation::Reservation; use super::metrics::record_classification; use super::metrics::record_classification_risk; use super::metrics::sampler_failure_reason; @@ -46,12 +48,15 @@ use super::sampler::LunaSampler; use super::sampler::LunaSamplerError; use super::sampler::LunaSamplingRequest; use super::score::GuardianV2ScoreProgress; +use super::transcript::CollectedTranscript; use super::transcript::ContextInput; use super::truncation::ClassificationTruncations; use super::trusted_skills::TrustedSkillInvocations; use super::trusted_tools::trusted_tool_context; +use codex_guardian_context::ComposedContext; pub(super) struct Classification { + pub(super) reservation: Option, pub(super) classification_started_at: Instant, pub(super) sampler: Arc, pub(super) guardian_config: GuardianV2Config, @@ -86,6 +91,14 @@ pub(super) struct Classification { pub(super) score_authorization: ScoreAuthorization, } +enum ClassificationContext { + Snapshot(ComposedContext), + Conversation { + evidence: CollectedTranscript, + reservation: Reservation, + }, +} + enum ClassificationOutcome { Scored, Superseded, @@ -94,6 +107,7 @@ enum ClassificationOutcome { impl Classification { pub(super) async fn run(self) { let Self { + reservation, classification_started_at, sampler, guardian_config, @@ -127,6 +141,9 @@ impl Classification { root_snapshot, score_authorization, } = self; + let generation = reservation + .as_ref() + .map(|reservation| reservation.generation.clone()); let mut truncations = ClassificationTruncations::default(); let trusted_tool_context = match mcp_tool.as_ref() { Some(tool) => { @@ -192,7 +209,7 @@ impl Classification { .collect::>(); let transcript = PreviousReviews::try_from_fragments(review_fragments).and_then(|reviews| { - guardian_config.transcript.build_context(ContextInput { + let context_input = ContextInput { target: ContextTarget::Async, history: history.as_ref(), root_conversation: root_conversation.as_deref().unwrap_or_default(), @@ -203,7 +220,20 @@ impl Classification { trusted_tool: trusted_tool_context.as_ref(), trusted_skill_paths: &trusted_skill_paths, node_repl_images: Some(&node_repl_images), - }) + }; + match reservation { + Some(reservation) => guardian_config + .transcript + .collect_context(context_input) + .map(|evidence| ClassificationContext::Conversation { + evidence, + reservation, + }), + None => guardian_config + .transcript + .build_context(context_input) + .map(ClassificationContext::Snapshot), + } }); let mut transcript = match transcript { Ok(transcript) => transcript, @@ -226,24 +256,10 @@ impl Classification { }; drop(history); drop(node_repl_images); - truncations.extend(std::mem::take(&mut transcript.truncations)); - if let Some(metrics) = metrics.as_deref() { - for (section, cost) in transcript.section_costs() { - for (measurement, value) in cost.measurements() { - metrics.histogram_with_boundaries( - codex_guardian_context::SECTION_COST_METRIC, - i64::try_from(value).unwrap_or(i64::MAX), - codex_guardian_context::SECTION_COST_BOUNDARIES, - &[ - ("target", "async"), - ("section", section), - ("measurement", measurement), - ], - ); - } - } + if let ClassificationContext::Snapshot(context) = &mut transcript { + truncations.extend(std::mem::take(&mut context.truncations)); + super::metrics::record_section_costs(metrics.as_deref(), context.section_costs()); } - let classification_input = transcript.into_messages(); let mut failure_reason = "invalid_output"; let mut classification_risk = None; let mut classification_finished_at = None; @@ -275,19 +291,40 @@ impl Classification { let policy = config.resolve_guardian_policy(model_messages); let extra_policy = config.guardian_extra_policy.as_deref().unwrap_or_default(); let instructions = guardian_config.render_classifier_instructions(policy, extra_policy); - let output = match sampler - .sample(LunaSamplingRequest { - parent_response_id, - instructions, - input: classification_input, - parent_compaction, - parent_compaction_hash, - reasoning_effort: guardian_config.reasoning_effort.clone(), - parent_turn_id: turn_id.clone(), - root_turn_id, - }) - .await - { + let mut sampling = LunaSamplingRequest { + parent_response_id, + instructions, + input: Vec::new(), + parent_compaction, + parent_compaction_hash, + reasoning_effort: guardian_config.reasoning_effort.clone(), + parent_turn_id: turn_id.clone(), + root_turn_id, + }; + let result = match transcript { + ClassificationContext::Snapshot(context) => { + sampling.input = context.into_messages(); + sampler.sample(sampling).await + } + ClassificationContext::Conversation { + evidence, + reservation, + } => { + let (ready, score) = tokio::sync::oneshot::channel(); + reservation.submit(ConversationRequest { + evidence, + sampling, + reset_token_limit: guardian_config + .async_classifier_conversation_token_limit, + ready, + authorization: score_authorization.clone(), + thread: Arc::clone(&thread), + metrics: metrics.clone(), + }); + score.await.unwrap_or(Err(LunaSamplerError::Superseded)) + } + }; + let output = match result { Ok(output) => output, Err(LunaSamplerError::Superseded) => { return Ok(ClassificationOutcome::Superseded); @@ -312,8 +349,10 @@ impl Classification { ), sampled_at: Some(sampled_at.into()), }; - if score_authorization - != ScoreAuthorization::current(&thread, &score_authorization.permissions).await + if !score_authorization.is_current(&thread).await + || generation + .as_ref() + .is_some_and(|generation| generation.strong_count() == 0) { return Ok(ClassificationOutcome::Superseded); } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/config.rs b/codex-rs/ext/guardian-v2/src/async_scorer/config.rs index 896b8e445e..f0bf74386c 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/config.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/config.rs @@ -28,6 +28,7 @@ const MAX_MODEL_CONTEXT_ITEM_TOKENS: usize = 100_000; const DEFAULT_REVIEW_THRESHOLD: f64 = 0.5; const LEGACY_REVIEW_THRESHOLD: f64 = 0.8; const DEFAULT_MAX_TOOL_CALL_LAG: usize = 2; +const DEFAULT_ASYNC_CLASSIFIER_CONVERSATION_TOKEN_LIMIT: usize = 100_000; pub(crate) const CLASSIFICATION_OUTPUT_INSTRUCTIONS: &str = "Your first output token is the entire classification: `high` for high risk or `low` for low risk. Output that token immediately and nothing else."; #[derive(Clone, Debug, PartialEq)] @@ -37,6 +38,7 @@ pub(crate) struct GuardianV2Config { pub(crate) classifier_instructions: String, pub(crate) review_threshold: f64, pub(crate) max_tool_call_lag: usize, + pub(crate) async_classifier_conversation_token_limit: usize, pub(crate) reasoning_effort: ReasoningEffort, pub(crate) max_action_tokens: usize, /// No truncation limit is applied unless local or model configuration supplies one. @@ -48,6 +50,16 @@ pub(crate) struct GuardianV2Config { } impl GuardianV2Config { + pub(super) fn classifier_mode( + &self, + model_defaults: Option<&GuardianV2ModelConfig>, + ) -> codex_protocol::openai_models::AsyncClassifierMode { + self.local_overrides + .async_classifier_mode + .or_else(|| model_defaults.and_then(|defaults| defaults.async_classifier_mode)) + .unwrap_or_default() + } + pub(crate) fn resolve(config: &Config) -> Result { let effective_config = config.config_layer_stack.effective_config(); let configured = match effective_config @@ -102,6 +114,9 @@ impl GuardianV2Config { configured.max_tool_call_lag = configured .max_tool_call_lag .or(model_defaults.max_tool_call_lag); + configured.async_classifier_conversation_token_limit = configured + .async_classifier_conversation_token_limit + .or(model_defaults.async_classifier_conversation_token_limit); configured.reasoning_effort = configured .reasoning_effort .or_else(|| model_defaults.reasoning_effort.clone()); @@ -167,6 +182,14 @@ impl GuardianV2Config { } fn from_overrides(configured: GuardianV2ConfigToml) -> Result { + let async_classifier_conversation_token_limit = configured + .async_classifier_conversation_token_limit + .unwrap_or(DEFAULT_ASYNC_CLASSIFIER_CONVERSATION_TOKEN_LIMIT); + if async_classifier_conversation_token_limit == 0 { + return Err( + "Guardian v2 async_classifier_conversation_token_limit must be positive".to_owned(), + ); + } // Existing custom and model-owned prompts retain their original calibration // unless their owner explicitly supplies a matching threshold. let review_threshold = configured.review_threshold.unwrap_or_else(|| { @@ -253,6 +276,7 @@ impl GuardianV2Config { ); Ok(Self { local_overrides: configured.clone(), + async_classifier_conversation_token_limit, persist_scores: configured.persist_scores.unwrap_or(false), classifier_instructions: configured.classifier_instructions.unwrap_or_else(|| { ResolvedModelMessages::bundled() diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/config_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/config_tests.rs index ac02d87d60..7ad9acbd5d 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/config_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/config_tests.rs @@ -213,6 +213,7 @@ fn model_runtime_settings_preserve_local_overrides() { classifier_instructions: Some(prompt.clone()), max_classifier_instruction_tokens: Some(256), max_tool_call_lag: Some(1), + async_classifier_conversation_token_limit: Some(80_000), reuse_parent_compaction: Some(false), transcript: Some(GuardianV2TranscriptModelConfig { include_images: Some(true), @@ -228,15 +229,17 @@ fn model_runtime_settings_preserve_local_overrides() { ( inherited.max_classifier_instruction_tokens, inherited.max_tool_call_lag, + inherited.async_classifier_conversation_token_limit, inherited.reuse_parent_compaction, inherited.transcript.include_images, ), - (Some(256), 1, false, true) + (Some(256), 1, 80_000, false, true) ); let overridden = GuardianV2Config::from_overrides(GuardianV2ConfigToml { max_classifier_instruction_tokens: Some(512), max_tool_call_lag: Some(4), + async_classifier_conversation_token_limit: Some(120_000), reuse_parent_compaction: Some(true), transcript: Some(GuardianV2TranscriptConfigToml { include_images: Some(false), @@ -251,10 +254,11 @@ fn model_runtime_settings_preserve_local_overrides() { ( overridden.max_classifier_instruction_tokens, overridden.max_tool_call_lag, + overridden.async_classifier_conversation_token_limit, overridden.reuse_parent_compaction, overridden.transcript.include_images, ), - (Some(512), 4, true, false) + (Some(512), 4, 120_000, true, false) ); let uncapped_defaults = GuardianV2ModelConfig { @@ -271,3 +275,31 @@ fn model_runtime_settings_preserve_local_overrides() { ) ); } + +#[test] +fn classifier_mode_uses_local_override_then_model_default() { + use codex_protocol::openai_models::AsyncClassifierMode; + for local in [ + None, + Some(AsyncClassifierMode::Snapshot), + Some(AsyncClassifierMode::Conversation), + ] { + let config = GuardianV2Config::from_overrides(GuardianV2ConfigToml { + async_classifier_mode: local, + ..Default::default() + }) + .unwrap(); + let defaults = GuardianV2ModelConfig { + async_classifier_mode: Some(AsyncClassifierMode::Conversation), + ..Default::default() + }; + assert_eq!( + config.classifier_mode(Some(&defaults)), + local.unwrap_or(AsyncClassifierMode::Conversation) + ); + assert_eq!( + config.classifier_mode(/*model_defaults*/ None), + local.unwrap_or_default() + ); + } +} diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/conversation.rs b/codex-rs/ext/guardian-v2/src/async_scorer/conversation.rs new file mode 100644 index 0000000000..e4af93e6fd --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/async_scorer/conversation.rs @@ -0,0 +1,194 @@ +//! One worker owns each treatment conversation and commits only completed requests. +//! Admission reserves order before async preparation; permits bound active and queued work. +//! Backend replacement invalidates its generation without changing snapshot shutdown policy. +//! Oversized retained requests rebuild from fresh evidence before the model's hard limit. +//! Prepared requests own their history; only progress remains in state during sampling. +//! Host-only instruction delivery metadata is committed with its completed input messages. + +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::Weak; + +use codex_context_fragments::RenderedFragment; +use codex_core::CodexThread; +use codex_extension_api::ExtensionMetrics; +use codex_guardian_reviewer::ConversationState; +use codex_history::ResponseItemEnvelope; +use codex_protocol::models::ResponseItem; +use tokio::sync::OwnedSemaphorePermit; +use tokio::sync::Semaphore; +use tokio::sync::mpsc; +use tokio::sync::oneshot; + +use super::authorization::ScoreAuthorization; +use super::sampler::LunaSampler; +use super::sampler::LunaSamplerError; +use super::sampler::LunaSamplingRequest; +use super::transcript::CollectedTranscript; + +const MAX_OUTSTANDING: usize = 16; + +pub(super) struct ConversationBackend { + queue: Mutex>, + capacity: Arc, + generation: Arc<()>, +} + +struct QueuedRequest { + request: oneshot::Receiver, + _permit: OwnedSemaphorePermit, +} + +pub(super) struct Reservation { + request: oneshot::Sender, + pub(super) generation: Weak<()>, +} + +pub(super) struct ConversationRequest { + pub(super) evidence: CollectedTranscript, + pub(super) sampling: LunaSamplingRequest, + pub(super) reset_token_limit: usize, + pub(super) authorization: ScoreAuthorization, + pub(super) thread: Arc, + pub(super) ready: oneshot::Sender>, + pub(super) metrics: Option>, +} + +// Authorization covers model settings; local config changes replace the backend. +#[derive(PartialEq)] +struct ReuseKey { + authorization: ScoreAuthorization, + instructions: RenderedFragment, + parent_compaction: Option, +} + +impl ConversationBackend { + pub(super) fn new(sampler: Arc) -> Self { + let (sender, receiver) = mpsc::unbounded_channel(); + let generation = Arc::new(()); + tokio::spawn(run(receiver, sampler, Arc::downgrade(&generation))); + Self { + queue: Mutex::new(sender), + capacity: Arc::new(Semaphore::new(MAX_OUTSTANDING)), + generation, + } + } + + pub(super) fn reserve( + &self, + observe: impl FnOnce() -> usize, + ) -> (usize, Result) { + // Keep the tool index and queue position in the same admission order. + let queue = self + .queue + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let index = observe(); + let reservation = (|| { + let permit = Arc::clone(&self.capacity) + .try_acquire_owned() + .map_err(|_| LunaSamplerError::QueueFull)?; + let (sender, request) = oneshot::channel(); + queue + .send(QueuedRequest { + request, + _permit: permit, + }) + .map_err(|_| LunaSamplerError::Superseded)?; + Ok(Reservation { + request: sender, + generation: Arc::downgrade(&self.generation), + }) + })(); + (index, reservation) + } +} + +impl Reservation { + pub(super) fn submit(self, request: ConversationRequest) { + if let Err(request) = self.request.send(request) { + let _ = request.ready.send(Err(LunaSamplerError::Superseded)); + } + } +} + +async fn run( + mut queue: mpsc::UnboundedReceiver, + sampler: Arc, + generation: Weak<()>, +) { + // Only completed requests have reusable history and a matching reuse key. + let mut committed: Option<(ReuseKey, ConversationState>)> = None; + while let Some(queued) = queue.recv().await { + // A preparation error drops its sender, preserving the order of later observations. + let Ok(request) = queued.request.await else { + continue; + }; + if generation.strong_count() == 0 + || !request.authorization.is_current(&request.thread).await + { + committed = None; + let _ = request.ready.send(Err(LunaSamplerError::Superseded)); + continue; + } + let key = ReuseKey { + authorization: request.authorization.clone(), + instructions: request.sampling.instructions.clone(), + parent_compaction: request.sampling.parent_compaction.clone(), + }; + let mut state = committed + .take() + .filter(|(previous, _)| previous == &key) + .map(|(_, state)| state) + .unwrap_or_default(); + let had_history = state.snapshot().is_some(); + let prepared = request + .sampling + .prepare_retained(&request.evidence, &mut state, sampler.max_input_tokens()) + .filter(|prepared| { + prepared.existing_context_tokens == 0 + || prepared.input_tokens <= request.reset_token_limit + }) + .or_else(|| { + if !had_history { + return None; + } + state = ConversationState::default(); + request.sampling.prepare_retained( + &request.evidence, + &mut state, + sampler.max_input_tokens(), + ) + }); + let Some(mut prepared) = prepared else { + let _ = request.ready.send(Err(LunaSamplerError::InputTooLarge)); + continue; + }; + super::metrics::record_section_costs( + request.metrics.as_deref(), + prepared.section_costs.iter().copied(), + ); + super::metrics::record_request_tokens( + request.metrics.as_deref(), + prepared.existing_context_tokens, + prepared.input_tokens, + ); + drop(request.evidence); + let cursor = prepared.cursor; + let pending_truncations = std::mem::take(&mut prepared.truncations); + let history = sampler.sample_retained(prepared, request.ready).await; + if let Some(history) = history + && generation.strong_count() > 0 + && request.authorization.is_current(&request.thread).await + { + let mut truncations = super::truncation::ClassificationTruncations::default(); + truncations.extend(pending_truncations); + truncations.emit(request.metrics.as_deref()); + state.complete_review(cursor); + state.commit_snapshot(history); + committed = Some((key, state)); + } + // The permit is released only after completion and commit, not after the early score. + drop(queued._permit); + } +} diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/conversation_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/conversation_tests.rs new file mode 100644 index 0000000000..c79f20fea9 --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/async_scorer/conversation_tests.rs @@ -0,0 +1,218 @@ +use super::*; +use crate::async_scorer::authorization::ScoreAuthorization; +use crate::async_scorer::conversation::ConversationBackend; +use crate::async_scorer::conversation::ConversationRequest; +use crate::async_scorer::sampler::LunaSampler; +use crate::async_scorer::sampler::LunaSamplerError; +use crate::async_scorer::sampler::LunaSamplingRequest; +use crate::async_scorer::sampler::tests::sample_request; +use crate::async_scorer::sampler::tests::sampler_config; +use crate::async_scorer::transcript::ContextInput; +use codex_extension_api::ThreadStopInput; +use codex_guardian_context::ContextTarget; +use pretty_assertions::assert_eq; + +async fn pending_request( + fixture: &GuardianFailureFixture, + label: &str, +) -> Result<( + ConversationRequest, + tokio::sync::oneshot::Receiver>, +)> { + let thread = Arc::clone(&fixture.test.codex); + let config = thread + .thread_extension_data() + .get::() + .unwrap(); + let history = thread.conversation_history_snapshot().await; + let instructions = config.render_classifier_instructions(TEST_GUARDIAN_POLICY, ""); + let evidence = config.transcript.collect_context(ContextInput { + target: ContextTarget::Async, + history: history.as_ref(), + root_conversation: &[], + trusted_user_answers: &[], + planned_action: None, + permissions: None, + previous_reviews: None, + trusted_tool: None, + trusted_skill_paths: &[], + node_repl_images: None, + })?; + let (ready, score) = tokio::sync::oneshot::channel(); + let authorization = ScoreAuthorization::current(&thread, &Default::default()).await; + Ok(( + ConversationRequest { + evidence, + reset_token_limit: config.async_classifier_conversation_token_limit, + authorization, + thread, + ready, + metrics: None, + sampling: LunaSamplingRequest { + instructions, + input: Vec::new(), + ..sample_request(label) + }, + }, + score, + )) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn configured_reset_limit_rebuilds_history_without_rejecting_fresh_evidence() -> Result<()> { + skip_if_no_network!(Ok(())); + for (limit, reused) in [(1, false), (100_000, true)] { + let server = responses::start_mock_server().await; + let fixture = GuardianFailureFixture::with_config(&format!( + "[features.guardianv2]\nasync_classifier_mode = 'conversation'\nasync_classifier_conversation_token_limit = {limit}" + )) + .await?; + let sampler = Arc::new(LunaSampler::new(sampler_config(server.uri()))); + let backend = ConversationBackend::new(sampler); + let events = responses::sse(vec![ + ev_assistant_message("score", "low"), + ev_completed("score"), + ]); + let mock = responses::mount_sse_sequence(&server, vec![events; 2]).await; + for index in 0..2 { + let (request, score) = pending_request(&fixture, "parent").await?; + backend.reserve(|| index).1.unwrap().submit(request); + assert_eq!(score.await??, "low"); + } + let requests = mock.requests(); + assert_eq!(requests.len(), 2); + let input = requests[1].input(); + assert_eq!(input.iter().any(|item| item["id"] == "score"), reused); + assert_eq!( + serde_json::to_string(&input)?.contains("TRANSCRIPT DELTA START"), + reused + ); + drop(backend); + fixture.test.codex.shutdown_and_wait().await?; + } + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn admission_orders_delayed_preparation_and_bounds_active_plus_pending_work() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = responses::start_mock_server().await; + let fixture = GuardianFailureFixture::with_config( + "[features.guardianv2]\nasync_classifier_mode = 'conversation'", + ) + .await?; + let sampler = Arc::new(LunaSampler::new(sampler_config(server.uri()))); + fixture + .test + .codex + .thread_extension_data() + .insert(ConversationBackend::new(sampler)); + let events = responses::sse(vec![ + ev_assistant_message("score", "low"), + ev_completed("score"), + ]); + let mock = responses::mount_sse_sequence(&server, vec![events; 3]).await; + let backend = fixture + .test + .codex + .thread_extension_data() + .get::() + .unwrap(); + let store = fixture.test.codex.thread_extension_data(); + let progress = store.get::().unwrap(); + let mut reservations = (0..16) + .map(|index| backend.reserve(|| index).1.unwrap()) + .collect::>(); + fixture.score_tool(ToolName::plain("read_file")).await; + let cached = progress.inspect(/*call_id*/ None); + assert!(cached.has_unscored_failure); + assert_eq!(cached.action_risk, Some(1.0)); + assert_eq!( + cached_approval( + &fixture.registry, + store, + "review action", + /*metrics*/ None + ) + .await, + None + ); + let c = reservations.remove(/*index*/ 2); + let b = reservations.remove(/*index*/ 1); + let a = reservations.remove(/*index*/ 0); + drop(reservations); + let mut scores = Vec::new(); + for (reservation, label) in [(c, "C"), (b, "B"), (a, "A")] { + let (request, score) = pending_request(&fixture, label).await?; + reservation.submit(request); + scores.push(score); + } + for score in scores { + assert_eq!(score.await??, "low"); + } + let turns = mock + .requests() + .iter() + .map(|request| { + request.body_json()["client_metadata"]["parent_turn_id"] + .as_str() + .unwrap() + .to_owned() + }) + .collect::>(); + assert_eq!(turns, ["A", "B", "C"]); + fixture.test.codex.shutdown_and_wait().await?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn queued_authorization_and_stopped_generations_cannot_continue() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = responses::start_mock_server().await; + let fixture = GuardianFailureFixture::with_config( + "[features.guardianv2]\nasync_classifier_mode = 'conversation'", + ) + .await?; + let sampler = Arc::new(LunaSampler::new(sampler_config(server.uri()))); + fixture + .test + .codex + .thread_extension_data() + .insert(ConversationBackend::new(sampler)); + let thread_store = fixture.test.codex.thread_extension_data(); + let backend = thread_store.get::().unwrap(); + let blocker = backend.reserve(|| 0).1.unwrap(); + let stale = backend.reserve(|| 0).1.unwrap(); + let (request, score) = pending_request(&fixture, "stale").await?; + fixture + .test + .codex + .inject_response_items(vec![user_instruction( + "Stop. Do not inspect any more files.", + )]) + .await?; + stale.submit(request); + drop(blocker); + assert!(matches!(score.await?, Err(LunaSamplerError::Superseded))); + let stopped = backend.reserve(|| 0).1.unwrap(); + let (request, score) = pending_request(&fixture, "stopped").await?; + drop(backend); + fixture.registry.thread_lifecycle_contributors()[0] + .on_thread_stop(ThreadStopInput { + session_store: &fixture.session_store, + thread_store, + }) + .await; + stopped.submit(request); + assert!(matches!(score.await?, Err(LunaSamplerError::Superseded))); + assert!( + server + .received_requests() + .await + .unwrap() + .iter() + .all(|request| request.method.as_str() != "POST") + ); + fixture.test.codex.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs index f6af94c22d..86020d125e 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs @@ -16,6 +16,7 @@ use codex_extension_api::SkillInvocationContributor; use codex_extension_api::SkillInvocationInput; use codex_extension_api::ThreadLifecycleContributor; use codex_extension_api::ThreadStartInput; +use codex_extension_api::ThreadStopInput; use codex_extension_api::ToolFinishInput; use codex_extension_api::ToolLifecycleContributor; use codex_extension_api::ToolLifecycleFuture; @@ -23,10 +24,12 @@ use codex_extension_api::ToolStartInput; use codex_features::Feature; use codex_login::AuthManager; use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::openai_models::AsyncClassifierMode; use codex_protocol::openai_models::ModelInfo; use codex_protocol::protocol::has_full_access; use super::config::GuardianV2Config; +use super::conversation::ConversationBackend; use super::sampler::LunaSampler; use super::score::GuardianV2ScoreProgress; use super::trusted_skills::TrustedSkillRoots; @@ -83,10 +86,22 @@ impl ThreadLifecycleContributor for GuardianV2Extension { .get_or_init(NodeReplReviewEvidence::default) .enable_image_capture(); } + let mode = guardian_config.classifier_mode( + model + .as_ref() + .and_then(|model| model.model_messages.as_ref()) + .and_then(|messages| messages.guardian_v2.as_ref()), + ); + input.thread_store.remove::(); input.thread_store.remove::(); let sampler = input .thread_store .get_or_init(|| LunaSampler::new(sampler_config)); + if mode == AsyncClassifierMode::Conversation { + input + .thread_store + .insert(ConversationBackend::new(Arc::clone(&sampler))); + } input.thread_store.insert(guardian_config); input.thread_store.insert(GuardianV2ScoreProgress::new( input.extension_metrics.clone(), @@ -121,6 +136,11 @@ impl ThreadLifecycleContributor for GuardianV2Extension { } }) } + fn on_thread_stop<'a>(&'a self, input: ThreadStopInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + input.thread_store.remove::(); + }) + } } impl SkillInvocationContributor for GuardianV2Extension { diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs index e9cb616503..4dd694228f 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs @@ -71,6 +71,9 @@ use pretty_assertions::assert_eq; use serde_json::json; use crate::async_scorer::authorization::ScoreAuthorization; + +#[path = "conversation_tests.rs"] +mod conversation; use crate::async_scorer::config::CLASSIFICATION_OUTPUT_INSTRUCTIONS; use crate::async_scorer::config::DEFAULT_PARENT_COMPACTION_TOKENS; use crate::async_scorer::config::GuardianV2Config; @@ -1830,6 +1833,8 @@ async fn contributor_uses_model_defaults_and_preserves_local_overrides() -> Resu skip_if_no_network!(Ok(())); let model_defaults = GuardianV2ModelConfig { + async_classifier_mode: None, + async_classifier_conversation_token_limit: None, classifier_instructions: Some("Use the experimental model-owned prompt.".to_owned()), review_threshold_basis_points: Some(6_000), max_tool_call_lag: Some(2), diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs b/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs index 92e7828fe1..fcea468a4d 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs @@ -27,6 +27,7 @@ pub(super) fn sampler_failure_reason(error: &LunaSamplerError) -> &'static str { LunaSamplerError::Superseded => "superseded", LunaSamplerError::IncompatibleCompaction => "incompatible_compaction", LunaSamplerError::InputTooLarge => "input_too_large", + LunaSamplerError::QueueFull => "queue_full", LunaSamplerError::Api(error) => match error { ApiError::Transport(TransportError::Http { status, .. }) | ApiError::Api { status, .. } => match status.as_u16() { @@ -110,3 +111,48 @@ pub(super) fn record_fast_decision( &[("decision", decision), ("reason", reason)], ); } + +pub(super) fn record_section_costs( + metrics: Option<&dyn ExtensionMetrics>, + costs: impl IntoIterator, +) { + let Some(metrics) = metrics else { + return; + }; + for (section, cost) in costs { + for (measurement, value) in cost.measurements() { + metrics.histogram_with_boundaries( + codex_guardian_context::SECTION_COST_METRIC, + i64::try_from(value).unwrap_or(i64::MAX), + codex_guardian_context::SECTION_COST_BOUNDARIES, + &[ + ("target", "async"), + ("section", section), + ("measurement", measurement), + ], + ); + } + } +} + +pub(super) fn record_request_tokens( + metrics: Option<&dyn ExtensionMetrics>, + existing: usize, + total: usize, +) { + let Some(metrics) = metrics else { + return; + }; + for (component, tokens) in [ + ("existing_context", existing), + ("new_input", total.saturating_sub(existing)), + ("total", total), + ] { + metrics.histogram_with_boundaries( + codex_guardian_context::REQUEST_TOKENS_METRIC, + i64::try_from(tokens).unwrap_or(i64::MAX), + codex_guardian_context::REQUEST_TOKENS_BOUNDARIES, + &[("target", "async"), ("component", component)], + ); + } +} diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/mod.rs b/codex-rs/ext/guardian-v2/src/async_scorer/mod.rs index dea634a6ea..d5207f7c10 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/mod.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/mod.rs @@ -3,13 +3,12 @@ mod approval; mod authorization; mod classification; mod config; +mod conversation; mod coverage; mod extension; mod metrics; mod observation; mod parent_compaction; -// The treatment backend wires this builder into sampling in the next PR. -#[allow(dead_code)] mod request; mod sampler; mod score; diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/observation.rs b/codex-rs/ext/guardian-v2/src/async_scorer/observation.rs index 7c1189a133..3c13957a90 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/observation.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/observation.rs @@ -26,6 +26,7 @@ use super::action::GuardianAction; use super::authorization::ScoreAuthorization; use super::classification::Classification; use super::config::GuardianV2Config; +use super::conversation::ConversationBackend; use super::coverage::scores_tool; use super::extension::GuardianV2Extension; use super::metrics::record_classification; @@ -121,7 +122,28 @@ impl GuardianV2Extension { let context_mode = GuardianContextMode::from_history(input.conversation_history.as_ref()); let analytics = input.session_store.get::(); let sampled_at = SystemTime::now(); - let tool_call_index = score_progress.observe(&input); + let (tool_call_index, reservation) = match input.thread_store.get::() { + Some(conversation) => { + let (index, reservation) = conversation.reserve(|| score_progress.observe(&input)); + (index, reservation.map(Some)) + } + None => (score_progress.observe(&input), Ok(None)), + }; + let reservation = match reservation { + Ok(reservation) => reservation, + Err(error) => { + score_progress.invalidate(tool_call_index); + score_progress.fail_closed(sampled_at); + record_classification( + metrics.as_deref(), + context_mode, + classification_started_at.elapsed(), + "failure", + Some(super::metrics::sampler_failure_reason(&error)), + ); + return; + } + }; let event_sink = Arc::clone(&self.event_sink); let thread_id = input.thread_store.level_id().to_owned(); let turn_id = input.turn_id.to_owned(); @@ -328,6 +350,7 @@ impl GuardianV2Extension { }; let score_authorization = ScoreAuthorization::current(&thread, &permissions).await; let classification = Classification { + reservation, classification_started_at, sampler, guardian_config, diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/request.rs b/codex-rs/ext/guardian-v2/src/async_scorer/request.rs index 9cdf67cc4d..3e5f27c6fa 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/request.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/request.rs @@ -35,7 +35,7 @@ impl LunaSamplingRequest { }) { return None; } - let (mut context, _) = evidence.compose(cursor).ok()?; + let (mut context, _) = evidence.clone().compose(cursor).ok()?; context.retain_new_instructions( conversation .snapshot() @@ -53,6 +53,7 @@ impl LunaSamplingRequest { }); } let truncations = std::mem::take(&mut context.truncations); + let section_costs = context.section_costs().collect(); let history = conversation.snapshot().map(ConversationCheckpoint::history); let existing_context_tokens = history.map_or(0, |history| { history @@ -80,6 +81,7 @@ impl LunaSamplingRequest { input, cursor: next_cursor, truncations, + section_costs, existing_context_tokens, input_tokens, }) @@ -92,6 +94,7 @@ pub(super) struct PreparedRequest<'a> { pub(super) input: Vec, pub(super) cursor: TranscriptCursor, pub(super) truncations: Vec, + pub(super) section_costs: Vec<(&'static str, codex_guardian_context::SectionCost)>, pub(super) existing_context_tokens: usize, pub(super) input_tokens: usize, } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/retained_sampling_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/retained_sampling_tests.rs index 86e57dbcc0..a4e2baab40 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/retained_sampling_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/retained_sampling_tests.rs @@ -47,6 +47,7 @@ async fn retained_websocket_requests_preserve_model_output_and_ids() -> Result<( transcript_entry_count: 0, }, truncations: Vec::new(), + section_costs: Vec::new(), }; let (completed, score) = tokio::join!(sampler.sample_retained(prepared, ready), score); assert_eq!(score??, "low"); @@ -142,6 +143,7 @@ async fn retention_overflow_or_incomplete_stream_keeps_early_score() -> Result<( transcript_entry_count: 0, }, truncations: Vec::new(), + section_costs: Vec::new(), }; let (history, score) = tokio::join!(sampler.sample_retained(prepared, ready), score); assert_eq!(score??, "low"); diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs index ec29d65cb2..e5bad22b83 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs @@ -115,6 +115,9 @@ pub enum LunaSamplerError { /// The complete classifier input exceeded the model allowance. #[error("Luna input exceeds the complete request budget")] InputTooLarge, + /// The retained conversation already owns its maximum outstanding work. + #[error("Guardian conversation queue is full")] + QueueFull, } struct ActiveRequest { @@ -277,20 +280,11 @@ impl LunaSampler { .into_iter() .map(ResponseItemEnvelope::into_item) .collect(); - if let Some(metrics) = self.config.metrics.as_deref() { - for (component, tokens) in [ - ("existing_context", 0), - ("new_input", total_tokens), - ("total", total_tokens), - ] { - metrics.histogram_with_boundaries( - codex_guardian_context::REQUEST_TOKENS_METRIC, - i64::try_from(tokens).unwrap_or(i64::MAX), - codex_guardian_context::REQUEST_TOKENS_BOUNDARIES, - &[("target", "async"), ("component", component)], - ); - } - } + super::metrics::record_request_tokens( + self.config.metrics.as_deref(), + /*existing*/ 0, + total_tokens, + ); // Oversized classifications defer to sync with the existing failure score. if total_tokens > self.config.max_input_tokens.saturating_sub(/*rhs*/ 256) { return Err(LunaSamplerError::InputTooLarge); @@ -322,8 +316,11 @@ impl LunaSampler { .await } + pub(super) fn max_input_tokens(&self) -> usize { + self.config.max_input_tokens + } + /// Publishes the early score and returns only validated, completed history. - #[cfg_attr(not(test), allow(dead_code))] pub(super) async fn sample_retained( &self, prepared: PreparedRequest<'_>, diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs index 0646955ff3..2affc8c3df 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs @@ -161,6 +161,7 @@ impl SamplingExecution { | LunaSamplerError::Superseded | LunaSamplerError::IncompatibleCompaction | LunaSamplerError::InputTooLarge + | LunaSamplerError::QueueFull | LunaSamplerError::Api( ApiError::Transport( TransportError::Build(_) diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs index 674f859210..aa9eb7e9b8 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs @@ -202,7 +202,7 @@ pub(in crate::async_scorer) async fn proxy_websocket_servers_with_http( Ok(format!("http://{address}/v1")) } -pub(super) fn sampler_config(base_url: String) -> LunaSamplerConfig { +pub(in crate::async_scorer) fn sampler_config(base_url: String) -> LunaSamplerConfig { LunaSamplerConfig { workspace_routing: codex_model_provider::WorkspaceRoutingContext::new( "https://chatgpt.com/backend-api".into(), @@ -262,7 +262,7 @@ fn assert_classifier_instructions(request: &serde_json::Value) { ); } -pub(super) fn sample_request(parent_turn_id: &str) -> LunaSamplingRequest { +pub(in crate::async_scorer) fn sample_request(parent_turn_id: &str) -> LunaSamplingRequest { LunaSamplingRequest { parent_response_id: None, instructions: classifier_instructions(), diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs b/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs index 016a6fa5eb..23a25653f7 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs @@ -88,15 +88,7 @@ impl TranscriptConfig { &self, input: ContextInput<'_>, ) -> Result { - let collected = self.collect_context(input)?; - let transcript = collected.profile.render_transcript( - collected.context.transcript_entries(), - /*entry_number_offset*/ 0, - ); - let mut context = collected - .context - .compose(ContextPresentation::Async, transcript)?; - // Each snapshot deduplicates only against the evidence in its own request. + let (mut context, _) = self.collect_context(input)?.compose(/*cursor*/ None)?; context.deduplicate_transcript_instructions(); Ok(context) } @@ -186,7 +178,7 @@ impl CollectedTranscript { } pub(super) fn compose( - &self, + self, cursor: Option, ) -> Result<(ComposedContext, TranscriptCursor), SectionError> { let (selection, next) = self.select(cursor); @@ -197,7 +189,7 @@ impl CollectedTranscript { } }; let transcript = self.profile.render_transcript(entries, offset); - let context = self.context.clone().compose(presentation, transcript)?; + let context = self.context.compose(presentation, transcript)?; Ok((context, next)) } } diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index f6b660d5dd..22d04dc54a 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -149,6 +149,14 @@ pub struct GuardianV2ReviewScopeConfigToml { #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, JsonSchema)] #[serde(deny_unknown_fields)] pub struct GuardianV2ConfigToml { + /// Classifier experiment override. Otherwise use the model default, then snapshots. + #[serde(skip_serializing_if = "Option::is_none")] + pub async_classifier_mode: Option, + /// Reset retained async history when the next request exceeds this token estimate. + /// Defaults to 100,000. Fresh requests remain subject to the model's input limit. + #[serde(skip_serializing_if = "Option::is_none")] + #[schemars(range(min = 1))] + pub async_classifier_conversation_token_limit: Option, #[serde(skip_serializing_if = "Option::is_none")] pub enabled: Option, /// Legacy setting retained for config compatibility; the backend now controls Guardian billing. @@ -215,6 +223,11 @@ where } let transcript = config.transcript.as_ref(); + if config.async_classifier_conversation_token_limit == Some(0) { + return Err(serde::de::Error::custom( + "Guardian v2 async_classifier_conversation_token_limit must be positive", + )); + } let message_entry = transcript.and_then(|value| value.max_message_entry_tokens); let tool_entry = transcript.and_then(|value| value.max_tool_entry_tokens); let message_total = transcript.and_then(|value| value.max_message_transcript_tokens); diff --git a/codex-rs/features/src/tests.rs b/codex-rs/features/src/tests.rs index ae0a10bd7e..f9b235db74 100644 --- a/codex-rs/features/src/tests.rs +++ b/codex-rs/features/src/tests.rs @@ -232,6 +232,7 @@ persist_scores = true classifier_instructions = "Review this action" review_threshold = 0.65 max_tool_call_lag = 2 +async_classifier_conversation_token_limit = 120000 reasoning_effort = "minimal" max_action_tokens = 512 max_classifier_instruction_tokens = 256 @@ -257,6 +258,8 @@ max_recent_non_user_entries = 12 assert_eq!( features.guardianv2, Some(FeatureToml::Config(crate::GuardianV2ConfigToml { + async_classifier_mode: None, + async_classifier_conversation_token_limit: Some(120_000), enabled: Some(true), free_guardian: Some(true), thread_context: None, @@ -317,6 +320,7 @@ fn guardian_v2_feature_config_rejects_invalid_settings() { "review_threshold = -0.1", "review_threshold = 1.1", "review_threshold = nan", + "async_classifier_conversation_token_limit = 0", "transcript.max_recent_non_user_entries = 0", "transcript.max_message_entry_tokens = 200\ntranscript.max_message_transcript_tokens = 100", "transcript.max_tool_entry_tokens = 200\ntranscript.max_tool_transcript_tokens = 100", diff --git a/codex-rs/guardian-context/src/retained_instructions.rs b/codex-rs/guardian-context/src/retained_instructions.rs index b2aae41076..9e1299b575 100644 --- a/codex-rs/guardian-context/src/retained_instructions.rs +++ b/codex-rs/guardian-context/src/retained_instructions.rs @@ -2,8 +2,8 @@ //! The host owns storage and lifecycle. Omitted records remain explicit; restrictions //! are never truncated into partial permissions, and retained source order is preserved. //! Section omissions do not change fast-approval eligibility. -//! Sync delivery compares against admitted reviewer history, so forks and compaction -//! need no separate retained-evidence cursor. Async deduplication is request-local. +//! Retained delivery compares against admitted reviewer history, so forks and compaction +//! need no separate retained-evidence cursor. Snapshot deduplication is request-local. use std::collections::HashSet; @@ -175,7 +175,7 @@ impl ComposedContext { } } - /// Each async sample carries its own originals and ordering guidance. + /// Each independent async sample carries its own originals and ordering guidance. pub fn deduplicate_transcript_instructions(&mut self) { self.remove_delivered_instructions(&[]); } diff --git a/codex-rs/protocol/src/openai_models.rs b/codex-rs/protocol/src/openai_models.rs index e5f02ac9e8..b2dde1499e 100644 --- a/codex-rs/protocol/src/openai_models.rs +++ b/codex-rs/protocol/src/openai_models.rs @@ -46,6 +46,7 @@ mod guardian_v2; mod reasoning_effort; pub use access_programs::ModelAccessPrograms; +pub use guardian_v2::AsyncClassifierMode; pub use guardian_v2::GuardianV2ModelConfig; pub use guardian_v2::GuardianV2TranscriptModelConfig; diff --git a/codex-rs/protocol/src/openai_models/guardian_v2.rs b/codex-rs/protocol/src/openai_models/guardian_v2.rs index 953a89a4e7..c13a0dca9b 100644 --- a/codex-rs/protocol/src/openai_models/guardian_v2.rs +++ b/codex-rs/protocol/src/openai_models/guardian_v2.rs @@ -5,9 +5,24 @@ use ts_rs::TS; use super::ReasoningEffort; +/// Async classifier experiment, resolved once for each parent task. +#[derive(Debug, Default, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, TS, JsonSchema)] +#[serde(rename_all = "snake_case")] +#[ts(rename_all = "snake_case")] +pub enum AsyncClassifierMode { + #[default] + Snapshot, + Conversation, +} + /// Optional model-owned defaults for Guardian v2 classification experiments. #[derive(Debug, Default, Serialize, Deserialize, Clone, PartialEq, Eq, TS, JsonSchema)] pub struct GuardianV2ModelConfig { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub async_classifier_mode: Option, + /// Reset retained async history when the next request exceeds this token estimate. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub async_classifier_conversation_token_limit: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub classifier_instructions: Option, #[serde(default, skip_serializing_if = "Option::is_none")]