diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index f219096e35..e562edd0f1 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2374,6 +2374,7 @@ dependencies = [ "codex-feedback", "codex-git-utils", "codex-hooks", + "codex-kernel", "codex-login", "codex-mcp", "codex-memories-read", @@ -2815,6 +2816,20 @@ dependencies = [ "tempfile", ] +[[package]] +name = "codex-kernel" +version = "0.0.0" +dependencies = [ + "codex-protocol", + "codex-tools", + "pretty_assertions", + "rand 0.9.3", + "serde", + "serde_json", + "tokio", + "tokio-util", +] + [[package]] name = "codex-keyring-store" version = "0.0.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index a9d6ea6c23..a5aaee5905 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -43,6 +43,7 @@ members = [ "external-agent-migration", "external-agent-sessions", "keyring-store", + "kernel", "file-search", "linux-sandbox", "lmstudio", @@ -158,6 +159,7 @@ codex-file-search = { path = "file-search" } codex-git-utils = { path = "git-utils" } codex-hooks = { path = "hooks" } codex-keyring-store = { path = "keyring-store" } +codex-kernel = { path = "kernel" } codex-linux-sandbox = { path = "linux-sandbox" } codex-lmstudio = { path = "lmstudio" } codex-login = { path = "login" } diff --git a/codex-rs/app-server/tests/suite/v2/turn_start.rs b/codex-rs/app-server/tests/suite/v2/turn_start.rs index 3ff04d5022..61b37954bd 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start.rs @@ -331,10 +331,20 @@ async fn turn_start_emits_thread_scoped_warning_notification_for_trimmed_skills( let warning: WarningNotification = serde_json::from_value(params).expect("deserialize warning notification"); assert_eq!(warning.thread_id.as_deref(), Some(thread.id.as_str())); - assert_eq!( - warning.message, - "Exceeded skills context budget of 2%. All skill descriptions were removed and 7 additional skills were not included in the model-visible skills list." - ); + let prefix = "Exceeded skills context budget of 2%. All skill descriptions were removed and "; + let plural_suffix = " additional skills were not included in the model-visible skills list."; + let singular_suffix = " additional skill was not included in the model-visible skills list."; + let omitted_count = warning + .message + .strip_prefix(prefix) + .and_then(|message| { + message + .strip_suffix(plural_suffix) + .or_else(|| message.strip_suffix(singular_suffix)) + }) + .and_then(|count| count.parse::().ok()) + .expect("warning should report an omitted skill count"); + assert!(omitted_count >= 1, "expected at least one omitted skill"); timeout( DEFAULT_READ_TIMEOUT, diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 44c6aacac5..f08a5d4e1f 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -47,6 +47,7 @@ codex-shell-command = { workspace = true } codex-execpolicy = { workspace = true } codex-git-utils = { workspace = true } codex-hooks = { workspace = true } +codex-kernel = { workspace = true } codex-network-proxy = { workspace = true } codex-otel = { workspace = true } codex-plugin = { workspace = true } diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index efe2670652..8eeba26f39 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -1,14 +1,7 @@ pub use codex_api::ResponseEvent; -use codex_config::types::Personality; +pub use codex_kernel::Prompt; use codex_protocol::error::Result; -use codex_protocol::models::BaseInstructions; -use codex_protocol::models::FunctionCallOutputBody; -use codex_protocol::models::ResponseItem; -use codex_tools::ToolSpec; use futures::Stream; -use serde::Deserialize; -use serde_json::Value; -use std::collections::HashSet; use std::pin::Pin; use std::task::Context; use std::task::Poll; @@ -23,157 +16,6 @@ pub const REVIEW_EXIT_SUCCESS_TMPL: &str = include_str!("../templates/review/exi pub const REVIEW_EXIT_INTERRUPTED_TMPL: &str = include_str!("../templates/review/exit_interrupted.xml"); -/// API request payload for a single model turn -#[derive(Debug, Clone)] -pub struct Prompt { - /// Conversation context input items. - pub input: Vec, - - /// Tools available to the model, including additional tools sourced from - /// external MCP servers. - pub(crate) tools: Vec, - - /// Whether parallel tool calls are permitted for this prompt. - pub(crate) parallel_tool_calls: bool, - - pub base_instructions: BaseInstructions, - - /// Optionally specify the personality of the model. - pub personality: Option, - - /// Optional the output schema for the model's response. - pub output_schema: Option, - - /// Whether the Responses API should strictly validate `output_schema`. - pub output_schema_strict: bool, -} - -impl Default for Prompt { - fn default() -> Self { - Self { - input: Vec::new(), - tools: Vec::new(), - parallel_tool_calls: false, - base_instructions: BaseInstructions::default(), - personality: None, - output_schema: None, - output_schema_strict: true, - } - } -} - -impl Prompt { - pub(crate) fn get_formatted_input(&self) -> Vec { - let mut input = self.input.clone(); - - // when using the *Freeform* apply_patch tool specifically, tool outputs - // should be structured text, not json. Do NOT reserialize when using - // the Function tool - note that this differs from the check above for - // instructions. We declare the result as a named variable for clarity. - let is_freeform_apply_patch_tool_present = self.tools.iter().any(|tool| match tool { - ToolSpec::Freeform(f) => f.name == "apply_patch", - _ => false, - }); - if is_freeform_apply_patch_tool_present { - reserialize_shell_outputs(&mut input); - } - - input - } -} - -fn reserialize_shell_outputs(items: &mut [ResponseItem]) { - let mut shell_call_ids: HashSet = HashSet::new(); - - items.iter_mut().for_each(|item| match item { - ResponseItem::LocalShellCall { call_id, id, .. } => { - if let Some(identifier) = call_id.clone().or_else(|| id.clone()) { - shell_call_ids.insert(identifier); - } - } - ResponseItem::CustomToolCall { - id: _, - status: _, - call_id, - name, - input: _, - } => { - if name == "apply_patch" { - shell_call_ids.insert(call_id.clone()); - } - } - ResponseItem::FunctionCall { name, call_id, .. } - if is_shell_tool_name(name) || name == "apply_patch" => - { - shell_call_ids.insert(call_id.clone()); - } - ResponseItem::FunctionCallOutput { - call_id, output, .. - } - | ResponseItem::CustomToolCallOutput { - call_id, output, .. - } => { - if shell_call_ids.remove(call_id) - && let Some(structured) = output - .text_content() - .and_then(parse_structured_shell_output) - { - output.body = FunctionCallOutputBody::Text(structured); - } - } - _ => {} - }) -} - -fn is_shell_tool_name(name: &str) -> bool { - matches!(name, "shell" | "container.exec") -} - -#[derive(Deserialize)] -struct ExecOutputJson { - output: String, - metadata: ExecOutputMetadataJson, -} - -#[derive(Deserialize)] -struct ExecOutputMetadataJson { - exit_code: i32, - duration_seconds: f32, -} - -fn parse_structured_shell_output(raw: &str) -> Option { - let parsed: ExecOutputJson = serde_json::from_str(raw).ok()?; - Some(build_structured_output(&parsed)) -} - -fn build_structured_output(parsed: &ExecOutputJson) -> String { - let mut sections = Vec::new(); - sections.push(format!("Exit code: {}", parsed.metadata.exit_code)); - sections.push(format!( - "Wall time: {} seconds", - parsed.metadata.duration_seconds - )); - - let mut output = parsed.output.clone(); - if let Some((stripped, total_lines)) = strip_total_output_header(&parsed.output) { - sections.push(format!("Total output lines: {total_lines}")); - output = stripped.to_string(); - } - - sections.push("Output:".to_string()); - sections.push(output); - - sections.join("\n") -} - -fn strip_total_output_header(output: &str) -> Option<(&str, u32)> { - let after_prefix = output.strip_prefix("Total output lines: ")?; - let (total_segment, remainder) = after_prefix.split_once('\n')?; - let total_lines = total_segment.parse::().ok()?; - let remainder = remainder.strip_prefix('\n').unwrap_or(remainder); - Some((remainder, total_lines)) -} - pub struct ResponseStream { pub(crate) rx_event: mpsc::Receiver>, /// Signals the mapper task that the consumer stopped polling before the diff --git a/codex-rs/core/src/client_common_tests.rs b/codex-rs/core/src/client_common_tests.rs index f67e4f1fd8..e615318008 100644 --- a/codex-rs/core/src/client_common_tests.rs +++ b/codex-rs/core/src/client_common_tests.rs @@ -3,11 +3,9 @@ use codex_api::ResponsesApiRequest; use codex_api::TextControls; use codex_api::create_text_param_for_request; use codex_protocol::config_types::ServiceTier; -use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::ResponseItem; use pretty_assertions::assert_eq; -use super::*; - #[test] fn serializes_text_verbosity_when_set() { let input: Vec = vec![]; @@ -166,65 +164,3 @@ fn serializes_flex_service_tier_when_set() { Some("flex") ); } - -#[test] -fn reserializes_shell_outputs_for_function_and_custom_tool_calls() { - let raw_output = r#"{"output":"hello","metadata":{"exit_code":0,"duration_seconds":0.5}}"#; - let expected_output = "Exit code: 0\nWall time: 0.5 seconds\nOutput:\nhello"; - let mut items = vec![ - ResponseItem::FunctionCall { - id: None, - name: "shell".to_string(), - namespace: None, - arguments: "{}".to_string(), - call_id: "call-1".to_string(), - }, - ResponseItem::FunctionCallOutput { - call_id: "call-1".to_string(), - output: FunctionCallOutputPayload::from_text(raw_output.to_string()), - }, - ResponseItem::CustomToolCall { - id: None, - status: None, - call_id: "call-2".to_string(), - name: "apply_patch".to_string(), - input: "*** Begin Patch".to_string(), - }, - ResponseItem::CustomToolCallOutput { - call_id: "call-2".to_string(), - name: None, - output: FunctionCallOutputPayload::from_text(raw_output.to_string()), - }, - ]; - - reserialize_shell_outputs(&mut items); - - assert_eq!( - items, - vec![ - ResponseItem::FunctionCall { - id: None, - name: "shell".to_string(), - namespace: None, - arguments: "{}".to_string(), - call_id: "call-1".to_string(), - }, - ResponseItem::FunctionCallOutput { - call_id: "call-1".to_string(), - output: FunctionCallOutputPayload::from_text(expected_output.to_string()), - }, - ResponseItem::CustomToolCall { - id: None, - status: None, - call_id: "call-2".to_string(), - name: "apply_patch".to_string(), - input: "*** Begin Patch".to_string(), - }, - ResponseItem::CustomToolCallOutput { - call_id: "call-2".to_string(), - name: None, - output: FunctionCallOutputPayload::from_text(expected_output.to_string()), - }, - ] - ); -} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 415b89c0de..57dfbea868 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -188,6 +188,7 @@ mod mcp; mod multi_agents; mod review; mod rollout_reconstruction; +mod sampling_loop; #[allow(clippy::module_inception)] pub(crate) mod session; pub(crate) mod turn; diff --git a/codex-rs/core/src/session/sampling_loop.rs b/codex-rs/core/src/session/sampling_loop.rs new file mode 100644 index 0000000000..457ddb013b --- /dev/null +++ b/codex-rs/core/src/session/sampling_loop.rs @@ -0,0 +1,221 @@ +use std::collections::HashSet; +use std::sync::Arc; +use std::time::Duration; + +use crate::SkillLoadOutcome; +use crate::client::ModelClientSession; +use crate::session::session::Session; +use crate::session::turn::build_prompt_config; +use crate::session::turn::build_tool_config; +use crate::session::turn::built_tools; +use crate::session::turn::try_run_sampling_request; +use crate::session::turn_context::TurnContext; +use crate::tools::context::SharedTurnDiffTracker; +use crate::tools::parallel::ToolCallRuntime; +use codex_kernel::PreparedSamplingRequest; +use codex_kernel::Prompt; +use codex_kernel::SamplingLoopHost; +use codex_kernel::SamplingRequestResult; +use codex_kernel::run_sampling_request_loop; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::CodexErrorInfo; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::RateLimitSnapshot; +use codex_protocol::protocol::StreamErrorEvent; +use codex_protocol::protocol::WarningEvent; +use tokio_util::sync::CancellationToken; +use tracing::instrument; +use tracing::warn; + +struct SamplingRequestRuntime { + _code_mode_worker: Option, +} + +#[allow(clippy::too_many_arguments)] +#[instrument( + level = "trace", + skip_all, + fields( + turn_id = %turn_context.sub_id, + model = %turn_context.model_info.slug, + cwd = %turn_context.cwd.display() + ) +)] +pub(super) async fn run_sampling_request( + sess: Arc, + turn_context: Arc, + turn_diff_tracker: SharedTurnDiffTracker, + client_session: &mut ModelClientSession, + turn_metadata_header: Option<&str>, + input: Vec, + explicitly_enabled_connectors: &HashSet, + skills_outcome: Option<&SkillLoadOutcome>, + cancellation_token: CancellationToken, +) -> CodexResult { + let host = CoreSamplingLoopHost { + sess, + turn_context, + turn_diff_tracker, + turn_metadata_header, + explicitly_enabled_connectors, + skills_outcome, + }; + + run_sampling_request_loop(&host, client_session, input, cancellation_token).await +} + +struct CoreSamplingLoopHost<'a> { + sess: Arc, + turn_context: Arc, + turn_diff_tracker: SharedTurnDiffTracker, + turn_metadata_header: Option<&'a str>, + explicitly_enabled_connectors: &'a HashSet, + skills_outcome: Option<&'a SkillLoadOutcome>, +} + +impl SamplingLoopHost for CoreSamplingLoopHost<'_> { + type ClientSession = ModelClientSession; + type Runtime = SamplingRequestRuntime; + type Tools = ToolCallRuntime; + + async fn prepare_sampling_request( + &self, + input: &[ResponseItem], + cancellation_token: &CancellationToken, + ) -> CodexResult> { + let router = built_tools( + self.sess.as_ref(), + self.turn_context.as_ref(), + input, + self.explicitly_enabled_connectors, + self.skills_outcome, + cancellation_token, + ) + .await?; + let base_instructions = self.sess.get_base_instructions().await; + let tool_runtime = ToolCallRuntime::new( + Arc::clone(&router), + Arc::clone(&self.sess), + Arc::clone(&self.turn_context), + Arc::clone(&self.turn_diff_tracker), + ); + let code_mode_worker = self + .sess + .services + .code_mode_service + .start_turn_worker( + &self.sess, + &self.turn_context, + Arc::clone(&router), + Arc::clone(&self.turn_diff_tracker), + ) + .await; + + Ok(PreparedSamplingRequest { + prompt_config: build_prompt_config(self.turn_context.as_ref(), base_instructions), + tool_config: build_tool_config(router.as_ref(), self.turn_context.as_ref()), + runtime: SamplingRequestRuntime { + _code_mode_worker: code_mode_worker, + }, + tools: tool_runtime, + }) + } + + async fn history_prompt_input(&self) -> Vec { + self.sess + .clone_history() + .await + .for_prompt(&self.turn_context.model_info.input_modalities) + } + + async fn run_single_sampling_request( + &self, + _runtime: &Self::Runtime, + tools: &Self::Tools, + client_session: &mut Self::ClientSession, + prompt: &Prompt, + cancellation_token: CancellationToken, + ) -> CodexResult { + try_run_sampling_request( + tools.clone(), + Arc::clone(&self.sess), + Arc::clone(&self.turn_context), + client_session, + self.turn_metadata_header, + Arc::clone(&self.turn_diff_tracker), + prompt, + cancellation_token, + ) + .await + } + + fn stream_max_retries(&self) -> u64 { + self.turn_context.provider.info().stream_max_retries() + } + + fn try_switch_fallback_transport(&self, client_session: &mut Self::ClientSession) -> bool { + client_session.try_switch_fallback_transport( + &self.turn_context.session_telemetry, + &self.turn_context.model_info, + ) + } + + fn should_notify_stream_retry(&self, retries: u64, _err: &CodexErr) -> bool { + retries > 1 + || cfg!(debug_assertions) + || !self + .sess + .services + .model_client + .responses_websocket_enabled() + } + + async fn handle_context_window_exceeded(&self) { + self.sess.set_total_tokens_full(&self.turn_context).await; + } + + async fn handle_usage_limit_reached(&self, rate_limits: Option) { + if let Some(rate_limits) = rate_limits { + self.sess + .update_rate_limits(&self.turn_context, rate_limits) + .await; + } + } + + async fn notify_fallback_to_http(&self, err: &CodexErr) { + self.sess + .send_event( + &self.turn_context, + EventMsg::Warning(WarningEvent { + message: format!("Falling back from WebSockets to HTTPS transport. {err:#}"), + }), + ) + .await; + } + + async fn notify_stream_retry( + &self, + retries: u64, + max_retries: u64, + delay: Duration, + err: &CodexErr, + ) { + warn!( + "stream disconnected - retrying sampling request ({retries}/{max_retries} in {delay:?})..." + ); + self.sess + .send_event( + &self.turn_context, + EventMsg::StreamError(StreamErrorEvent { + message: format!("Reconnecting... {retries}/{max_retries}"), + codex_error_info: Some(CodexErrorInfo::ResponseStreamDisconnected { + http_status_code: err.http_status_code_value(), + }), + additional_details: Some(err.to_string()), + }), + ) + .await; + } +} diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 823dcbfc3d..b7adabbe9c 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -40,6 +40,7 @@ use crate::parse_turn_item; use crate::plugins::build_plugin_injections; use crate::resolve_skill_dependencies_for_turn; use crate::session::PreviousTurnSettings; +use crate::session::sampling_loop::run_sampling_request; use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::stream_events_utils::HandleOutputCtx; @@ -57,7 +58,6 @@ use crate::tools::router::ToolRouterParams; use crate::turn_diff_tracker::TurnDiffTracker; use crate::turn_timing::record_turn_ttft_metric; use crate::unavailable_tool::collect_unavailable_called_tools; -use crate::util::backoff; use crate::util::error_or_panic; use codex_analytics::AppInvocation; use codex_analytics::CompactionPhase; @@ -71,6 +71,9 @@ use codex_hooks::HookEvent; use codex_hooks::HookEventAfterAgent; use codex_hooks::HookPayload; use codex_hooks::HookResult; +use codex_kernel::PromptConfig; +use codex_kernel::SamplingRequestResult; +use codex_kernel::ToolConfig; use codex_protocol::config_types::ModeKind; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; @@ -932,16 +935,11 @@ fn connector_inserted_in_messages( connector_count == 1 && skill_count == 0 && mention_names_lower.contains(&mention_slug) } -pub(crate) fn build_prompt( - input: Vec, - router: &ToolRouter, +pub(super) fn build_prompt_config( turn_context: &TurnContext, base_instructions: BaseInstructions, -) -> Prompt { - Prompt { - input, - tools: router.model_visible_specs(), - parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls, +) -> PromptConfig { + PromptConfig { base_instructions, personality: turn_context.personality, output_schema: turn_context.final_output_json_schema.clone(), @@ -951,156 +949,23 @@ pub(crate) fn build_prompt( } } -#[allow(clippy::too_many_arguments)] -#[instrument(level = "trace", - skip_all, - fields( - turn_id = %turn_context.sub_id, - model = %turn_context.model_info.slug, - cwd = %turn_context.cwd.display() - ) -)] -async fn run_sampling_request( - sess: Arc, - turn_context: Arc, - turn_diff_tracker: SharedTurnDiffTracker, - client_session: &mut ModelClientSession, - turn_metadata_header: Option<&str>, - input: Vec, - explicitly_enabled_connectors: &HashSet, - skills_outcome: Option<&SkillLoadOutcome>, - cancellation_token: CancellationToken, -) -> CodexResult { - let router = built_tools( - sess.as_ref(), - turn_context.as_ref(), - &input, - explicitly_enabled_connectors, - skills_outcome, - &cancellation_token, - ) - .await?; - - let base_instructions = sess.get_base_instructions().await; - - let tool_runtime = ToolCallRuntime::new( - Arc::clone(&router), - Arc::clone(&sess), - Arc::clone(&turn_context), - Arc::clone(&turn_diff_tracker), - ); - let _code_mode_worker = sess - .services - .code_mode_service - .start_turn_worker( - &sess, - &turn_context, - Arc::clone(&router), - Arc::clone(&turn_diff_tracker), - ) - .await; - let mut retries = 0; - let mut initial_input = Some(input); - loop { - let prompt_input = if let Some(input) = initial_input.take() { - input - } else { - sess.clone_history() - .await - .for_prompt(&turn_context.model_info.input_modalities) - }; - let prompt = build_prompt( - prompt_input, - router.as_ref(), - turn_context.as_ref(), - base_instructions.clone(), - ); - let err = match try_run_sampling_request( - tool_runtime.clone(), - Arc::clone(&sess), - Arc::clone(&turn_context), - client_session, - turn_metadata_header, - Arc::clone(&turn_diff_tracker), - &prompt, - cancellation_token.child_token(), - ) - .await - { - Ok(output) => { - return Ok(output); - } - Err(CodexErr::ContextWindowExceeded) => { - sess.set_total_tokens_full(&turn_context).await; - return Err(CodexErr::ContextWindowExceeded); - } - Err(CodexErr::UsageLimitReached(e)) => { - let rate_limits = e.rate_limits.clone(); - if let Some(rate_limits) = rate_limits { - sess.update_rate_limits(&turn_context, *rate_limits).await; - } - return Err(CodexErr::UsageLimitReached(e)); - } - Err(err) => err, - }; - - if !err.is_retryable() { - return Err(err); - } - - // Use the configured provider-specific stream retry budget. - let max_retries = turn_context.provider.info().stream_max_retries(); - if retries >= max_retries - && client_session.try_switch_fallback_transport( - &turn_context.session_telemetry, - &turn_context.model_info, - ) - { - sess.send_event( - &turn_context, - EventMsg::Warning(WarningEvent { - message: format!("Falling back from WebSockets to HTTPS transport. {err:#}"), - }), - ) - .await; - retries = 0; - continue; - } - if retries < max_retries { - retries += 1; - let delay = match &err { - CodexErr::Stream(_, requested_delay) => { - requested_delay.unwrap_or_else(|| backoff(retries)) - } - _ => backoff(retries), - }; - warn!( - "stream disconnected - retrying sampling request ({retries}/{max_retries} in {delay:?})...", - ); - - // In release builds, hide the first websocket retry notification to reduce noisy - // transient reconnect messages. In debug builds, keep full visibility for diagnosis. - let report_error = retries > 1 - || cfg!(debug_assertions) - || !sess.services.model_client.responses_websocket_enabled(); - if report_error { - // Surface retry information to any UI/front‑end so the - // user understands what is happening instead of staring - // at a seemingly frozen screen. - sess.notify_stream_error( - &turn_context, - format!("Reconnecting... {retries}/{max_retries}"), - err, - ) - .await; - } - tokio::time::sleep(delay).await; - } else { - return Err(err); - } +pub(super) fn build_tool_config(router: &ToolRouter, turn_context: &TurnContext) -> ToolConfig { + ToolConfig { + tools: router.model_visible_specs(), + parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls, } } +pub(crate) fn build_prompt( + input: Vec, + router: &ToolRouter, + turn_context: &TurnContext, + base_instructions: BaseInstructions, +) -> Prompt { + build_prompt_config(turn_context, base_instructions) + .build_prompt(input, &build_tool_config(router, turn_context)) +} + #[expect( clippy::await_holding_invalid_type, reason = "tool router construction reads through the session-owned manager guard" @@ -1243,12 +1108,6 @@ pub(crate) async fn built_tools( ))) } -#[derive(Debug)] -struct SamplingRequestResult { - needs_follow_up: bool, - last_agent_message: Option, -} - /// Ephemeral per-response state for streaming a single proposed plan. /// This is intentionally not persisted or stored in session/state since it /// only exists while a response is actively streaming. The final plan text @@ -1806,7 +1665,7 @@ async fn drain_in_flight( model = %turn_context.model_info.slug ) )] -async fn try_run_sampling_request( +pub(super) async fn try_run_sampling_request( tool_runtime: ToolCallRuntime, sess: Arc, turn_context: Arc, diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 5a31d18020..903d8a9025 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -10,18 +10,18 @@ use tokio_util::sync::CancellationToken; use crate::context::ContextualUserFragment; use crate::context::ImageGenerationInstructions; -use crate::function_tool::FunctionCallError; use crate::parse_turn_item; use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::tools::parallel::ToolCallRuntime; -use crate::tools::router::ToolRouter; +use codex_kernel::CompletedResponseItem; +use codex_kernel::KernelToolExecutor; +use codex_kernel::execute_tool_call_with_default_output; +use codex_kernel::response_input_to_response_item; use codex_memories_read::citations::parse_memory_citation; use codex_memories_read::citations::thread_ids_from_memory_citation; use codex_protocol::error::CodexErr; use codex_protocol::error::Result; -use codex_protocol::models::FunctionCallOutputBody; -use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::MessagePhase; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; @@ -225,9 +225,8 @@ pub(crate) async fn handle_output_item_done( let mut output = OutputItemResult::default(); let plan_mode = ctx.turn_context.collaboration_mode.mode == ModeKind::Plan; - match ToolRouter::build_tool_call(ctx.sess.as_ref(), item.clone()).await { - // The model emitted a tool call; log it, persist the item immediately, and queue the tool execution. - Ok(Some(call)) => { + match ctx.tool_runtime.classify_response_item(item).await? { + CompletedResponseItem::ToolCall { item, call } => { ctx.sess .accept_mailbox_delivery_for_current_turn(&ctx.turn_context.sub_id) .await; @@ -244,17 +243,15 @@ pub(crate) async fn handle_output_item_done( .await; let cancellation_token = ctx.cancellation_token.child_token(); - let tool_future: InFlightFuture<'static> = Box::pin( - ctx.tool_runtime - .clone() - .handle_tool_call(call, cancellation_token), - ); + let tool_runtime = ctx.tool_runtime.clone(); + let tool_future: InFlightFuture<'static> = Box::pin(async move { + execute_tool_call_with_default_output(&tool_runtime, call, cancellation_token).await + }); output.needs_follow_up = true; output.tool_future = Some(tool_future); } - // No tool call: convert messages/reasoning into turn items and mark them as complete. - Ok(None) => { + CompletedResponseItem::NonTool(item) => { if let Some(turn_item) = handle_non_tool_response_item( ctx.sess.as_ref(), ctx.turn_context.as_ref(), @@ -286,21 +283,7 @@ pub(crate) async fn handle_output_item_done( output.last_agent_message = last_agent_message; } - // Guardrail: the model issued a LocalShellCall without an id; surface the error back into history. - Err(FunctionCallError::MissingLocalShellCallId) => { - let msg = "LocalShellCall without call_id or id"; - ctx.turn_context - .session_telemetry - .log_tool_failed("local_shell", msg); - tracing::error!(msg); - - let response = ResponseInputItem::FunctionCallOutput { - call_id: String::new(), - output: FunctionCallOutputPayload { - body: FunctionCallOutputBody::Text(msg.to_string()), - ..Default::default() - }, - }; + CompletedResponseItem::ImmediateResponse { item, response } => { record_completed_response_item(ctx.sess.as_ref(), ctx.turn_context.as_ref(), &item) .await; if let Some(response_item) = response_input_to_response_item(&response) { @@ -314,32 +297,6 @@ pub(crate) async fn handle_output_item_done( output.needs_follow_up = true; } - // The tool request should be answered directly (or was denied); push that response into the transcript. - Err(FunctionCallError::RespondToModel(message)) => { - let response = ResponseInputItem::FunctionCallOutput { - call_id: String::new(), - output: FunctionCallOutputPayload { - body: FunctionCallOutputBody::Text(message), - ..Default::default() - }, - }; - record_completed_response_item(ctx.sess.as_ref(), ctx.turn_context.as_ref(), &item) - .await; - if let Some(response_item) = response_input_to_response_item(&response) { - ctx.sess - .record_conversation_items( - &ctx.turn_context, - std::slice::from_ref(&response_item), - ) - .await; - } - - output.needs_follow_up = true; - } - // A fatal error occurred; surface it back into history. - Err(FunctionCallError::Fatal(message)) => { - return Err(CodexErr::Fatal(message)); - } } Ok(output) @@ -465,45 +422,6 @@ fn completed_item_defers_mailbox_delivery_to_next_turn( } } -pub(crate) fn response_input_to_response_item(input: &ResponseInputItem) -> Option { - match input { - ResponseInputItem::FunctionCallOutput { call_id, output } => { - Some(ResponseItem::FunctionCallOutput { - call_id: call_id.clone(), - output: output.clone(), - }) - } - ResponseInputItem::CustomToolCallOutput { - call_id, - name, - output, - } => Some(ResponseItem::CustomToolCallOutput { - call_id: call_id.clone(), - name: name.clone(), - output: output.clone(), - }), - ResponseInputItem::McpToolCallOutput { call_id, output } => { - let output = output.as_function_call_output_payload(); - Some(ResponseItem::FunctionCallOutput { - call_id: call_id.clone(), - output, - }) - } - ResponseInputItem::ToolSearchOutput { - call_id, - status, - execution, - tools, - } => Some(ResponseItem::ToolSearchOutput { - call_id: Some(call_id.clone()), - status: status.clone(), - execution: execution.clone(), - tools: tools.clone(), - }), - _ => None, - } -} - #[cfg(test)] #[path = "stream_events_utils_tests.rs"] mod tests; diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index 384a800b4b..ddf547ec90 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -20,8 +20,14 @@ use crate::tools::registry::ToolArgumentDiffConsumer; use crate::tools::router::ToolCall; use crate::tools::router::ToolCallSource; use crate::tools::router::ToolRouter; +use codex_kernel::CompletedResponseItem; +use codex_kernel::KernelToolCall; +use codex_kernel::KernelToolExecutor; +use codex_kernel::ToolCallError; use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; use codex_protocol::models::ResponseInputItem; +use codex_protocol::models::ResponseItem; use codex_tools::ToolSpec; #[derive(Clone)] @@ -60,25 +66,6 @@ impl ToolCallRuntime { self.router.create_diff_consumer(tool_name) } - #[instrument(level = "trace", skip_all)] - pub(crate) fn handle_tool_call( - self, - call: ToolCall, - cancellation_token: CancellationToken, - ) -> impl std::future::Future> { - let error_call = call.clone(); - let future = - self.handle_tool_call_with_source(call, ToolCallSource::Direct, cancellation_token); - async move { - match future.await { - Ok(response) => Ok(response.into_response()), - Err(FunctionCallError::Fatal(message)) => Err(CodexErr::Fatal(message)), - Err(other) => Ok(Self::failure_response(error_call, other)), - } - } - .in_current_span() - } - #[instrument(level = "trace", skip_all)] pub(crate) fn handle_tool_call_with_source( self, @@ -144,33 +131,6 @@ impl ToolCallRuntime { } impl ToolCallRuntime { - fn failure_response(call: ToolCall, err: FunctionCallError) -> ResponseInputItem { - let message = err.to_string(); - match call.payload { - ToolPayload::ToolSearch { .. } => ResponseInputItem::ToolSearchOutput { - call_id: call.call_id, - status: "completed".to_string(), - execution: "client".to_string(), - tools: Vec::new(), - }, - ToolPayload::Custom { .. } => ResponseInputItem::CustomToolCallOutput { - call_id: call.call_id, - name: None, - output: codex_protocol::models::FunctionCallOutputPayload { - body: codex_protocol::models::FunctionCallOutputBody::Text(message), - success: Some(false), - }, - }, - _ => ResponseInputItem::FunctionCallOutput { - call_id: call.call_id, - output: codex_protocol::models::FunctionCallOutputPayload { - body: codex_protocol::models::FunctionCallOutputBody::Text(message), - success: Some(false), - }, - }, - } - } - fn aborted_response(call: &ToolCall, secs: f32) -> AnyToolResult { AnyToolResult { call_id: call.call_id.clone(), @@ -195,3 +155,87 @@ impl ToolCallRuntime { } } } + +impl KernelToolCall for ToolCall { + fn error_response(&self, message: String) -> ResponseInputItem { + failure_response(self, message) + } +} + +impl KernelToolExecutor for ToolCallRuntime { + type Call = ToolCall; + + async fn classify_response_item( + &self, + item: ResponseItem, + ) -> CodexResult> { + match ToolRouter::build_tool_call(self.session.as_ref(), item.clone()).await { + Ok(Some(call)) => Ok(CompletedResponseItem::ToolCall { item, call }), + Ok(None) => Ok(CompletedResponseItem::NonTool(item)), + Err(FunctionCallError::MissingLocalShellCallId) => { + let response = ResponseInputItem::FunctionCallOutput { + call_id: String::new(), + output: codex_protocol::models::FunctionCallOutputPayload::from_text( + "LocalShellCall without call_id or id".to_string(), + ), + }; + Ok(CompletedResponseItem::ImmediateResponse { item, response }) + } + Err(FunctionCallError::RespondToModel(message)) => { + let response = ResponseInputItem::FunctionCallOutput { + call_id: String::new(), + output: codex_protocol::models::FunctionCallOutputPayload::from_text(message), + }; + Ok(CompletedResponseItem::ImmediateResponse { item, response }) + } + Err(FunctionCallError::Fatal(message)) => Err(CodexErr::Fatal(message)), + } + } + + async fn execute_tool_call( + &self, + call: Self::Call, + cancellation_token: CancellationToken, + ) -> Result { + match self + .clone() + .handle_tool_call_with_source(call, ToolCallSource::Direct, cancellation_token) + .await + { + Ok(result) => Ok(result.into_response()), + Err(FunctionCallError::RespondToModel(message)) => { + Err(ToolCallError::RespondToModel(message)) + } + Err(FunctionCallError::MissingLocalShellCallId) => Err(ToolCallError::RespondToModel( + "LocalShellCall without call_id or id".to_string(), + )), + Err(FunctionCallError::Fatal(message)) => Err(ToolCallError::Fatal(message)), + } + } +} + +fn failure_response(call: &ToolCall, message: String) -> ResponseInputItem { + match &call.payload { + ToolPayload::ToolSearch { .. } => ResponseInputItem::ToolSearchOutput { + call_id: call.call_id.clone(), + status: "completed".to_string(), + execution: "client".to_string(), + tools: Vec::new(), + }, + ToolPayload::Custom { .. } => ResponseInputItem::CustomToolCallOutput { + call_id: call.call_id.clone(), + name: None, + output: codex_protocol::models::FunctionCallOutputPayload { + body: codex_protocol::models::FunctionCallOutputBody::Text(message), + success: Some(false), + }, + }, + _ => ResponseInputItem::FunctionCallOutput { + call_id: call.call_id.clone(), + output: codex_protocol::models::FunctionCallOutputPayload { + body: codex_protocol::models::FunctionCallOutputBody::Text(message), + success: Some(false), + }, + }, + } +} diff --git a/codex-rs/core/src/tools/runtimes/mod_tests.rs b/codex-rs/core/src/tools/runtimes/mod_tests.rs index 6aa64d1e37..dd6f38ca1f 100644 --- a/codex-rs/core/src/tools/runtimes/mod_tests.rs +++ b/codex-rs/core/src/tools/runtimes/mod_tests.rs @@ -675,10 +675,13 @@ fn maybe_wrap_shell_lc_with_snapshot_keeps_user_proxy_env_when_proxy_inactive() &HashMap::new(), &HashMap::new(), ); - let output = Command::new(&rewritten[0]) - .args(&rewritten[1..]) - .output() - .expect("run rewritten command"); + let mut process = Command::new(&rewritten[0]); + process.args(&rewritten[1..]); + for key in PROXY_ENV_KEYS { + process.env_remove(key); + } + process.env_remove(PROXY_ACTIVE_ENV_KEY); + let output = process.output().expect("run rewritten command"); assert!(output.status.success(), "command failed: {output:?}"); assert_eq!( diff --git a/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs b/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs index 2ec5ede4a8..6975d3caf5 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation_tests.rs @@ -338,23 +338,25 @@ async fn execve_permission_request_hook_short_circuits_prompt() -> anyhow::Resul std::fs::write( &script_path, format!( - "#!/bin/sh\ncat > {log_path}\nprintf '%s\\n' '{response}'\n", - log_path = shlex::try_quote(log_path.to_string_lossy().as_ref())?, + r#"import json +from pathlib import Path +import sys + +payload = json.load(sys.stdin) +with Path(r"{log_path}").open("a", encoding="utf-8") as handle: + handle.write(json.dumps(payload) + "\n") + +print({response:?}) +"#, + log_path = log_path.display(), response = "{\"hookSpecificOutput\":{\"hookEventName\":\"PermissionRequest\",\"decision\":{\"behavior\":\"allow\"}}}", ), ) .with_context(|| format!("write hook script to {}", script_path.display()))?; - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - - let mut permissions = std::fs::metadata(&script_path) - .with_context(|| format!("read hook script metadata from {}", script_path.display()))? - .permissions(); - permissions.set_mode(0o755); - std::fs::set_permissions(&script_path, permissions) - .with_context(|| format!("set hook script permissions on {}", script_path.display()))?; - } + let script_path_arg = format!( + "'{}'", + script_path.display().to_string().replace('\'', "'\\''") + ); std::fs::write( turn_context.config.codex_home.join("hooks.json"), serde_json::json!({ @@ -362,7 +364,7 @@ async fn execve_permission_request_hook_short_circuits_prompt() -> anyhow::Resul "PermissionRequest": [{ "hooks": [{ "type": "command", - "command": script_path.display().to_string(), + "command": format!("python3 {script_path_arg}"), }] }] } diff --git a/codex-rs/kernel/Cargo.toml b/codex-rs/kernel/Cargo.toml new file mode 100644 index 0000000000..cfc54e89e1 --- /dev/null +++ b/codex-rs/kernel/Cargo.toml @@ -0,0 +1,25 @@ +[package] +edition.workspace = true +license.workspace = true +name = "codex-kernel" +version.workspace = true + +[lib] +doctest = false +name = "codex_kernel" +path = "src/lib.rs" + +[lints] +workspace = true + +[dependencies] +codex-protocol = { workspace = true } +codex-tools = { workspace = true } +rand = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } +tokio = { workspace = true, features = ["time"] } +tokio-util = { workspace = true, features = ["rt"] } + +[dev-dependencies] +pretty_assertions = { workspace = true } diff --git a/codex-rs/kernel/src/lib.rs b/codex-rs/kernel/src/lib.rs new file mode 100644 index 0000000000..29cf6b8b5b --- /dev/null +++ b/codex-rs/kernel/src/lib.rs @@ -0,0 +1,17 @@ +mod prompt; +mod sampling; +mod tools; + +pub use prompt::Prompt; +pub use prompt::PromptConfig; +pub use sampling::PreparedSamplingRequest; +pub use sampling::SamplingLoopHost; +pub use sampling::SamplingRequestResult; +pub use sampling::run_sampling_request_loop; +pub use tools::CompletedResponseItem; +pub use tools::KernelToolCall; +pub use tools::KernelToolExecutor; +pub use tools::ToolCallError; +pub use tools::ToolConfig; +pub use tools::execute_tool_call_with_default_output; +pub use tools::response_input_to_response_item; diff --git a/codex-rs/kernel/src/prompt.rs b/codex-rs/kernel/src/prompt.rs new file mode 100644 index 0000000000..31d5e4c2ae --- /dev/null +++ b/codex-rs/kernel/src/prompt.rs @@ -0,0 +1,257 @@ +use crate::ToolConfig; +use codex_protocol::config_types::Personality; +use codex_protocol::models::BaseInstructions; +use codex_protocol::models::FunctionCallOutputBody; +use codex_protocol::models::ResponseItem; +use codex_tools::ToolSpec; +use serde::Deserialize; +use serde_json::Value; +use std::collections::HashSet; + +/// API request payload for a single model turn. +#[derive(Debug, Clone)] +pub struct Prompt { + pub input: Vec, + pub tools: Vec, + pub parallel_tool_calls: bool, + pub base_instructions: BaseInstructions, + pub personality: Option, + pub output_schema: Option, + pub output_schema_strict: bool, +} + +/// Retry-stable prompt settings used to rebuild a [`Prompt`] after tool calls or stream retries. +#[derive(Debug, Clone)] +pub struct PromptConfig { + pub base_instructions: BaseInstructions, + pub personality: Option, + pub output_schema: Option, + pub output_schema_strict: bool, +} + +impl Prompt { + pub fn get_formatted_input(&self) -> Vec { + let mut input = self.input.clone(); + let is_freeform_apply_patch_tool_present = self.tools.iter().any(|tool| match tool { + ToolSpec::Freeform(f) => f.name == "apply_patch", + _ => false, + }); + if is_freeform_apply_patch_tool_present { + reserialize_shell_outputs(&mut input); + } + + input + } +} + +impl PromptConfig { + pub fn build_prompt(&self, input: Vec, tool_config: &ToolConfig) -> Prompt { + Prompt { + input, + tools: tool_config.tools.clone(), + parallel_tool_calls: tool_config.parallel_tool_calls, + base_instructions: self.base_instructions.clone(), + personality: self.personality, + output_schema: self.output_schema.clone(), + output_schema_strict: self.output_schema_strict, + } + } +} + +impl Default for Prompt { + fn default() -> Self { + Self { + input: Vec::new(), + tools: Vec::new(), + parallel_tool_calls: false, + base_instructions: BaseInstructions::default(), + personality: None, + output_schema: None, + output_schema_strict: true, + } + } +} + +impl Default for PromptConfig { + fn default() -> Self { + Self { + base_instructions: BaseInstructions::default(), + personality: None, + output_schema: None, + output_schema_strict: true, + } + } +} + +fn reserialize_shell_outputs(items: &mut [ResponseItem]) { + let mut shell_call_ids: HashSet = HashSet::new(); + + items.iter_mut().for_each(|item| match item { + ResponseItem::LocalShellCall { call_id, id, .. } => { + if let Some(identifier) = call_id.clone().or_else(|| id.clone()) { + shell_call_ids.insert(identifier); + } + } + ResponseItem::CustomToolCall { + id: _, + status: _, + call_id, + name, + input: _, + } => { + if name == "apply_patch" { + shell_call_ids.insert(call_id.clone()); + } + } + ResponseItem::FunctionCall { name, call_id, .. } + if is_shell_tool_name(name) || name == "apply_patch" => + { + shell_call_ids.insert(call_id.clone()); + } + ResponseItem::FunctionCallOutput { + call_id, output, .. + } + | ResponseItem::CustomToolCallOutput { + call_id, output, .. + } => { + if shell_call_ids.remove(call_id) + && let Some(structured) = output + .text_content() + .and_then(parse_structured_shell_output) + { + output.body = FunctionCallOutputBody::Text(structured); + } + } + _ => {} + }); +} + +fn is_shell_tool_name(name: &str) -> bool { + matches!(name, "shell" | "container.exec") +} + +#[derive(Deserialize)] +struct ExecOutputJson { + output: String, + metadata: ExecOutputMetadataJson, +} + +#[derive(Deserialize)] +struct ExecOutputMetadataJson { + exit_code: i32, + duration_seconds: f32, +} + +fn parse_structured_shell_output(raw: &str) -> Option { + let parsed: ExecOutputJson = serde_json::from_str(raw).ok()?; + Some(build_structured_output(&parsed)) +} + +fn build_structured_output(parsed: &ExecOutputJson) -> String { + let mut sections = Vec::new(); + let exit_code = parsed.metadata.exit_code; + sections.push(format!("Exit code: {exit_code}")); + let duration_seconds = parsed.metadata.duration_seconds; + sections.push(format!("Wall time: {duration_seconds} seconds")); + + let mut output = parsed.output.clone(); + if let Some((stripped, total_lines)) = strip_total_output_header(&parsed.output) { + sections.push(format!("Total output lines: {total_lines}")); + output = stripped.to_string(); + } + + sections.push("Output:".to_string()); + sections.push(output); + sections.join("\n") +} + +fn strip_total_output_header(output: &str) -> Option<(&str, u32)> { + let after_prefix = output.strip_prefix("Total output lines: ")?; + let (total_segment, remainder) = after_prefix.split_once('\n')?; + let total_lines = total_segment.parse::().ok()?; + let remainder = remainder.strip_prefix('\n').unwrap_or(remainder); + Some((remainder, total_lines)) +} + +#[cfg(test)] +mod tests { + use super::Prompt; + use codex_protocol::models::FunctionCallOutputPayload; + use codex_protocol::models::ResponseItem; + use codex_tools::FreeformTool; + use codex_tools::FreeformToolFormat; + use codex_tools::ToolSpec; + use pretty_assertions::assert_eq; + + #[test] + fn reserializes_shell_outputs_for_function_and_custom_tool_calls() { + let raw_output = r#"{"output":"hello","metadata":{"exit_code":0,"duration_seconds":0.5}}"#; + let expected_output = "Exit code: 0\nWall time: 0.5 seconds\nOutput:\nhello"; + let prompt = Prompt { + input: vec![ + ResponseItem::FunctionCall { + id: None, + name: "shell".to_string(), + namespace: None, + arguments: "{}".to_string(), + call_id: "call-1".to_string(), + }, + ResponseItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload::from_text(raw_output.to_string()), + }, + ResponseItem::CustomToolCall { + id: None, + status: None, + call_id: "call-2".to_string(), + name: "apply_patch".to_string(), + input: "*** Begin Patch".to_string(), + }, + ResponseItem::CustomToolCallOutput { + call_id: "call-2".to_string(), + name: None, + output: FunctionCallOutputPayload::from_text(raw_output.to_string()), + }, + ], + tools: vec![ToolSpec::Freeform(FreeformTool { + name: "apply_patch".to_string(), + description: "patch".to_string(), + format: FreeformToolFormat { + r#type: "grammar".to_string(), + syntax: "lark".to_string(), + definition: "patch".to_string(), + }, + })], + ..Prompt::default() + }; + + assert_eq!( + prompt.get_formatted_input(), + vec![ + ResponseItem::FunctionCall { + id: None, + name: "shell".to_string(), + namespace: None, + arguments: "{}".to_string(), + call_id: "call-1".to_string(), + }, + ResponseItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload::from_text(expected_output.to_string()), + }, + ResponseItem::CustomToolCall { + id: None, + status: None, + call_id: "call-2".to_string(), + name: "apply_patch".to_string(), + input: "*** Begin Patch".to_string(), + }, + ResponseItem::CustomToolCallOutput { + call_id: "call-2".to_string(), + name: None, + output: FunctionCallOutputPayload::from_text(expected_output.to_string()), + }, + ] + ); + } +} diff --git a/codex-rs/kernel/src/sampling.rs b/codex-rs/kernel/src/sampling.rs new file mode 100644 index 0000000000..7c9ed674c3 --- /dev/null +++ b/codex-rs/kernel/src/sampling.rs @@ -0,0 +1,170 @@ +use crate::KernelToolExecutor; +use crate::Prompt; +use crate::PromptConfig; +use crate::ToolConfig; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::RateLimitSnapshot; +use rand::Rng; +use std::time::Duration; +use tokio_util::sync::CancellationToken; + +const INITIAL_DELAY_MS: u64 = 200; +const BACKOFF_FACTOR: f64 = 2.0; + +/// Host-facing bundle created once before the kernel begins retrying a logical sampling request. +/// +/// Implementations should place retry-stable state here so retries can rebuild prompts without +/// repeating expensive setup work such as tool routing or per-turn worker startup. +pub struct PreparedSamplingRequest { + pub prompt_config: PromptConfig, + pub tool_config: ToolConfig, + pub runtime: Runtime, + pub tools: Tools, +} + +/// Final outcome of a logical sampling request after the model/tool loop reaches a stable state. +#[derive(Debug)] +pub struct SamplingRequestResult { + pub needs_follow_up: bool, + pub last_agent_message: Option, +} + +/// Host adapter for the kernel's retryable sampling loop. +/// +/// Implementations provide the concrete model transport, tool runtime, history access, and +/// user-facing retry notifications while allowing `codex-kernel` to stay agnostic of local +/// filesystem, network, or tool-execution details. +#[allow(async_fn_in_trait)] +pub trait SamplingLoopHost { + type ClientSession; + type Runtime; + type Tools: KernelToolExecutor; + + async fn prepare_sampling_request( + &self, + input: &[ResponseItem], + cancellation_token: &CancellationToken, + ) -> Result>; + + async fn history_prompt_input(&self) -> Vec; + + async fn run_single_sampling_request( + &self, + runtime: &Self::Runtime, + tools: &Self::Tools, + client_session: &mut Self::ClientSession, + prompt: &Prompt, + cancellation_token: CancellationToken, + ) -> Result; + + fn stream_max_retries(&self) -> u64; + + fn try_switch_fallback_transport(&self, client_session: &mut Self::ClientSession) -> bool; + + fn should_notify_stream_retry(&self, retries: u64, err: &CodexErr) -> bool; + + async fn handle_context_window_exceeded(&self); + + async fn handle_usage_limit_reached(&self, rate_limits: Option); + + async fn notify_fallback_to_http(&self, err: &CodexErr); + + async fn notify_stream_retry( + &self, + retries: u64, + max_retries: u64, + delay: Duration, + err: &CodexErr, + ); +} + +pub async fn run_sampling_request_loop( + host: &Host, + client_session: &mut Host::ClientSession, + initial_input: Vec, + cancellation_token: CancellationToken, +) -> Result +where + Host: SamplingLoopHost, +{ + let prepared = host + .prepare_sampling_request(initial_input.as_slice(), &cancellation_token) + .await?; + let mut retries = 0; + let mut initial_input = Some(initial_input); + loop { + let prompt_input = if let Some(input) = initial_input.take() { + input + } else { + host.history_prompt_input().await + }; + let prompt = prepared + .prompt_config + .build_prompt(prompt_input, &prepared.tool_config); + let err = match host + .run_single_sampling_request( + &prepared.runtime, + &prepared.tools, + client_session, + &prompt, + cancellation_token.child_token(), + ) + .await + { + Ok(output) => { + return Ok(output); + } + Err(CodexErr::ContextWindowExceeded) => { + host.handle_context_window_exceeded().await; + return Err(CodexErr::ContextWindowExceeded); + } + Err(CodexErr::UsageLimitReached(error)) => { + host.handle_usage_limit_reached( + error + .rate_limits + .as_ref() + .map(|snapshot| (**snapshot).clone()), + ) + .await; + return Err(CodexErr::UsageLimitReached(error)); + } + Err(err) => err, + }; + + if !err.is_retryable() { + return Err(err); + } + + let max_retries = host.stream_max_retries(); + if retries >= max_retries && host.try_switch_fallback_transport(client_session) { + host.notify_fallback_to_http(&err).await; + retries = 0; + continue; + } + if retries < max_retries { + retries += 1; + let delay = match &err { + CodexErr::Stream(_, requested_delay) => { + requested_delay.unwrap_or_else(|| backoff(retries)) + } + _ => backoff(retries), + }; + if host.should_notify_stream_retry(retries, &err) { + host.notify_stream_retry(retries, max_retries, delay, &err) + .await; + } + tokio::time::sleep(delay).await; + } else { + return Err(err); + } + } +} + +fn backoff(attempt: u64) -> Duration { + let exp = BACKOFF_FACTOR.powi(attempt.saturating_sub(1) as i32); + let base = (INITIAL_DELAY_MS as f64 * exp) as u64; + let jitter = rand::rng().random_range(0.9..1.1); + Duration::from_millis((base as f64 * jitter) as u64) +} diff --git a/codex-rs/kernel/src/tools.rs b/codex-rs/kernel/src/tools.rs new file mode 100644 index 0000000000..131421da4f --- /dev/null +++ b/codex-rs/kernel/src/tools.rs @@ -0,0 +1,226 @@ +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result; +use codex_protocol::models::ResponseInputItem; +use codex_protocol::models::ResponseItem; +use codex_tools::ToolSpec; +use std::fmt; +use tokio_util::sync::CancellationToken; + +/// Retry-stable tool registration for a logical model request. +/// +/// The host chooses which tools are exposed to the model and whether the model +/// may issue them in parallel. `codex-kernel` treats this as opaque model-facing +/// configuration and uses it to rebuild prompts across retries. +#[derive(Debug, Clone, Default)] +pub struct ToolConfig { + pub tools: Vec, + pub parallel_tool_calls: bool, +} + +/// Metadata a host-specific tool call must expose so `codex-kernel` can +/// synthesize a model-visible failure response when execution cannot complete. +pub trait KernelToolCall: Clone + Send + Sync { + fn error_response(&self, message: String) -> ResponseInputItem; +} + +/// Non-fatal tool execution failures are surfaced back to the model as a +/// complementary output item, while fatal failures abort the turn. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ToolCallError { + RespondToModel(String), + Fatal(String), +} + +impl fmt::Display for ToolCallError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::RespondToModel(message) => write!(f, "{message}"), + Self::Fatal(message) => write!(f, "Fatal error: {message}"), + } + } +} + +/// Classification of a completed model output item relative to the host's tool +/// catalog. +pub enum CompletedResponseItem { + NonTool(ResponseItem), + ToolCall { + item: ResponseItem, + call: Call, + }, + ImmediateResponse { + item: ResponseItem, + response: ResponseInputItem, + }, +} + +/// Host-provided tool adapter used by `codex-kernel`'s default model/tool loop. +/// +/// Implementations own tool-call parsing and execution, while +/// `codex-kernel` owns the generic behavior of attempting a complementary +/// tool output item for every tool call the model emits. +#[allow(async_fn_in_trait)] +pub trait KernelToolExecutor { + type Call: KernelToolCall; + + async fn classify_response_item( + &self, + item: ResponseItem, + ) -> Result>; + + async fn execute_tool_call( + &self, + call: Self::Call, + cancellation_token: CancellationToken, + ) -> std::result::Result; +} + +pub async fn execute_tool_call_with_default_output( + executor: &Executor, + call: Executor::Call, + cancellation_token: CancellationToken, +) -> Result +where + Executor: KernelToolExecutor, +{ + match executor + .execute_tool_call(call.clone(), cancellation_token) + .await + { + Ok(response) => Ok(response), + Err(ToolCallError::RespondToModel(message)) => Ok(call.error_response(message)), + Err(ToolCallError::Fatal(message)) => Err(CodexErr::Fatal(message)), + } +} + +pub fn response_input_to_response_item(input: &ResponseInputItem) -> Option { + match input { + ResponseInputItem::FunctionCallOutput { call_id, output } => { + Some(ResponseItem::FunctionCallOutput { + call_id: call_id.clone(), + output: output.clone(), + }) + } + ResponseInputItem::CustomToolCallOutput { + call_id, + name, + output, + } => Some(ResponseItem::CustomToolCallOutput { + call_id: call_id.clone(), + name: name.clone(), + output: output.clone(), + }), + ResponseInputItem::McpToolCallOutput { call_id, output } => { + let output = output.as_function_call_output_payload(); + Some(ResponseItem::FunctionCallOutput { + call_id: call_id.clone(), + output, + }) + } + ResponseInputItem::ToolSearchOutput { + call_id, + status, + execution, + tools, + } => Some(ResponseItem::ToolSearchOutput { + call_id: Some(call_id.clone()), + status: status.clone(), + execution: execution.clone(), + tools: tools.clone(), + }), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::CompletedResponseItem; + use super::KernelToolCall; + use super::KernelToolExecutor; + use super::ToolCallError; + use super::execute_tool_call_with_default_output; + use codex_protocol::error::Result; + use codex_protocol::models::FunctionCallOutputPayload; + use codex_protocol::models::ResponseInputItem; + use codex_protocol::models::ResponseItem; + use pretty_assertions::assert_eq; + use tokio_util::sync::CancellationToken; + + #[derive(Clone, Debug, PartialEq, Eq)] + struct TestCall { + call_id: String, + } + + impl KernelToolCall for TestCall { + fn error_response(&self, message: String) -> ResponseInputItem { + ResponseInputItem::FunctionCallOutput { + call_id: self.call_id.clone(), + output: FunctionCallOutputPayload::from_text(message), + } + } + } + + struct TestExecutor; + + impl KernelToolExecutor for TestExecutor { + type Call = TestCall; + + async fn classify_response_item( + &self, + item: ResponseItem, + ) -> Result> { + Ok(CompletedResponseItem::NonTool(item)) + } + + async fn execute_tool_call( + &self, + call: Self::Call, + _cancellation_token: CancellationToken, + ) -> std::result::Result { + match call.call_id.as_str() { + "ok" => Ok(ResponseInputItem::FunctionCallOutput { + call_id: call.call_id, + output: FunctionCallOutputPayload::from_text("done".to_string()), + }), + "retry" => Err(ToolCallError::RespondToModel("try again".to_string())), + _ => Err(ToolCallError::Fatal("boom".to_string())), + } + } + } + + #[tokio::test] + async fn execute_tool_call_with_default_output_surfaces_non_fatal_errors() { + let response = execute_tool_call_with_default_output( + &TestExecutor, + TestCall { + call_id: "retry".to_string(), + }, + CancellationToken::new(), + ) + .await + .expect("tool response"); + + assert_eq!( + response, + ResponseInputItem::FunctionCallOutput { + call_id: "retry".to_string(), + output: FunctionCallOutputPayload::from_text("try again".to_string()), + } + ); + } + + #[tokio::test] + async fn execute_tool_call_with_default_output_propagates_fatal_errors() { + let err = execute_tool_call_with_default_output( + &TestExecutor, + TestCall { + call_id: "fatal".to_string(), + }, + CancellationToken::new(), + ) + .await + .expect_err("fatal tool error"); + + assert_eq!(err.to_string(), "Fatal error: boom"); + } +}