diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/approval.rs b/codex-rs/ext/guardian-v2/src/async_scorer/approval.rs index b6440cb17d..e61fda2bff 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/approval.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/approval.rs @@ -4,11 +4,11 @@ use super::authorization::ScoreAuthorization; use super::config::GuardianV2Config; use super::coverage::GuardianPolicy; -use super::extension::GuardianV2ScoreProgress; use super::metrics::TOOL_CALL_LAG_METRIC; use super::metrics::record_fast_decision; use super::parent_compaction::select_parent_compaction; use super::sampler::LunaSampler; +use super::score::GuardianV2ScoreProgress; use codex_core::CodexThread; use codex_core::ThreadManager; use codex_core::context::GuardianContextMode; diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/classification.rs b/codex-rs/ext/guardian-v2/src/async_scorer/classification.rs new file mode 100644 index 0000000000..6e83d1542a --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/async_scorer/classification.rs @@ -0,0 +1,400 @@ +//! Builds context and samples a captured observation in the background. +//! Successful results require current authorization and a newer sample timestamp. + +use std::collections::BTreeMap; +use std::sync::Arc; +use std::sync::atomic::Ordering; +use std::time::Instant; +use std::time::SystemTime; + +use codex_analytics::AnalyticsEventsClient; +use codex_analytics::GuardianV2Event; +use codex_analytics::GuardianV2EventKind; +use codex_core::CodexThread; +use codex_core::GuardianAuthorizationVersion; +use codex_core::GuardianRootSnapshot; +use codex_core::ThreadManager; +use codex_core::config::Config; +use codex_core::context::ContextualUserFragment; +use codex_core::context::GuardianContextMode; +use codex_core::context::GuardianReviewEvidenceFragment; +use codex_core::context::GuardianReviewEvidenceRecord; +use codex_extension_api::ConversationHistorySnapshot; +use codex_extension_api::ExtensionEventSink; +use codex_extension_api::ExtensionMetrics; +use codex_extension_api::ExtensionWarning; +use codex_extension_api::McpToolContext; +use codex_extension_api::ResponseItem; +use codex_guardian_context::ContextTarget; +use codex_guardian_context::PlannedAction; +use codex_guardian_context::PlannedActionKind; +use codex_guardian_context::PreviousReviews; +use codex_guardian_context::ReviewEvidence; +use codex_guardian_context::render_review_evidence; +use codex_history::RolloutItem; +use codex_model_provider::create_model_provider; +use codex_protocol::models::ContentItem; +use codex_protocol::openai_models::ModelInfo; +use codex_protocol::security_risk::SecurityRiskScore; + +use super::authorization::ScoreAuthorization; +use super::config::GuardianV2Config; +use super::metrics::record_classification; +use super::metrics::record_classification_risk; +use super::metrics::sampler_failure_reason; +use super::sampler::LunaSampler; +use super::sampler::LunaSamplerError; +use super::sampler::LunaSamplingRequest; +use super::score::GuardianV2ScoreProgress; +use super::score::record_fail_closed_score; +use super::transcript::ContextInput; +use super::truncation::ClassificationTruncations; +use super::trusted_skills::TrustedSkillInvocations; +use super::trusted_tools::trusted_tool_context; + +pub(super) struct Classification { + pub(super) classification_started_at: Instant, + pub(super) sampler: Arc, + pub(super) guardian_config: GuardianV2Config, + pub(super) score_progress: Arc, + pub(super) parent_model: Option>, + pub(super) metrics: Option>, + pub(super) analytics: Option>, + pub(super) sampled_at: SystemTime, + pub(super) tool_call_index: usize, + pub(super) event_sink: Arc, + pub(super) thread_id: String, + pub(super) turn_id: String, + pub(super) root_turn_id: Option, + pub(super) parent_response_id: Option, + pub(super) manager: Arc, + pub(super) thread: Arc, + pub(super) config: Arc, + pub(super) context_mode: GuardianContextMode, + pub(super) parent_compaction: Option, + pub(super) parent_compaction_hash: Option, + pub(super) call_id: String, + pub(super) mcp_tool: Option, + pub(super) planned_action: String, + pub(super) review_model_override: Option, + pub(super) sync_reviews: Vec>, + pub(super) trusted_user_inputs: Vec, + pub(super) authorization_version: GuardianAuthorizationVersion, + pub(super) history: Arc, + pub(super) local_trusted_skill_paths: Vec, + pub(super) node_repl_images: Vec, + pub(super) root_snapshot: Option, + pub(super) score_authorization: ScoreAuthorization, +} + +enum ClassificationOutcome { + Scored, + Superseded, +} + +impl Classification { + pub(super) async fn run(self) { + let Self { + classification_started_at, + sampler, + guardian_config, + score_progress, + parent_model, + metrics, + analytics, + sampled_at, + tool_call_index, + event_sink, + thread_id, + turn_id, + root_turn_id, + parent_response_id, + manager, + thread, + config, + context_mode, + parent_compaction, + parent_compaction_hash, + call_id, + mcp_tool, + planned_action, + review_model_override, + sync_reviews, + trusted_user_inputs, + authorization_version, + history, + local_trusted_skill_paths, + node_repl_images, + root_snapshot, + score_authorization, + } = self; + let mut truncations = ClassificationTruncations::default(); + let trusted_tool_context = match mcp_tool.as_ref() { + Some(tool) => { + trusted_tool_context(tool.tool_info(), tool.source(), &manager, &config).await + } + None => None, + }; + let root_snapshot = if context_mode == GuardianContextMode::ThreadOwned { + root_snapshot + } else { + thread.guardian_root_snapshot().await + }; + let mut trusted_skills = TrustedSkillInvocations::default(); + for path in local_trusted_skill_paths.iter().chain( + root_snapshot + .as_ref() + .into_iter() + .flat_map(|snapshot| snapshot.trusted_skill_paths.iter()), + ) { + trusted_skills.record(path.clone()); + } + let trusted_skill_paths = trusted_skills.into_paths(); + let root_authorization_version = root_snapshot + .as_ref() + .map(|snapshot| snapshot.authorization_version); + let root_conversation = root_snapshot.map(|snapshot| snapshot.messages); + let score_authorization = ScoreAuthorization { + local: authorization_version, + root: root_authorization_version, + model: parent_model.clone(), + ..score_authorization + }; + let action_section = PlannedAction { + json: planned_action.clone(), + tool_descriptions: None, + kind: PlannedActionKind::Command, + reason: None, + }; + let review_fragments = sync_reviews + .iter() + .filter(|review| { + review.authorization_version == authorization_version + && review.root_authorization_version == root_authorization_version + }) + .map(|review| { + let review = render_review_evidence(ReviewEvidence { + correlation: &review.correlation, + decision: &review.decision, + action: &review.action, + rationale: review.rationale.as_deref(), + }); + truncations.extend(review.truncations); + GuardianReviewEvidenceFragment::new(review.body).render() + }) + .collect::>(); + let transcript = + PreviousReviews::try_from_fragments(review_fragments).and_then(|reviews| { + guardian_config.transcript.build_context(ContextInput { + target: ContextTarget::Async, + history: history.as_ref(), + root_conversation: root_conversation.as_deref().unwrap_or_default(), + trusted_user_answers: &trusted_user_inputs, + planned_action: Some(&action_section), + previous_reviews: Some(&reviews), + trusted_tool: trusted_tool_context.as_ref(), + trusted_skill_paths: &trusted_skill_paths, + node_repl_images: Some(&node_repl_images), + }) + }); + let mut transcript = match transcript { + Ok(transcript) => transcript, + Err(error) => { + record_fail_closed_score(thread.thread_extension_data(), sampled_at); + record_classification( + metrics.as_deref(), + classification_started_at.elapsed(), + "failure", + Some("context_build_error"), + ); + event_sink.emit_warning(ExtensionWarning { + thread_id, + turn_id: Some(turn_id), + message: format!("Guardian V2 context collection failed: {error}"), + }); + return; + } + }; + 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), + ], + ); + } + } + } + let classification_input = transcript.into_messages(); + let mut failure_reason = "invalid_output"; + let mut classification_risk = None; + let mut classification_finished_at = None; + let result: Result = async { + let review_model_messages = if config.guardian_policy_config.is_none() { + let review_model_id = review_model_override.as_deref().unwrap_or_else(|| { + create_model_provider( + config.model_provider.clone(), + Some(manager.auth_manager()), + ) + .approval_review_preferred_model() + }); + let review_model = manager + .get_models_manager() + .get_model_info(review_model_id, &config.to_models_manager_config()) + .await; + if review_model.used_fallback_model_metadata && review_model_override.is_none() { + parent_model + .as_ref() + .and_then(|model| model.model_messages.clone()) + } else { + review_model.model_messages + } + } else { + None + }; + let policy = config.resolve_guardian_policy(review_model_messages.as_ref()); + let instructions = guardian_config.render_classifier_instructions(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 + { + Ok(output) => output, + Err(LunaSamplerError::Superseded) => { + return Ok(ClassificationOutcome::Superseded); + } + Err(error) => { + failure_reason = sampler_failure_reason(&error); + return Err(error.to_string()); + } + }; + let (action_risk, risk_level) = match output.as_str() { + "high" => (1.0, "high"), + "low" => (0.0, "low"), + _ => return Err("invalid Guardian V2 classification".to_owned()), + }; + classification_risk = Some(risk_level); + failure_reason = "action_deserialization_error"; + let score = SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), action_risk)]), + call_id: Some(call_id.clone()), + action: Some( + serde_json::from_str(&planned_action).map_err(|error| error.to_string())?, + ), + sampled_at: Some(sampled_at.into()), + }; + if score_authorization != ScoreAuthorization::current(&thread).await { + return Ok(ClassificationOutcome::Superseded); + } + let accepted = { + let mut scored_authorization = score_progress + .authorization + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let accepted = + thread + .thread_extension_data() + .insert_if(score.clone(), |previous| { + previous.is_none_or(|previous| previous.sampled_at < score.sampled_at) + }); + if accepted { + *scored_authorization = Some(score_authorization); + } + accepted + }; + tracing::info!( + %thread_id, + %turn_id, + %call_id, + tool_call_index, + action_risk = score.scores.get("action_risk").copied(), + review_threshold = guardian_config.review_threshold, + sampled_at = ?score.sampled_at, + accepted, + "Guardian V2 classification result" + ); + if !accepted { + return Ok(ClassificationOutcome::Superseded); + } + score_progress + .latest_scored_tool_call + .fetch_max(tool_call_index, Ordering::Release); + classification_finished_at = Some(Instant::now()); + record_classification_risk(metrics.as_deref(), output.as_str()); + if guardian_config.persist_scores + && !config.ephemeral + && let Err(error) = thread + .append_rollout_items(&[RolloutItem::SecurityRiskScore(score)]) + .await + { + tracing::warn!( + %thread_id, + %turn_id, + %call_id, + %error, + "failed to persist Guardian V2 classification result" + ); + } + Ok(ClassificationOutcome::Scored) + } + .await; + if result.is_err() { + record_fail_closed_score(thread.thread_extension_data(), sampled_at); + } + let duration = classification_finished_at + .map(|finished_at: Instant| finished_at.duration_since(classification_started_at)) + .unwrap_or_else(|| classification_started_at.elapsed()); + let outcome = match &result { + Ok(ClassificationOutcome::Scored) => "success", + Ok(ClassificationOutcome::Superseded) => "superseded", + Err(_) => "failure", + }; + record_classification( + metrics.as_deref(), + duration, + outcome, + result.is_err().then_some(failure_reason), + ); + if let Some(analytics) = analytics { + analytics.track_guardian_v2_event(GuardianV2Event { + thread_id: thread_id.clone(), + turn_id: turn_id.clone(), + item_id: Some(call_id), + model: parent_model.as_ref().map(|model| model.slug.clone()), + occurred_at_ms: codex_analytics::now_unix_millis(), + kind: GuardianV2EventKind::Classification { + outcome, + risk_level: classification_risk, + duration_ms: u64::try_from(duration.as_millis()).unwrap_or(u64::MAX), + }, + }); + } + if matches!(result, Ok(ClassificationOutcome::Scored)) { + truncations.emit(metrics.as_deref()); + } + if let Err(error) = result { + event_sink.emit_warning(ExtensionWarning { + thread_id, + turn_id: Some(turn_id), + message: format!("Guardian V2 risk scoring failed: {error}"), + }); + } + } +} 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 964329567c..ebdc867f11 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs @@ -1,27 +1,14 @@ -use super::transcript::ContextInput; -use codex_core::context::ContextualUserFragment; -use std::collections::BTreeMap; -use std::collections::BTreeSet; +//! Registers Guardian classification and handles thread, skill, and tool lifecycle hooks. + use std::sync::Arc; -use std::sync::Mutex; use std::sync::Weak; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::time::Instant; -use std::time::SystemTime; -use codex_analytics::AnalyticsEventsClient; -use codex_analytics::GuardianV2Event; -use codex_analytics::GuardianV2EventKind; use codex_core::ThreadManager; use codex_core::config::Config; -use codex_core::context::GuardianContextMode; use codex_core::context::GuardianReviewEvidence; use codex_core::context::NodeReplReviewEvidence; -use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionEventSink; use codex_extension_api::ExtensionFuture; -use codex_extension_api::ExtensionMetrics; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ExtensionWarning; use codex_extension_api::GuardianV2Enabled; @@ -32,72 +19,23 @@ use codex_extension_api::ThreadStartInput; use codex_extension_api::ToolFinishInput; use codex_extension_api::ToolLifecycleContributor; use codex_extension_api::ToolLifecycleFuture; -use codex_extension_api::ToolPayload; use codex_extension_api::ToolStartInput; use codex_features::Feature; -use codex_guardian_context::ContextTarget; -use codex_guardian_context::PlannedAction; -use codex_guardian_context::PlannedActionKind; -use codex_history::RolloutItem; use codex_login::AuthManager; -use codex_model_provider::create_model_provider; -use codex_protocol::ThreadId; use codex_protocol::config_types::ApprovalsReviewer; -use codex_protocol::mcp::is_node_repl_backed_server; -use codex_protocol::openai_models::GuardianScope; use codex_protocol::openai_models::ModelInfo; use codex_protocol::protocol::has_full_access; -use codex_protocol::security_risk::SecurityRiskScore; -use super::action::ActionRenderError; -use super::action::GuardianAction; -use super::authorization::ScoreAuthorization; use super::config::GuardianV2Config; -use super::coverage::UnscoredAction; -use super::metrics::record_classification; -use super::metrics::record_classification_risk; -use super::metrics::sampler_failure_reason; -use super::parent_compaction::ParentCompactionError; -use super::parent_compaction::select_parent_compaction; use super::sampler::LunaSampler; -use super::sampler::LunaSamplerError; -use super::sampler::LunaSamplingRequest; -use super::truncation::ClassificationTruncations; -use super::trusted_skills::TrustedSkillInvocations; +use super::score::GuardianV2ScoreProgress; use super::trusted_skills::TrustedSkillRoots; -use super::trusted_tools::trusted_tool_context; -use super::wrapper_lag::WrapperLag; -use codex_core::context::GuardianReviewEvidenceFragment; -use codex_guardian_context::PreviousReviews; -use codex_guardian_context::ReviewEvidence; -use codex_guardian_context::render_review_evidence; - -enum ClassificationOutcome { - Scored, - Superseded, -} - -#[derive(Default)] -pub(super) struct GuardianV2ScoreProgress { - pub(super) wrapper_lag: WrapperLag, - pub(super) latest_tool_call: AtomicUsize, - // Setup and reset calls must not consume the first JS execution allowance. - pub(super) js_executions: AtomicUsize, - pub(super) latest_scored_tool_call: AtomicUsize, - pub(super) latest_failed_tool_call: AtomicUsize, - // Keep overflow attached to each active call even after a newer score succeeds. - // The host's finish callback removes entries on completion, failure, or cancellation. - pub(super) oversized_tool_calls: Mutex>, - // Serialize successful score publication with its authorization metadata. - pub(super) authorization: Mutex>, - metrics: Option>, -} #[derive(Clone)] -struct GuardianV2Extension { +pub(super) struct GuardianV2Extension { auth_manager: Arc, - event_sink: Arc, - thread_manager: Weak, + pub(super) event_sink: Arc, + pub(super) thread_manager: Weak, } impl ThreadLifecycleContributor for GuardianV2Extension { @@ -230,578 +168,6 @@ impl ToolLifecycleContributor for GuardianV2Extension { } } -impl GuardianV2Extension { - fn record_fail_closed_score(thread_store: &ExtensionData, sampled_at: SystemTime) { - let score = SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), 1.0)]), - call_id: None, - action: None, - sampled_at: Some(sampled_at.into()), - }; - thread_store.insert_if(score.clone(), |previous| { - previous.is_none_or(|previous| previous.sampled_at <= score.sampled_at) - }); - } - - async fn score_tool(&self, input: ToolStartInput<'_>) { - // Polling a code cell does not introduce another action or age its score. - if input.tool_name.is_default_namespace() && input.tool_name.name == "wait" { - return; - } - let classification_started_at = Instant::now(); - let Some(sampler) = input.thread_store.get::() else { - return; - }; - let Some(guardian_config) = input.thread_store.get::() else { - return; - }; - let Some(score_progress) = input.thread_store.get::() else { - return; - }; - let parent_model = input.thread_store.get::(); - let policy = guardian_config.policy_for_model(parent_model.as_deref()); - if !policy.scoring_enabled() { - input.thread_store.remove::(); - } - let mcp_server = input - .mcp_tool - .map(|tool| tool.tool_info().server_name.as_str()); - let scope = mcp_server - .map(GuardianScope::for_mcp_server) - .or_else(|| GuardianScope::for_tool(input.tool_name)); - if !policy.scores_tool(input.tool_name, input.payload, scope) { - match policy.unscored_action { - UnscoredAction::Ignore => {} - UnscoredAction::AgeScore => { - let index = score_progress - .latest_tool_call - .fetch_add(/*val*/ 1, Ordering::Relaxed) - .saturating_add(/*rhs*/ 1); - score_progress.wrapper_lag.record(&input, index); - // Unscored permission widening must not reuse an earlier approval score. - if input.tool_name.is_default_namespace() - && input.tool_name.name == "exec_command" - && let ToolPayload::Function { arguments } = input.payload - && let Ok(arguments) = serde_json::from_str::(arguments) - && arguments - .get("sandbox_permissions") - .and_then(serde_json::Value::as_str) - == Some("with_additional_permissions") - { - score_progress - .latest_failed_tool_call - .fetch_max(index, Ordering::Release); - } - } - UnscoredAction::InvalidateScore => { - let index = score_progress - .latest_tool_call - .fetch_add(/*val*/ 1, Ordering::Relaxed) - .saturating_add(/*rhs*/ 1); - score_progress.wrapper_lag.record(&input, index); - score_progress - .latest_failed_tool_call - .fetch_max(index, Ordering::Release); - } - } - return; - } - if input.mcp_tool.is_some_and(|tool| { - let info = tool.tool_info(); - is_node_repl_backed_server(&info.server_name) && info.tool.name == "js" - }) { - score_progress - .js_executions - .fetch_add(/*val*/ 1, Ordering::Relaxed); - } - let metrics = score_progress.metrics.clone(); - let analytics = input.session_store.get::(); - let sampled_at = SystemTime::now(); - let tool_call_index = score_progress - .latest_tool_call - .fetch_add(/*val*/ 1, Ordering::Relaxed) - .saturating_add(/*rhs*/ 1); - score_progress.wrapper_lag.record(&input, tool_call_index); - 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(); - let root_turn_id = input.root_turn_id.map(str::to_owned); - let parent_response_id = input - .turn_store - .get::() - .map(|id| id.0.clone()); - let thread_context: Result<_, String> = async { - let parsed_thread_id = - ThreadId::from_string(&thread_id).map_err(|error| error.to_string())?; - let manager = self - .thread_manager - .upgrade() - .ok_or_else(|| "thread manager is unavailable".to_string())?; - let thread = manager - .get_thread(parsed_thread_id) - .await - .map_err(|error| error.to_string())?; - let config = thread.config().await; - Ok((manager, thread, config)) - } - .await; - let (manager, thread, config) = match thread_context { - Ok(context) => context, - Err(error) => { - score_progress - .latest_failed_tool_call - .fetch_max(tool_call_index, Ordering::Release); - record_classification( - metrics.as_deref(), - classification_started_at.elapsed(), - "failure", - Some("thread_context_error"), - ); - event_sink.emit_warning(ExtensionWarning { - thread_id, - turn_id: Some(turn_id), - message: format!("Guardian V2 risk scoring failed: {error}"), - }); - return; - } - }; - // Use the live reviewer, not the startup config or per-app reviewer overrides. - let snapshot = thread.config_snapshot().await; - if snapshot.full_access - || thread.approvals_reviewer_for_turn(input.turn_id).await == ApprovalsReviewer::User - { - // A skipped call invalidates older scores, including ones still in flight. - score_progress - .latest_failed_tool_call - .fetch_max(tool_call_index, Ordering::Release); - return; - } - // A required model keeps synchronous review outside its CUA allowance. - if !(scope == Some(GuardianScope::ComputerUse) && policy.initial_cua_call) - && parent_model.as_ref().is_some_and(|model| { - config - .config_layer_stack - .requirements() - .auto_review_required_for_model(&model.slug) - }) - { - input.thread_store.remove::(); - return; - } - input.thread_store.insert(GuardianV2Enabled); - let model_defaults = parent_model - .as_ref() - .and_then(|model| model.model_messages.as_ref()) - .and_then(|messages| messages.guardian_v2.as_ref()); - let guardian_config = match guardian_config.with_model_defaults(model_defaults) { - Ok(config) => config, - Err(error) => { - Self::record_fail_closed_score(input.thread_store, sampled_at); - record_classification( - metrics.as_deref(), - classification_started_at.elapsed(), - "failure", - Some("configuration_error"), - ); - self.event_sink.emit_warning(ExtensionWarning { - thread_id: input.thread_store.level_id().to_owned(), - turn_id: Some(input.turn_id.to_owned()), - message: error, - }); - return; - } - }; - if guardian_config.transcript.include_images { - input - .thread_store - .get_or_init(NodeReplReviewEvidence::default) - .enable_image_capture(); - } - input.thread_store.insert(guardian_config.clone()); - let guardian_evidence = input - .thread_store - .get_or_init(GuardianReviewEvidence::default); - let context_mode = guardian_evidence.context_mode(); - let selected_compaction = match select_parent_compaction( - context_mode, - &guardian_config, - input.conversation_history.as_ref(), - &sampler, - parent_model - .as_ref() - .and_then(|model| model.comp_hash.as_deref()), - ) { - Ok(compaction) => compaction, - Err(error) => { - let (outcome, failure_reason) = if error == ParentCompactionError::RequiresSync { - score_progress - .latest_failed_tool_call - .fetch_max(tool_call_index, Ordering::Release); - ("skipped", None) - } else { - ("failure", Some("parent_compaction_error")) - }; - Self::record_fail_closed_score(input.thread_store, sampled_at); - record_classification( - metrics.as_deref(), - classification_started_at.elapsed(), - outcome, - failure_reason, - ); - return; - } - }; - let parent_compaction = selected_compaction.item; - let parent_compaction_hash = selected_compaction.model_hash; - let call_id = input.call_id.to_owned(); - let mcp_tool = input.mcp_tool.cloned(); - let action = GuardianAction { - tool_name: input.tool_name.clone(), - payload: input.payload.clone(), - }; - let planned_action = match action.render(guardian_config.max_action_tokens) { - Ok(text) => text, - Err(ActionRenderError::TooLarge { .. }) => { - score_progress - .oversized_tool_calls - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert(input.call_id.to_owned()); - score_progress - .latest_failed_tool_call - .fetch_max(tool_call_index, Ordering::Release); - Self::record_fail_closed_score(input.thread_store, sampled_at); - record_classification( - metrics.as_deref(), - classification_started_at.elapsed(), - "failure", - Some("input_too_large"), - ); - return; - } - Err(error) => { - score_progress - .latest_failed_tool_call - .fetch_max(tool_call_index, Ordering::Release); - Self::record_fail_closed_score(input.thread_store, sampled_at); - record_classification( - metrics.as_deref(), - classification_started_at.elapsed(), - "failure", - Some("action_serialization_error"), - ); - event_sink.emit_warning(ExtensionWarning { - thread_id, - turn_id: Some(turn_id), - message: format!("Guardian V2 action serialization failed: {error}"), - }); - return; - } - }; - let review_model_override = parent_model - .as_ref() - .and_then(|model| model.auto_review_model_override.clone()); - // Snapshot before spawning so a delayed sample cannot see later reviews. - let sync_reviews = guardian_evidence.snapshot(); - let codex_core::context::GuardianUserInputSnapshot { - fragments: trusted_user_inputs, - authorization_version, - } = guardian_evidence.user_input_snapshot(input.conversation_history.as_ref()); - let history = Arc::clone(&input.conversation_history); - let local_trusted_skill_paths = guardian_evidence.trusted_skill_paths(input.turn_id); - let node_repl_images = if guardian_config.transcript.include_images { - input - .thread_store - .get::() - .map(|evidence| evidence.images()) - .unwrap_or_default() - } else { - Vec::new() - }; - // Capture root evidence before background metadata resolution or model I/O. - // Later root changes invalidate this sample through its captured authorization version. - let root_snapshot = if context_mode == GuardianContextMode::ThreadOwned { - thread.guardian_root_snapshot().await - } else { - None - }; - - let score_authorization = ScoreAuthorization::current(&thread).await; - tokio::spawn(async move { - let mut truncations = ClassificationTruncations::default(); - let trusted_tool_context = match mcp_tool.as_ref() { - Some(tool) => { - trusted_tool_context(tool.tool_info(), tool.source(), &manager, &config).await - } - None => None, - }; - let root_snapshot = if context_mode == GuardianContextMode::ThreadOwned { - root_snapshot - } else { - thread.guardian_root_snapshot().await - }; - let mut trusted_skills = TrustedSkillInvocations::default(); - for path in local_trusted_skill_paths.iter().chain( - root_snapshot - .as_ref() - .into_iter() - .flat_map(|snapshot| snapshot.trusted_skill_paths.iter()), - ) { - trusted_skills.record(path.clone()); - } - let trusted_skill_paths = trusted_skills.into_paths(); - let root_authorization_version = root_snapshot - .as_ref() - .map(|snapshot| snapshot.authorization_version); - let root_conversation = root_snapshot.map(|snapshot| snapshot.messages); - let score_authorization = ScoreAuthorization { - local: authorization_version, - root: root_authorization_version, - model: parent_model.clone(), - ..score_authorization - }; - let action_section = PlannedAction { - json: planned_action.clone(), - tool_descriptions: None, - kind: PlannedActionKind::Command, - reason: None, - }; - let review_fragments = sync_reviews - .iter() - .filter(|review| { - review.authorization_version == authorization_version - && review.root_authorization_version == root_authorization_version - }) - .map(|review| { - let review = render_review_evidence(ReviewEvidence { - correlation: &review.correlation, - decision: &review.decision, - action: &review.action, - rationale: review.rationale.as_deref(), - }); - truncations.extend(review.truncations); - GuardianReviewEvidenceFragment::new(review.body).render() - }) - .collect::>(); - let transcript = - PreviousReviews::try_from_fragments(review_fragments).and_then(|reviews| { - guardian_config.transcript.build_context(ContextInput { - target: ContextTarget::Async, - history: history.as_ref(), - root_conversation: root_conversation.as_deref().unwrap_or_default(), - trusted_user_answers: &trusted_user_inputs, - planned_action: Some(&action_section), - previous_reviews: Some(&reviews), - trusted_tool: trusted_tool_context.as_ref(), - trusted_skill_paths: &trusted_skill_paths, - node_repl_images: Some(&node_repl_images), - }) - }); - let mut transcript = match transcript { - Ok(transcript) => transcript, - Err(error) => { - Self::record_fail_closed_score(thread.thread_extension_data(), sampled_at); - record_classification( - metrics.as_deref(), - classification_started_at.elapsed(), - "failure", - Some("context_build_error"), - ); - event_sink.emit_warning(ExtensionWarning { - thread_id, - turn_id: Some(turn_id), - message: format!("Guardian V2 context collection failed: {error}"), - }); - return; - } - }; - 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), - ], - ); - } - } - } - let classification_input = transcript.into_messages(); - let mut failure_reason = "invalid_output"; - let mut classification_risk = None; - let mut classification_finished_at = None; - let result: Result = async { - let review_model_messages = if config.guardian_policy_config.is_none() { - let review_model_id = review_model_override.as_deref().unwrap_or_else(|| { - create_model_provider( - config.model_provider.clone(), - Some(manager.auth_manager()), - ) - .approval_review_preferred_model() - }); - let review_model = manager - .get_models_manager() - .get_model_info(review_model_id, &config.to_models_manager_config()) - .await; - if review_model.used_fallback_model_metadata && review_model_override.is_none() - { - parent_model - .as_ref() - .and_then(|model| model.model_messages.clone()) - } else { - review_model.model_messages - } - } else { - None - }; - let policy = config.resolve_guardian_policy(review_model_messages.as_ref()); - let instructions = guardian_config.render_classifier_instructions(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 - { - Ok(output) => output, - Err(LunaSamplerError::Superseded) => { - return Ok(ClassificationOutcome::Superseded); - } - Err(error) => { - failure_reason = sampler_failure_reason(&error); - return Err(error.to_string()); - } - }; - let (action_risk, risk_level) = match output.as_str() { - "high" => (1.0, "high"), - "low" => (0.0, "low"), - _ => return Err("invalid Guardian V2 classification".to_owned()), - }; - classification_risk = Some(risk_level); - failure_reason = "action_deserialization_error"; - let score = SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), action_risk)]), - call_id: Some(call_id.clone()), - action: Some( - serde_json::from_str(&planned_action).map_err(|error| error.to_string())?, - ), - sampled_at: Some(sampled_at.into()), - }; - if score_authorization != ScoreAuthorization::current(&thread).await { - return Ok(ClassificationOutcome::Superseded); - } - let accepted = { - let mut scored_authorization = score_progress - .authorization - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - let accepted = - thread - .thread_extension_data() - .insert_if(score.clone(), |previous| { - previous - .is_none_or(|previous| previous.sampled_at < score.sampled_at) - }); - if accepted { - *scored_authorization = Some(score_authorization); - } - accepted - }; - tracing::info!( - %thread_id, - %turn_id, - %call_id, - tool_call_index, - action_risk = score.scores.get("action_risk").copied(), - review_threshold = guardian_config.review_threshold, - sampled_at = ?score.sampled_at, - accepted, - "Guardian V2 classification result" - ); - if !accepted { - return Ok(ClassificationOutcome::Superseded); - } - score_progress - .latest_scored_tool_call - .fetch_max(tool_call_index, Ordering::Release); - classification_finished_at = Some(Instant::now()); - record_classification_risk(metrics.as_deref(), output.as_str()); - if guardian_config.persist_scores - && !config.ephemeral - && let Err(error) = thread - .append_rollout_items(&[RolloutItem::SecurityRiskScore(score)]) - .await - { - tracing::warn!( - %thread_id, - %turn_id, - %call_id, - %error, - "failed to persist Guardian V2 classification result" - ); - } - Ok(ClassificationOutcome::Scored) - } - .await; - if result.is_err() { - Self::record_fail_closed_score(thread.thread_extension_data(), sampled_at); - } - let duration = classification_finished_at - .map(|finished_at: Instant| finished_at.duration_since(classification_started_at)) - .unwrap_or_else(|| classification_started_at.elapsed()); - let outcome = match &result { - Ok(ClassificationOutcome::Scored) => "success", - Ok(ClassificationOutcome::Superseded) => "superseded", - Err(_) => "failure", - }; - record_classification( - metrics.as_deref(), - duration, - outcome, - result.is_err().then_some(failure_reason), - ); - if let Some(analytics) = analytics { - analytics.track_guardian_v2_event(GuardianV2Event { - thread_id: thread_id.clone(), - turn_id: turn_id.clone(), - item_id: Some(call_id), - model: parent_model.as_ref().map(|model| model.slug.clone()), - occurred_at_ms: codex_analytics::now_unix_millis(), - kind: GuardianV2EventKind::Classification { - outcome, - risk_level: classification_risk, - duration_ms: u64::try_from(duration.as_millis()).unwrap_or(u64::MAX), - }, - }); - } - if matches!(result, Ok(ClassificationOutcome::Scored)) { - truncations.emit(metrics.as_deref()); - } - if let Err(error) = result { - event_sink.emit_warning(ExtensionWarning { - thread_id, - turn_id: Some(turn_id), - message: format!("Guardian V2 risk scoring failed: {error}"), - }); - } - }); - } -} - /// Installs feature-gated Guardian V2 tool classification for each thread. pub fn install( registry: &mut ExtensionRegistryBuilder, diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs index 438308b8a8..50540c617a 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 @@ -3,7 +3,6 @@ use std::sync::Arc; use std::sync::Mutex; use std::sync::atomic::Ordering; use std::time::Duration; -use std::time::SystemTime; use anyhow::Result; use codex_core::config::Config; @@ -67,9 +66,7 @@ use core_test_support::wait_for_event; use pretty_assertions::assert_eq; use serde_json::json; -use super::GuardianV2Extension; -use super::GuardianV2ScoreProgress; - +use crate::async_scorer::authorization::ScoreAuthorization; use crate::async_scorer::config::CLASSIFICATION_OUTPUT_INSTRUCTIONS; use crate::async_scorer::config::DEFAULT_PARENT_COMPACTION_TOKENS; use crate::async_scorer::config::GuardianV2Config; @@ -86,6 +83,7 @@ use crate::async_scorer::sampler::LunaSampler; use crate::async_scorer::sampler::MODEL; use crate::async_scorer::sampler::tests::ProxyPrewarmLimit; use crate::async_scorer::sampler::tests::proxy_websocket_servers_with_http; +use crate::async_scorer::score::GuardianV2ScoreProgress; use crate::async_scorer::transcript::MAX_MESSAGE_ENTRY_TOKENS; use crate::async_scorer::transcript::MAX_TOOL_ENTRY_TOKENS; use crate::async_scorer::transcript::truncate_entry; @@ -449,41 +447,6 @@ impl ConversationHistorySnapshot for TestConversationHistory { } } -#[test] -fn fail_closed_score_preserves_classification_order() { - let thread_store = ExtensionData::new("thread-1"); - let newer_sampled_at = SystemTime::UNIX_EPOCH + Duration::from_secs(1); - let newest_sampled_at = newer_sampled_at + Duration::from_secs(1); - let newer_score = SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), - call_id: None, - action: None, - sampled_at: Some(newer_sampled_at.into()), - }; - thread_store.insert(newer_score.clone()); - - GuardianV2Extension::record_fail_closed_score(&thread_store, SystemTime::UNIX_EPOCH); - assert_eq!( - thread_store.get::().as_deref(), - Some(&newer_score) - ); - - GuardianV2Extension::record_fail_closed_score(&thread_store, newest_sampled_at); - let fail_closed_score = SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), 1.0)]), - call_id: None, - action: None, - sampled_at: Some(newest_sampled_at.into()), - }; - assert!(!thread_store.insert_if(newer_score.clone(), |previous| { - previous.is_none_or(|previous| previous.sampled_at < newer_score.sampled_at) - })); - assert_eq!( - thread_store.get::().as_deref(), - Some(&fail_closed_score) - ); -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn sandboxed_shell_classification_respects_review_scope() -> Result<()> { let sandboxed = ToolPayload::Function { @@ -2240,7 +2203,7 @@ async fn contributor_skips_required_models_in_standard_scope() -> Result<()> { .await; model_info.slug = "protected-model".to_owned(); thread_store.insert(model_info); - let authorization = super::ScoreAuthorization::current(&test.codex).await; + let authorization = ScoreAuthorization::current(&test.codex).await; let progress = thread_store .get::() .expect("Guardian v2 should track score progress per thread"); @@ -2435,7 +2398,7 @@ async fn assert_compaction_approval_policy(thread_context_enabled: bool) -> Resu fixture.test.codex.guardian_authorization_version().await, authorization ); - let score_authorization = super::ScoreAuthorization::current(&fixture.test.codex).await; + let score_authorization = ScoreAuthorization::current(&fixture.test.codex).await; *thread_store .get::() .expect("score progress") 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 821d49485d..d271523231 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/mod.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/mod.rs @@ -1,12 +1,15 @@ mod action; mod approval; mod authorization; +mod classification; mod config; mod coverage; mod extension; mod metrics; +mod observation; mod parent_compaction; mod sampler; +mod score; mod startup; mod transcript; mod truncation; diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/observation.rs b/codex-rs/ext/guardian-v2/src/async_scorer/observation.rs new file mode 100644 index 0000000000..a0c222af39 --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/async_scorer/observation.rs @@ -0,0 +1,359 @@ +//! Selects tool observations for classification and captures their evidence before spawning. +//! Keep snapshots on their existing side of the background task boundary. + +use std::sync::Arc; +use std::sync::atomic::Ordering; +use std::time::Instant; +use std::time::SystemTime; + +use codex_analytics::AnalyticsEventsClient; +use codex_core::context::GuardianContextMode; +use codex_core::context::GuardianReviewEvidence; +use codex_core::context::NodeReplReviewEvidence; +use codex_extension_api::ExtensionWarning; +use codex_extension_api::GuardianV2Enabled; +use codex_extension_api::ToolPayload; +use codex_extension_api::ToolStartInput; +use codex_protocol::ThreadId; +use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::mcp::is_node_repl_backed_server; +use codex_protocol::openai_models::GuardianScope; +use codex_protocol::openai_models::ModelInfo; +use codex_protocol::security_risk::SecurityRiskScore; + +use super::action::ActionRenderError; +use super::action::GuardianAction; +use super::authorization::ScoreAuthorization; +use super::classification::Classification; +use super::config::GuardianV2Config; +use super::coverage::UnscoredAction; +use super::extension::GuardianV2Extension; +use super::metrics::record_classification; +use super::parent_compaction::ParentCompactionError; +use super::parent_compaction::select_parent_compaction; +use super::sampler::LunaSampler; +use super::score::GuardianV2ScoreProgress; +use super::score::record_fail_closed_score; + +impl GuardianV2Extension { + pub(super) async fn score_tool(&self, input: ToolStartInput<'_>) { + // Polling a code cell does not introduce another action or age its score. + if input.tool_name.is_default_namespace() && input.tool_name.name == "wait" { + return; + } + let classification_started_at = Instant::now(); + let Some(sampler) = input.thread_store.get::() else { + return; + }; + let Some(guardian_config) = input.thread_store.get::() else { + return; + }; + let Some(score_progress) = input.thread_store.get::() else { + return; + }; + let parent_model = input.thread_store.get::(); + let policy = guardian_config.policy_for_model(parent_model.as_deref()); + if !policy.scoring_enabled() { + input.thread_store.remove::(); + } + let mcp_server = input + .mcp_tool + .map(|tool| tool.tool_info().server_name.as_str()); + let scope = mcp_server + .map(GuardianScope::for_mcp_server) + .or_else(|| GuardianScope::for_tool(input.tool_name)); + if !policy.scores_tool(input.tool_name, input.payload, scope) { + match policy.unscored_action { + UnscoredAction::Ignore => {} + UnscoredAction::AgeScore => { + let index = score_progress + .latest_tool_call + .fetch_add(/*val*/ 1, Ordering::Relaxed) + .saturating_add(/*rhs*/ 1); + score_progress.wrapper_lag.record(&input, index); + // Unscored permission widening must not reuse an earlier approval score. + if input.tool_name.is_default_namespace() + && input.tool_name.name == "exec_command" + && let ToolPayload::Function { arguments } = input.payload + && let Ok(arguments) = serde_json::from_str::(arguments) + && arguments + .get("sandbox_permissions") + .and_then(serde_json::Value::as_str) + == Some("with_additional_permissions") + { + score_progress + .latest_failed_tool_call + .fetch_max(index, Ordering::Release); + } + } + UnscoredAction::InvalidateScore => { + let index = score_progress + .latest_tool_call + .fetch_add(/*val*/ 1, Ordering::Relaxed) + .saturating_add(/*rhs*/ 1); + score_progress.wrapper_lag.record(&input, index); + score_progress + .latest_failed_tool_call + .fetch_max(index, Ordering::Release); + } + } + return; + } + if input.mcp_tool.is_some_and(|tool| { + let info = tool.tool_info(); + is_node_repl_backed_server(&info.server_name) && info.tool.name == "js" + }) { + score_progress + .js_executions + .fetch_add(/*val*/ 1, Ordering::Relaxed); + } + let metrics = score_progress.metrics.clone(); + let analytics = input.session_store.get::(); + let sampled_at = SystemTime::now(); + let tool_call_index = score_progress + .latest_tool_call + .fetch_add(/*val*/ 1, Ordering::Relaxed) + .saturating_add(/*rhs*/ 1); + score_progress.wrapper_lag.record(&input, tool_call_index); + 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(); + let root_turn_id = input.root_turn_id.map(str::to_owned); + let parent_response_id = input + .turn_store + .get::() + .map(|id| id.0.clone()); + let thread_context: Result<_, String> = async { + let parsed_thread_id = + ThreadId::from_string(&thread_id).map_err(|error| error.to_string())?; + let manager = self + .thread_manager + .upgrade() + .ok_or_else(|| "thread manager is unavailable".to_string())?; + let thread = manager + .get_thread(parsed_thread_id) + .await + .map_err(|error| error.to_string())?; + let config = thread.config().await; + Ok((manager, thread, config)) + } + .await; + let (manager, thread, config) = match thread_context { + Ok(context) => context, + Err(error) => { + score_progress + .latest_failed_tool_call + .fetch_max(tool_call_index, Ordering::Release); + record_classification( + metrics.as_deref(), + classification_started_at.elapsed(), + "failure", + Some("thread_context_error"), + ); + event_sink.emit_warning(ExtensionWarning { + thread_id, + turn_id: Some(turn_id), + message: format!("Guardian V2 risk scoring failed: {error}"), + }); + return; + } + }; + // Use the live reviewer, not the startup config or per-app reviewer overrides. + let snapshot = thread.config_snapshot().await; + if snapshot.full_access + || thread.approvals_reviewer_for_turn(input.turn_id).await == ApprovalsReviewer::User + { + // A skipped call invalidates older scores, including ones still in flight. + score_progress + .latest_failed_tool_call + .fetch_max(tool_call_index, Ordering::Release); + return; + } + // A required model keeps synchronous review outside its CUA allowance. + if !(scope == Some(GuardianScope::ComputerUse) && policy.initial_cua_call) + && parent_model.as_ref().is_some_and(|model| { + config + .config_layer_stack + .requirements() + .auto_review_required_for_model(&model.slug) + }) + { + input.thread_store.remove::(); + return; + } + input.thread_store.insert(GuardianV2Enabled); + let model_defaults = parent_model + .as_ref() + .and_then(|model| model.model_messages.as_ref()) + .and_then(|messages| messages.guardian_v2.as_ref()); + let guardian_config = match guardian_config.with_model_defaults(model_defaults) { + Ok(config) => config, + Err(error) => { + record_fail_closed_score(input.thread_store, sampled_at); + record_classification( + metrics.as_deref(), + classification_started_at.elapsed(), + "failure", + Some("configuration_error"), + ); + self.event_sink.emit_warning(ExtensionWarning { + thread_id: input.thread_store.level_id().to_owned(), + turn_id: Some(input.turn_id.to_owned()), + message: error, + }); + return; + } + }; + if guardian_config.transcript.include_images { + input + .thread_store + .get_or_init(NodeReplReviewEvidence::default) + .enable_image_capture(); + } + input.thread_store.insert(guardian_config.clone()); + let guardian_evidence = input + .thread_store + .get_or_init(GuardianReviewEvidence::default); + let context_mode = guardian_evidence.context_mode(); + let selected_compaction = match select_parent_compaction( + context_mode, + &guardian_config, + input.conversation_history.as_ref(), + &sampler, + parent_model + .as_ref() + .and_then(|model| model.comp_hash.as_deref()), + ) { + Ok(compaction) => compaction, + Err(error) => { + let (outcome, failure_reason) = if error == ParentCompactionError::RequiresSync { + score_progress + .latest_failed_tool_call + .fetch_max(tool_call_index, Ordering::Release); + ("skipped", None) + } else { + ("failure", Some("parent_compaction_error")) + }; + record_fail_closed_score(input.thread_store, sampled_at); + record_classification( + metrics.as_deref(), + classification_started_at.elapsed(), + outcome, + failure_reason, + ); + return; + } + }; + let parent_compaction = selected_compaction.item; + let parent_compaction_hash = selected_compaction.model_hash; + let call_id = input.call_id.to_owned(); + let mcp_tool = input.mcp_tool.cloned(); + let action = GuardianAction { + tool_name: input.tool_name.clone(), + payload: input.payload.clone(), + }; + let planned_action = match action.render(guardian_config.max_action_tokens) { + Ok(text) => text, + Err(ActionRenderError::TooLarge { .. }) => { + score_progress + .oversized_tool_calls + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(input.call_id.to_owned()); + score_progress + .latest_failed_tool_call + .fetch_max(tool_call_index, Ordering::Release); + record_fail_closed_score(input.thread_store, sampled_at); + record_classification( + metrics.as_deref(), + classification_started_at.elapsed(), + "failure", + Some("input_too_large"), + ); + return; + } + Err(error) => { + score_progress + .latest_failed_tool_call + .fetch_max(tool_call_index, Ordering::Release); + record_fail_closed_score(input.thread_store, sampled_at); + record_classification( + metrics.as_deref(), + classification_started_at.elapsed(), + "failure", + Some("action_serialization_error"), + ); + event_sink.emit_warning(ExtensionWarning { + thread_id, + turn_id: Some(turn_id), + message: format!("Guardian V2 action serialization failed: {error}"), + }); + return; + } + }; + let review_model_override = parent_model + .as_ref() + .and_then(|model| model.auto_review_model_override.clone()); + // Snapshot before spawning so a delayed sample cannot see later reviews. + let sync_reviews = guardian_evidence.snapshot(); + let codex_core::context::GuardianUserInputSnapshot { + fragments: trusted_user_inputs, + authorization_version, + } = guardian_evidence.user_input_snapshot(input.conversation_history.as_ref()); + let history = Arc::clone(&input.conversation_history); + let local_trusted_skill_paths = guardian_evidence.trusted_skill_paths(input.turn_id); + let node_repl_images = if guardian_config.transcript.include_images { + input + .thread_store + .get::() + .map(|evidence| evidence.images()) + .unwrap_or_default() + } else { + Vec::new() + }; + // Capture root evidence before background metadata resolution or model I/O. + // Later root changes invalidate this sample through its captured authorization version. + let root_snapshot = if context_mode == GuardianContextMode::ThreadOwned { + thread.guardian_root_snapshot().await + } else { + None + }; + + let score_authorization = ScoreAuthorization::current(&thread).await; + let classification = Classification { + classification_started_at, + sampler, + guardian_config, + score_progress, + parent_model, + metrics, + analytics, + sampled_at, + tool_call_index, + event_sink, + thread_id, + turn_id, + root_turn_id, + parent_response_id, + manager, + thread, + config, + context_mode, + parent_compaction, + parent_compaction_hash, + call_id, + mcp_tool, + planned_action, + review_model_override, + sync_reviews, + trusted_user_inputs, + authorization_version, + history, + local_trusted_skill_paths, + node_repl_images, + root_snapshot, + score_authorization, + }; + tokio::spawn(classification.run()); + } +} diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/score.rs b/codex-rs/ext/guardian-v2/src/async_scorer/score.rs new file mode 100644 index 0000000000..47729efd3c --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/async_scorer/score.rs @@ -0,0 +1,48 @@ +//! Tracks observation coverage and the authorization attached to the latest score. +//! Failed samples may replace an equally old score, but never a newer one. + +use std::collections::BTreeMap; +use std::collections::BTreeSet; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::time::SystemTime; + +use codex_extension_api::ExtensionData; +use codex_extension_api::ExtensionMetrics; +use codex_protocol::security_risk::SecurityRiskScore; + +use super::authorization::ScoreAuthorization; +use super::wrapper_lag::WrapperLag; + +#[derive(Default)] +pub(super) struct GuardianV2ScoreProgress { + pub(super) wrapper_lag: WrapperLag, + pub(super) latest_tool_call: AtomicUsize, + // Setup and reset calls must not consume the first JS execution allowance. + pub(super) js_executions: AtomicUsize, + pub(super) latest_scored_tool_call: AtomicUsize, + pub(super) latest_failed_tool_call: AtomicUsize, + // Keep overflow attached to each active call even after a newer score succeeds. + // The host's finish callback removes entries on completion, failure, or cancellation. + pub(super) oversized_tool_calls: Mutex>, + // Serialize successful score publication with its authorization metadata. + pub(super) authorization: Mutex>, + pub(super) metrics: Option>, +} + +pub(super) fn record_fail_closed_score(thread_store: &ExtensionData, sampled_at: SystemTime) { + let score = SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 1.0)]), + call_id: None, + action: None, + sampled_at: Some(sampled_at.into()), + }; + thread_store.insert_if(score.clone(), |previous| { + previous.is_none_or(|previous| previous.sampled_at <= score.sampled_at) + }); +} + +#[cfg(test)] +#[path = "score_tests.rs"] +mod tests; diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/score_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/score_tests.rs new file mode 100644 index 0000000000..8fcf3bff9f --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/async_scorer/score_tests.rs @@ -0,0 +1,44 @@ +use std::collections::BTreeMap; +use std::time::Duration; +use std::time::SystemTime; + +use codex_extension_api::ExtensionData; +use codex_protocol::security_risk::SecurityRiskScore; +use pretty_assertions::assert_eq; + +use super::record_fail_closed_score; + +#[test] +fn fail_closed_score_preserves_classification_order() { + let thread_store = ExtensionData::new("thread-1"); + let newer_sampled_at = SystemTime::UNIX_EPOCH + Duration::from_secs(1); + let newest_sampled_at = newer_sampled_at + Duration::from_secs(1); + let newer_score = SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), + call_id: None, + action: None, + sampled_at: Some(newer_sampled_at.into()), + }; + thread_store.insert(newer_score.clone()); + + record_fail_closed_score(&thread_store, SystemTime::UNIX_EPOCH); + assert_eq!( + thread_store.get::().as_deref(), + Some(&newer_score) + ); + + for sampled_at in [newer_sampled_at, newest_sampled_at] { + thread_store.insert(newer_score.clone()); + record_fail_closed_score(&thread_store, sampled_at); + let fail_closed_score = SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 1.0)]), + call_id: None, + action: None, + sampled_at: Some(sampled_at.into()), + }; + assert_eq!( + thread_store.get::().as_deref(), + Some(&fail_closed_score) + ); + } +}