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")]