diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 3ef79f7aa..ae93327b8 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -31,6 +31,7 @@ use async_trait::async_trait; use codex_apply_patch::ApplyPatchAction; use codex_apply_patch::ApplyPatchFileChange; use codex_utils_absolute_path::AbsolutePathBuf; +use std::sync::Arc; pub struct ApplyPatchHandler; @@ -139,8 +140,8 @@ impl ToolHandler for ApplyPatchHandler { let mut orchestrator = ToolOrchestrator::new(); let mut runtime = ApplyPatchRuntime::new(); let tool_ctx = ToolCtx { - session: session.as_ref(), - turn: turn.as_ref(), + session: session.clone(), + turn: turn.clone(), call_id: call_id.clone(), tool_name: tool_name.to_string(), }; @@ -149,7 +150,7 @@ impl ToolHandler for ApplyPatchHandler { &mut runtime, &req, &tool_ctx, - &turn, + turn.as_ref(), turn.approval_policy.value(), ) .await @@ -193,8 +194,8 @@ pub(crate) async fn intercept_apply_patch( command: &[String], cwd: &Path, timeout_ms: Option, - session: &Session, - turn: &TurnContext, + session: Arc, + turn: Arc, tracker: Option<&SharedTurnDiffTracker>, call_id: &str, tool_name: &str, @@ -203,11 +204,13 @@ pub(crate) async fn intercept_apply_patch( codex_apply_patch::MaybeApplyPatchVerified::Body(changes) => { session .record_model_warning( - format!("apply_patch was requested via {tool_name}. Use the apply_patch tool instead of exec_command."), - turn, + format!( + "apply_patch was requested via {tool_name}. Use the apply_patch tool instead of exec_command." + ), + turn.as_ref(), ) .await; - match apply_patch::apply_patch(turn, changes).await { + match apply_patch::apply_patch(turn.as_ref(), changes).await { InternalApplyPatchInvocation::Output(item) => { let content = item?; Ok(Some(ToolOutput::Function { @@ -219,8 +222,12 @@ pub(crate) async fn intercept_apply_patch( let changes = convert_apply_patch_to_protocol(&apply.action); let approval_keys = file_paths_for_action(&apply.action); let emitter = ToolEmitter::apply_patch(changes.clone(), apply.auto_approved); - let event_ctx = - ToolEventCtx::new(session, turn, call_id, tracker.as_ref().copied()); + let event_ctx = ToolEventCtx::new( + session.as_ref(), + turn.as_ref(), + call_id, + tracker.as_ref().copied(), + ); emitter.begin(event_ctx).await; let req = ApplyPatchRequest { @@ -235,8 +242,8 @@ pub(crate) async fn intercept_apply_patch( let mut orchestrator = ToolOrchestrator::new(); let mut runtime = ApplyPatchRuntime::new(); let tool_ctx = ToolCtx { - session, - turn, + session: session.clone(), + turn: turn.clone(), call_id: call_id.to_string(), tool_name: tool_name.to_string(), }; @@ -245,13 +252,17 @@ pub(crate) async fn intercept_apply_patch( &mut runtime, &req, &tool_ctx, - turn, + turn.as_ref(), turn.approval_policy.value(), ) .await .map(|result| result.output); - let event_ctx = - ToolEventCtx::new(session, turn, call_id, tracker.as_ref().copied()); + let event_ctx = ToolEventCtx::new( + session.as_ref(), + turn.as_ref(), + call_id, + tracker.as_ref().copied(), + ); let content = emitter.finish(event_ctx, out).await?; Ok(Some(ToolOutput::Function { body: FunctionCallOutputBody::Text(content), diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index 4cdbf454c..46a90681d 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -296,8 +296,8 @@ impl ShellHandler { &exec_params.command, &exec_params.cwd, exec_params.expiration.timeout_ms(), - session.as_ref(), - turn.as_ref(), + session.clone(), + turn.clone(), Some(&tracker), &call_id, tool_name.as_str(), @@ -343,8 +343,8 @@ impl ShellHandler { let mut orchestrator = ToolOrchestrator::new(); let mut runtime = ShellRuntime::new(); let tool_ctx = ToolCtx { - session: session.as_ref(), - turn: turn.as_ref(), + session: session.clone(), + turn: turn.clone(), call_id: call_id.clone(), tool_name, }; diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index c7d57731e..76dbade19 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -172,8 +172,8 @@ impl ToolHandler for UnifiedExecHandler { &command, &cwd, Some(yield_time_ms), - context.session.as_ref(), - context.turn.as_ref(), + context.session.clone(), + context.turn.clone(), Some(&tracker), &context.call_id, tool_name.as_str(), diff --git a/codex-rs/core/src/tools/orchestrator.rs b/codex-rs/core/src/tools/orchestrator.rs index afb24d4b3..db28c9313 100644 --- a/codex-rs/core/src/tools/orchestrator.rs +++ b/codex-rs/core/src/tools/orchestrator.rs @@ -48,7 +48,7 @@ impl ToolOrchestrator { async fn run_attempt( tool: &mut T, req: &Rq, - tool_ctx: &ToolCtx<'_>, + tool_ctx: &ToolCtx, attempt: &SandboxAttempt<'_>, has_managed_network_requirements: bool, ) -> (Result, Option) @@ -56,7 +56,7 @@ impl ToolOrchestrator { T: ToolRuntime, { let network_approval = begin_network_approval( - tool_ctx.session, + &tool_ctx.session, &tool_ctx.turn.sub_id, &tool_ctx.call_id, has_managed_network_requirements, @@ -65,8 +65,8 @@ impl ToolOrchestrator { .await; let attempt_tool_ctx = ToolCtx { - session: tool_ctx.session, - turn: tool_ctx.turn, + session: tool_ctx.session.clone(), + turn: tool_ctx.turn.clone(), call_id: tool_ctx.call_id.clone(), tool_name: tool_ctx.tool_name.clone(), }; @@ -79,7 +79,7 @@ impl ToolOrchestrator { match network_approval.mode() { NetworkApprovalMode::Immediate => { let finalize_result = - finish_immediate_network_approval(tool_ctx.session, network_approval).await; + finish_immediate_network_approval(&tool_ctx.session, network_approval).await; if let Err(err) = finalize_result { return (Err(err), None); } @@ -88,7 +88,7 @@ impl ToolOrchestrator { NetworkApprovalMode::Deferred => { let deferred = network_approval.into_deferred(); if run_result.is_err() { - finish_deferred_network_approval(tool_ctx.session, deferred).await; + finish_deferred_network_approval(&tool_ctx.session, deferred).await; return (run_result, None); } (run_result, deferred) @@ -100,7 +100,7 @@ impl ToolOrchestrator { &mut self, tool: &mut T, req: &Rq, - tool_ctx: &ToolCtx<'_>, + tool_ctx: &ToolCtx, turn_ctx: &crate::codex::TurnContext, approval_policy: AskForApproval, ) -> Result, ToolError> @@ -128,7 +128,7 @@ impl ToolOrchestrator { } ExecApprovalRequirement::NeedsApproval { reason, .. } => { let approval_ctx = ApprovalCtx { - session: tool_ctx.session, + session: &tool_ctx.session, turn: turn_ctx, call_id: &tool_ctx.call_id, retry_reason: reason, @@ -256,7 +256,7 @@ impl ToolOrchestrator { && network_approval_context.is_none(); if !bypass_retry_approval { let approval_ctx = ApprovalCtx { - session: tool_ctx.session, + session: &tool_ctx.session, turn: turn_ctx, call_id: &tool_ctx.call_id, retry_reason: Some(retry_reason), diff --git a/codex-rs/core/src/tools/runtimes/apply_patch.rs b/codex-rs/core/src/tools/runtimes/apply_patch.rs index d0ae2f61d..749e1abd2 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch.rs @@ -70,7 +70,7 @@ impl ApplyPatchRuntime { }) } - fn stdout_stream(ctx: &ToolCtx<'_>) -> Option { + fn stdout_stream(ctx: &ToolCtx) -> Option { Some(crate::exec::StdoutStream { sub_id: ctx.turn.sub_id.clone(), call_id: ctx.call_id.clone(), @@ -156,7 +156,7 @@ impl ToolRuntime for ApplyPatchRuntime { &mut self, req: &ApplyPatchRequest, attempt: &SandboxAttempt<'_>, - ctx: &ToolCtx<'_>, + ctx: &ToolCtx, ) -> Result { let spec = Self::build_command_spec(req)?; let env = attempt diff --git a/codex-rs/core/src/tools/runtimes/shell.rs b/codex-rs/core/src/tools/runtimes/shell.rs index bcbeafadf..11ad66916 100644 --- a/codex-rs/core/src/tools/runtimes/shell.rs +++ b/codex-rs/core/src/tools/runtimes/shell.rs @@ -60,7 +60,7 @@ impl ShellRuntime { Self } - fn stdout_stream(ctx: &ToolCtx<'_>) -> Option { + fn stdout_stream(ctx: &ToolCtx) -> Option { Some(crate::exec::StdoutStream { sub_id: ctx.turn.sub_id.clone(), call_id: ctx.call_id.clone(), @@ -150,7 +150,7 @@ impl ToolRuntime for ShellRuntime { fn network_approval_spec( &self, req: &ShellRequest, - _ctx: &ToolCtx<'_>, + _ctx: &ToolCtx, ) -> Option { req.network.as_ref()?; Some(NetworkApprovalSpec { @@ -163,7 +163,7 @@ impl ToolRuntime for ShellRuntime { &mut self, req: &ShellRequest, attempt: &SandboxAttempt<'_>, - ctx: &ToolCtx<'_>, + ctx: &ToolCtx, ) -> Result { let base_command = &req.command; let session_shell = ctx.session.user_shell(); @@ -207,7 +207,7 @@ impl ToolRuntime for ShellRuntime { .session .services .zsh_exec_bridge - .execute_shell_request(&env, ctx.session, ctx.turn, &ctx.call_id) + .execute_shell_request(&env, &ctx.session, &ctx.turn, &ctx.call_id) .await; } diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index 201de7ad8..5e06ca397 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -153,7 +153,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt fn network_approval_spec( &self, req: &UnifiedExecRequest, - _ctx: &ToolCtx<'_>, + _ctx: &ToolCtx, ) -> Option { req.network.as_ref()?; Some(NetworkApprovalSpec { @@ -166,7 +166,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt &mut self, req: &UnifiedExecRequest, attempt: &SandboxAttempt<'_>, - ctx: &ToolCtx<'_>, + ctx: &ToolCtx, ) -> Result { let base_command = &req.command; let session_shell = ctx.session.user_shell(); diff --git a/codex-rs/core/src/tools/sandboxing.rs b/codex-rs/core/src/tools/sandboxing.rs index 9721fed53..540eb1643 100644 --- a/codex-rs/core/src/tools/sandboxing.rs +++ b/codex-rs/core/src/tools/sandboxing.rs @@ -18,14 +18,14 @@ use codex_protocol::approvals::ExecPolicyAmendment; use codex_protocol::approvals::NetworkApprovalContext; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::ReviewDecision; +use futures::Future; +use futures::future::BoxFuture; +use serde::Serialize; use std::collections::HashMap; use std::fmt::Debug; use std::hash::Hash; use std::path::Path; - -use futures::Future; -use futures::future::BoxFuture; -use serde::Serialize; +use std::sync::Arc; #[derive(Clone, Default, Debug)] pub(crate) struct ApprovalStore { @@ -267,9 +267,9 @@ pub(crate) trait Sandboxable { } } -pub(crate) struct ToolCtx<'a> { - pub session: &'a Session, - pub turn: &'a TurnContext, +pub(crate) struct ToolCtx { + pub session: Arc, + pub turn: Arc, pub call_id: String, pub tool_name: String, } @@ -281,7 +281,7 @@ pub(crate) enum ToolError { } pub(crate) trait ToolRuntime: Approvable + Sandboxable { - fn network_approval_spec(&self, _req: &Req, _ctx: &ToolCtx<'_>) -> Option { + fn network_approval_spec(&self, _req: &Req, _ctx: &ToolCtx) -> Option { None } diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 68c4f95ce..a3da4ba6b 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -594,8 +594,8 @@ impl UnifiedExecProcessManager { exec_approval_requirement, }; let tool_ctx = ToolCtx { - session: context.session.as_ref(), - turn: context.turn.as_ref(), + session: context.session.clone(), + turn: context.turn.clone(), call_id: context.call_id.clone(), tool_name: "exec_command".to_string(), }; @@ -604,7 +604,7 @@ impl UnifiedExecProcessManager { &mut runtime, &req, &tool_ctx, - context.turn.as_ref(), + &context.turn, context.turn.approval_policy.value(), ) .await