From bd7e3ef0f38f724a1ae21f480d29a91e4d0cf0bd Mon Sep 17 00:00:00 2001 From: Roy Han Date: Thu, 19 Mar 2026 09:09:46 -0700 Subject: [PATCH] scale schema --- codex-rs/core/src/codex.rs | 69 ++++++++++++++++++-------------- codex-rs/core/src/codex_tests.rs | 67 +++++++++++++++++++++++++------ codex-rs/core/src/state/mod.rs | 2 + codex-rs/core/src/state/turn.rs | 68 +++++++++++++++++++++++++------ 4 files changed, 151 insertions(+), 55 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 58e99d42de..1791e2a42a 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -283,6 +283,8 @@ use crate::skills::injection::app_id_from_path; use crate::skills::injection::tool_kind_for_path; use crate::skills::resolve_skill_dependencies_for_turn; use crate::state::ActiveTurn; +use crate::state::ApprovalOutcomeMetadata; +use crate::state::PendingApprovalMetadata; use crate::state::SessionServices; use crate::state::SessionState; use crate::state_db; @@ -325,7 +327,6 @@ use codex_protocol::models::DeveloperInstructions; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; use codex_protocol::models::ResponseItemMetadata; -use codex_protocol::models::ReviewDecisionMetadata; use codex_protocol::models::SandboxPolicyMetadata; use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; use codex_protocol::protocol::CodexErrorInfo; @@ -1061,28 +1062,6 @@ fn stamp_message_metadata_on_response_item( } } -fn review_decision_to_metadata(decision: &ReviewDecision) -> ReviewDecisionMetadata { - match decision { - ReviewDecision::Approved => ReviewDecisionMetadata::Approved, - ReviewDecision::ApprovedExecpolicyAmendment { .. } => { - ReviewDecisionMetadata::ApprovedWithAmendment - } - ReviewDecision::ApprovedForSession => ReviewDecisionMetadata::ApprovedForSession, - ReviewDecision::NetworkPolicyAmendment { - network_policy_amendment, - } => match network_policy_amendment.action { - codex_protocol::protocol::NetworkPolicyRuleAction::Allow => { - ReviewDecisionMetadata::ApprovedWithNetworkPolicyAllow - } - codex_protocol::protocol::NetworkPolicyRuleAction::Deny => { - ReviewDecisionMetadata::DeniedWithNetworkPolicyDeny - } - }, - ReviewDecision::Denied => ReviewDecisionMetadata::Denied, - ReviewDecision::Abort => ReviewDecisionMetadata::Abort, - } -} - fn response_item_tool_call_id(item: &ResponseItem) -> Option<&str> { match item { ResponseItem::LocalShellCall { @@ -1135,7 +1114,7 @@ fn tool_call_metadata_or_default(item: &ResponseItem) -> Option, + approval_outcomes_by_call_id: HashMap, pending_approval_call_ids: HashSet, } @@ -1163,9 +1142,9 @@ fn stamp_tool_approval_metadata_with_snapshot( )); match outcome { - Some(review_decision) => { + Some(outcome) => { metadata.is_tool_call_escalated = Some(true); - metadata.review_decision = Some(review_decision); + metadata.review_decision = outcome.review_decision; } None if !has_pending_approval => { metadata.is_tool_call_escalated = Some(false); @@ -3067,7 +3046,9 @@ impl Session { let mut ts = at.turn_state.lock().await; ts.insert_pending_approval_call_id( effective_approval_id.clone(), - call_id.clone(), + PendingApprovalMetadata { + call_id: call_id.clone(), + }, ); ts.insert_pending_approval(effective_approval_id.clone(), tx_approve) } @@ -3134,7 +3115,12 @@ impl Session { match active.as_mut() { Some(at) => { let mut ts = at.turn_state.lock().await; - ts.insert_pending_approval_call_id(approval_id.clone(), call_id.clone()); + ts.insert_pending_approval_call_id( + approval_id.clone(), + PendingApprovalMetadata { + call_id: call_id.clone(), + }, + ); ts.insert_pending_approval(approval_id.clone(), tx_approve) } None => None, @@ -3426,13 +3412,34 @@ impl Session { return; }; let mut ts = at.turn_state.lock().await; - let call_id = ts + let pending_approval = ts .remove_pending_approval_call_id(approval_id) - .unwrap_or_else(|| approval_id.to_string()); + .unwrap_or_else(|| PendingApprovalMetadata { + call_id: approval_id.to_string(), + }); + drop(ts); + drop(active); + self.record_call_approval_outcome( + pending_approval.call_id, + ApprovalOutcomeMetadata::reviewed(decision), + ) + .await; + } + + pub(crate) async fn record_call_approval_outcome( + &self, + call_id: String, + outcome: ApprovalOutcomeMetadata, + ) { + let mut active = self.active_turn.lock().await; + let Some(at) = active.as_mut() else { + return; + }; if !self.enabled(Feature::ItemMetadata) { return; } - ts.record_approval_outcome(call_id, review_decision_to_metadata(decision)); + let mut ts = at.turn_state.lock().await; + ts.record_approval_outcome(call_id, outcome); } pub async fn notify_approval(&self, approval_id: &str, decision: ReviewDecision) { diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 75ab04cad8..c3c3167d16 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -43,6 +43,7 @@ use crate::protocol::UserMessageEvent; use crate::rollout::policy::EventPersistenceMode; use crate::rollout::recorder::RolloutRecorder; use crate::rollout::recorder::RolloutRecorderParams; +use crate::state::ApprovalOutcomeMetadata; use crate::state::TaskKind; use crate::tasks::SessionTask; use crate::tasks::SessionTaskContext; @@ -4285,20 +4286,56 @@ async fn task_finish_emits_prompt_queued_metadata_for_injected_user_input_when_f #[test] fn review_decision_metadata_mapping_is_stable() { assert_eq!( - review_decision_to_metadata(&ReviewDecision::Approved), - codex_protocol::models::ReviewDecisionMetadata::Approved + ApprovalOutcomeMetadata::reviewed(&ReviewDecision::Approved).review_decision, + Some(codex_protocol::models::ReviewDecisionMetadata::Approved) ); assert_eq!( - review_decision_to_metadata(&ReviewDecision::Denied), - codex_protocol::models::ReviewDecisionMetadata::Denied + ApprovalOutcomeMetadata::reviewed(&ReviewDecision::Denied).review_decision, + Some(codex_protocol::models::ReviewDecisionMetadata::Denied) ); assert_eq!( - review_decision_to_metadata(&ReviewDecision::Abort), - codex_protocol::models::ReviewDecisionMetadata::Abort + ApprovalOutcomeMetadata::reviewed(&ReviewDecision::Abort).review_decision, + Some(codex_protocol::models::ReviewDecisionMetadata::Abort) ); assert_eq!( - review_decision_to_metadata(&ReviewDecision::ApprovedForSession), - codex_protocol::models::ReviewDecisionMetadata::ApprovedForSession + ApprovalOutcomeMetadata::reviewed(&ReviewDecision::ApprovedForSession).review_decision, + Some(codex_protocol::models::ReviewDecisionMetadata::ApprovedForSession) + ); + assert_eq!( + ApprovalOutcomeMetadata::reviewed(&ReviewDecision::ApprovedExecpolicyAmendment { + proposed_execpolicy_amendment: codex_protocol::approvals::ExecPolicyAmendment::new( + vec!["echo".to_string(), "hi".to_string()], + codex_protocol::protocol::BehaviorRule { + action: codex_execpolicy::ExecAction::Allow, + tag: "test".to_string(), + users: Vec::new(), + always_apply: false, + reason: None, + }, + ), + }) + .review_decision, + Some(codex_protocol::models::ReviewDecisionMetadata::ApprovedWithAmendment) + ); + assert_eq!( + ApprovalOutcomeMetadata::reviewed(&ReviewDecision::NetworkPolicyAmendment { + network_policy_amendment: codex_protocol::approvals::NetworkPolicyAmendment { + host: "example.com".to_string(), + action: codex_protocol::protocol::NetworkPolicyRuleAction::Allow, + }, + }) + .review_decision, + Some(codex_protocol::models::ReviewDecisionMetadata::ApprovedWithNetworkPolicyAllow) + ); + assert_eq!( + ApprovalOutcomeMetadata::reviewed(&ReviewDecision::NetworkPolicyAmendment { + network_policy_amendment: codex_protocol::approvals::NetworkPolicyAmendment { + host: "example.com".to_string(), + action: codex_protocol::protocol::NetworkPolicyRuleAction::Deny, + }, + }) + .review_decision, + Some(codex_protocol::models::ReviewDecisionMetadata::DeniedWithNetworkPolicyDeny) ); } @@ -4342,8 +4379,11 @@ async fn tool_call_metadata_stamps_escalated_review_decision_when_feature_enable .await; while rx.try_recv().is_ok() {} - sess.record_approval_outcome("call-1", &ReviewDecision::Denied) - .await; + sess.record_call_approval_outcome( + "call-1".to_string(), + ApprovalOutcomeMetadata::reviewed(&ReviewDecision::Denied), + ) + .await; sess.record_response_item_and_emit_turn_item( tc.as_ref(), ResponseItem::FunctionCall { @@ -4526,8 +4566,11 @@ async fn tool_call_metadata_can_be_restamped_after_approval_outcome() { }, ) .await; - sess.record_approval_outcome("call-restamp-1", &ReviewDecision::Denied) - .await; + sess.record_call_approval_outcome( + "call-restamp-1".to_string(), + ApprovalOutcomeMetadata::reviewed(&ReviewDecision::Denied), + ) + .await; let original_item = sess .clone_history() diff --git a/codex-rs/core/src/state/mod.rs b/codex-rs/core/src/state/mod.rs index 642433a786..f9448cfee3 100644 --- a/codex-rs/core/src/state/mod.rs +++ b/codex-rs/core/src/state/mod.rs @@ -5,5 +5,7 @@ mod turn; pub(crate) use service::SessionServices; pub(crate) use session::SessionState; pub(crate) use turn::ActiveTurn; +pub(crate) use turn::ApprovalOutcomeMetadata; +pub(crate) use turn::PendingApprovalMetadata; pub(crate) use turn::RunningTask; pub(crate) use turn::TaskKind; diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index 82a85c55ae..5cc944cd5d 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -65,6 +65,43 @@ pub(crate) struct PendingInputItem { pub(crate) metadata: Option, } +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct PendingApprovalMetadata { + pub(crate) call_id: String, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct ApprovalOutcomeMetadata { + pub(crate) review_decision: Option, +} + +impl ApprovalOutcomeMetadata { + pub(crate) fn reviewed(decision: &ReviewDecision) -> Self { + let review_decision = match decision { + ReviewDecision::Approved => ReviewDecisionMetadata::Approved, + ReviewDecision::ApprovedExecpolicyAmendment { .. } => { + ReviewDecisionMetadata::ApprovedWithAmendment + } + ReviewDecision::Denied => ReviewDecisionMetadata::Denied, + ReviewDecision::Abort => ReviewDecisionMetadata::Abort, + ReviewDecision::ApprovedForSession => ReviewDecisionMetadata::ApprovedForSession, + ReviewDecision::NetworkPolicyAmendment { + network_policy_amendment, + } => match network_policy_amendment.action { + codex_protocol::protocol::NetworkPolicyRuleAction::Allow => { + ReviewDecisionMetadata::ApprovedWithNetworkPolicyAllow + } + codex_protocol::protocol::NetworkPolicyRuleAction::Deny => { + ReviewDecisionMetadata::DeniedWithNetworkPolicyDeny + } + }, + }; + Self { + review_decision: Some(review_decision), + } + } +} + impl ActiveTurn { pub(crate) fn add_task(&mut self, task: RunningTask) { let sub_id = task.turn_context.sub_id.clone(); @@ -85,8 +122,8 @@ impl ActiveTurn { #[derive(Default)] pub(crate) struct TurnState { pending_approvals: HashMap>, - pending_approval_call_ids: HashMap, - approval_outcomes_by_call_id: HashMap, + pending_approval_metadata_by_id: HashMap, + approval_outcomes_by_call_id: HashMap, pending_request_permissions: HashMap>, pending_user_input: HashMap>, pending_elicitations: HashMap<(String, RequestId), oneshot::Sender>, @@ -109,29 +146,36 @@ impl TurnState { pub(crate) fn insert_pending_approval_call_id( &mut self, approval_key: String, - call_id: String, - ) -> Option { - self.pending_approval_call_ids.insert(approval_key, call_id) + pending_metadata: PendingApprovalMetadata, + ) -> Option { + self.pending_approval_metadata_by_id + .insert(approval_key, pending_metadata) } - pub(crate) fn remove_pending_approval_call_id(&mut self, approval_key: &str) -> Option { - self.pending_approval_call_ids.remove(approval_key) + pub(crate) fn remove_pending_approval_call_id( + &mut self, + approval_key: &str, + ) -> Option { + self.pending_approval_metadata_by_id.remove(approval_key) } pub(crate) fn record_approval_outcome( &mut self, call_id: String, - decision: ReviewDecisionMetadata, + outcome: ApprovalOutcomeMetadata, ) { - self.approval_outcomes_by_call_id.insert(call_id, decision); + self.approval_outcomes_by_call_id.insert(call_id, outcome); } pub(crate) fn approval_metadata_snapshot( &self, - ) -> (HashMap, HashSet) { + ) -> (HashMap, HashSet) { ( self.approval_outcomes_by_call_id.clone(), - self.pending_approval_call_ids.values().cloned().collect(), + self.pending_approval_metadata_by_id + .values() + .map(|metadata| metadata.call_id.clone()) + .collect(), ) } @@ -144,7 +188,7 @@ impl TurnState { pub(crate) fn clear_pending(&mut self) { self.pending_approvals.clear(); - self.pending_approval_call_ids.clear(); + self.pending_approval_metadata_by_id.clear(); self.approval_outcomes_by_call_id.clear(); self.pending_request_permissions.clear(); self.pending_user_input.clear();