diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index 440d250b62..d8684648f8 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -5,15 +5,11 @@ use crate::model_family::ModelFamily; use crate::models::ContentItem; use crate::models::ResponseItem; use crate::openai_tools::OpenAiTool; -use crate::protocol::AskForApproval; -use crate::protocol::SandboxPolicy; use crate::protocol::TokenUsage; use codex_apply_patch::APPLY_PATCH_TOOL_INSTRUCTIONS; use futures::Stream; use serde::Serialize; use std::borrow::Cow; -use std::fmt::Display; -use std::path::PathBuf; use std::pin::Pin; use std::task::Context; use std::task::Poll; @@ -23,62 +19,19 @@ use tokio::sync::mpsc; /// with this content. const BASE_INSTRUCTIONS: &str = include_str!("../prompt.md"); -/// wraps environment context message in a tag for the model to parse more easily. -const ENVIRONMENT_CONTEXT_START: &str = "\n\n"; -const ENVIRONMENT_CONTEXT_END: &str = "\n\n"; - /// wraps user instructions message in a tag for the model to parse more easily. const USER_INSTRUCTIONS_START: &str = "\n\n"; const USER_INSTRUCTIONS_END: &str = "\n\n"; -#[derive(Debug, Clone)] -pub(crate) struct EnvironmentContext { - pub cwd: PathBuf, - pub approval_policy: AskForApproval, - pub sandbox_policy: SandboxPolicy, -} - -impl Display for EnvironmentContext { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - writeln!( - f, - "Current working directory: {}", - self.cwd.to_string_lossy() - )?; - writeln!(f, "Approval policy: {}", self.approval_policy)?; - writeln!(f, "Sandbox policy: {}", self.sandbox_policy)?; - - let network_access = match self.sandbox_policy.clone() { - SandboxPolicy::DangerFullAccess => "enabled", - SandboxPolicy::ReadOnly => "restricted", - SandboxPolicy::WorkspaceWrite { network_access, .. } => { - if network_access { - "enabled" - } else { - "restricted" - } - } - }; - writeln!(f, "Network access: {network_access}")?; - Ok(()) - } -} - -/// API request payload for a single model turn. +/// API request payload for a single model turn #[derive(Default, Debug, Clone)] pub struct Prompt { /// Conversation context input items. pub input: Vec, - /// Optional instructions from the user to amend to the built-in agent - /// instructions. - pub user_instructions: Option, + /// Whether to store response on server side (disable_response_storage = !store). pub store: bool, - /// A list of key-value pairs that will be added as a developer message - /// for the model to use - pub environment_context: Option, - /// Tools available to the model, including additional tools sourced from /// external MCP servers. pub tools: Vec, @@ -100,36 +53,19 @@ impl Prompt { Cow::Owned(sections.join("\n")) } - fn get_formatted_user_instructions(&self) -> Option { - self.user_instructions - .as_ref() - .map(|ui| format!("{USER_INSTRUCTIONS_START}{ui}{USER_INSTRUCTIONS_END}")) - } - - fn get_formatted_environment_context(&self) -> Option { - self.environment_context - .as_ref() - .map(|ec| format!("{ENVIRONMENT_CONTEXT_START}{ec}{ENVIRONMENT_CONTEXT_END}")) - } - pub(crate) fn get_formatted_input(&self) -> Vec { - let mut input_with_instructions = Vec::with_capacity(self.input.len() + 2); - if let Some(ec) = self.get_formatted_environment_context() { - input_with_instructions.push(ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ContentItem::InputText { text: ec }], - }); + self.input.clone() + } + + /// Creates a formatted user instructions message from a string + pub(crate) fn format_user_instructions_message(ui: &str) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: format!("{USER_INSTRUCTIONS_START}{ui}{USER_INSTRUCTIONS_END}"), + }], } - if let Some(ui) = self.get_formatted_user_instructions() { - input_with_instructions.push(ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ContentItem::InputText { text: ui }], - }); - } - input_with_instructions.extend(self.input.clone()); - input_with_instructions } } @@ -259,7 +195,6 @@ mod tests { #[test] fn get_full_instructions_no_user_content() { let prompt = Prompt { - user_instructions: Some("custom instruction".to_string()), ..Default::default() }; let expected = format!("{BASE_INSTRUCTIONS}\n{APPLY_PATCH_TOOL_INSTRUCTIONS}"); diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 2ae1db5bd5..098ff115fa 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -38,7 +38,6 @@ use crate::apply_patch::convert_apply_patch_to_protocol; use crate::apply_patch::get_writable_roots; use crate::apply_patch::{self}; use crate::client::ModelClient; -use crate::client_common::EnvironmentContext; use crate::client_common::Prompt; use crate::client_common::ResponseEvent; use crate::config::Config; @@ -46,6 +45,7 @@ use crate::config_types::ReasoningEffort as ReasoningEffortConfig; use crate::config_types::ReasoningSummary as ReasoningSummaryConfig; use crate::config_types::ShellEnvironmentPolicy; use crate::conversation_history::ConversationHistory; +use crate::environment_context::EnvironmentContext; use crate::error::CodexErr; use crate::error::Result as CodexResult; use crate::error::SandboxErr; @@ -137,9 +137,7 @@ impl Codex { let config = Arc::new(config); let resume_path = config.experimental_resume.clone(); - let session_id = Uuid::new_v4(); let configure_session = ConfigureSession { - session_id, provider: config.model_provider.clone(), model: config.model.clone(), model_reasoning_effort: config.model_reasoning_effort, @@ -161,6 +159,7 @@ impl Codex { error!("Failed to create session: {e:#}"); CodexErr::InternalAgentDied })?; + let session_id = session.session_id; // This task will run until Op::Shutdown is received. tokio::spawn(submission_loop(session, config, rx_sub)); @@ -253,8 +252,6 @@ pub(crate) struct Session { /// Configure the model session. struct ConfigureSession { - session_id: Uuid, - /// Provider identifier ("openai", "openrouter", ...). provider: ModelProviderInfo, @@ -302,7 +299,6 @@ impl Session { tx_event: Sender, ) -> anyhow::Result> { let ConfigureSession { - mut session_id, provider, model, model_reasoning_effort, @@ -324,53 +320,85 @@ impl Session { // Error messages to dispatch after SessionConfigured is sent. let mut post_session_configured_error_events = Vec::::new(); - // If `resume_path` is specified, then fail if we cannot resume the - // existing rollout. Though if `resume_path` is not specified and we - // fail to create a `RolloutRecorder`, potentially due to some sort of - // I/O error in attempting to create the log file, then we should still - // create the `Session`, but we do log an error. - let mut restored_items: Option> = None; - let rollout_recorder: Option = match resume_path.as_ref() { - Some(path) => match RolloutRecorder::resume(path, cwd.clone()).await { - Ok((rec, saved)) => { - session_id = saved.session_id; - if !saved.items.is_empty() { - restored_items = Some(saved.items); - } - Some(rec) - } - Err(e) => { - return Err(anyhow::anyhow!( - "failed to resume rollout from {path:?}: {e}" - )); - } - }, - None => { - match RolloutRecorder::new(&config, session_id, user_instructions.clone()).await { - Ok(r) => Some(r), - Err(e) => { - let message = format!("failed to initialise rollout recorder: {e}"); - post_session_configured_error_events.push(Event { - id: INITIAL_SUBMIT_ID.to_owned(), - msg: EventMsg::Error(ErrorEvent { - message: message.clone(), - }), - }); - warn!(message); - None - } + // Kick off independent async setup tasks in parallel to reduce startup latency. + // + // - initialize RolloutRecorder with new or resumed session info + // - spin up MCP connection manager + // - perform default shell discovery + // - load history metadata + let rollout_fut = async { + match resume_path.as_ref() { + Some(path) => RolloutRecorder::resume(path, cwd.clone()) + .await + .map(|(rec, saved)| (saved.session_id, Some(saved), rec)), + None => { + let session_id = Uuid::new_v4(); + RolloutRecorder::new(&config, session_id, user_instructions.clone()) + .await + .map(|rec| (session_id, None, rec)) } } }; - let client = ModelClient::new( - config.clone(), - auth.clone(), - provider.clone(), - model_reasoning_effort, - model_reasoning_summary, + let mcp_fut = McpConnectionManager::new(config.mcp_servers.clone()); + let default_shell_fut = shell::default_user_shell(); + let history_meta_fut = crate::message_history::history_metadata(&config); + + // Join all independent futures. + let (rollout_res, mcp_res, default_shell, (history_log_id, history_entry_count)) = + tokio::join!(rollout_fut, mcp_fut, default_shell_fut, history_meta_fut); + + // Handle rollout result, which determines the session_id. + struct RolloutResult { + session_id: Uuid, + rollout_recorder: Option, + restored_items: Option>, + } + let rollout_result = match rollout_res { + Ok((session_id, maybe_saved, recorder)) => { + let restored_items: Option> = + maybe_saved.and_then(|saved_session| { + if saved_session.items.is_empty() { + None + } else { + Some(saved_session.items) + } + }); + RolloutResult { + session_id, + rollout_recorder: Some(recorder), + restored_items, + } + } + Err(e) => { + if let Some(path) = resume_path.as_ref() { + return Err(anyhow::anyhow!( + "failed to resume rollout from {path:?}: {e}" + )); + } + + let message = format!("failed to initialize rollout recorder: {e}"); + post_session_configured_error_events.push(Event { + id: INITIAL_SUBMIT_ID.to_owned(), + msg: EventMsg::Error(ErrorEvent { + message: message.clone(), + }), + }); + warn!("{message}"); + + RolloutResult { + session_id: Uuid::new_v4(), + rollout_recorder: None, + restored_items: None, + } + } + }; + + let RolloutResult { session_id, - ); + rollout_recorder, + restored_items, + } = rollout_result; // Create the mutable state for the Session. let mut state = State { @@ -383,19 +411,19 @@ impl Session { let writable_roots = get_writable_roots(&cwd); - let (mcp_connection_manager, failed_clients) = - match McpConnectionManager::new(config.mcp_servers.clone()).await { - Ok((mgr, failures)) => (mgr, failures), - Err(e) => { - let message = format!("Failed to create MCP connection manager: {e:#}"); - error!("{message}"); - post_session_configured_error_events.push(Event { - id: INITIAL_SUBMIT_ID.to_owned(), - msg: EventMsg::Error(ErrorEvent { message }), - }); - (McpConnectionManager::default(), Default::default()) - } - }; + // Handle MCP manager result and record any startup failures. + let (mcp_connection_manager, failed_clients) = match mcp_res { + Ok((mgr, failures)) => (mgr, failures), + Err(e) => { + let message = format!("Failed to create MCP connection manager: {e:#}"); + error!("{message}"); + post_session_configured_error_events.push(Event { + id: INITIAL_SUBMIT_ID.to_owned(), + msg: EventMsg::Error(ErrorEvent { message }), + }); + (McpConnectionManager::default(), Default::default()) + } + }; // Surface individual client start-up failures to the user. if !failed_clients.is_empty() { @@ -409,7 +437,16 @@ impl Session { } } - let default_shell = shell::default_user_shell().await; + // Now that `session_id` is final (may have been updated by resume), + // construct the model client. + let client = ModelClient::new( + config.clone(), + auth.clone(), + provider.clone(), + model_reasoning_effort, + model_reasoning_summary, + session_id, + ); let sess = Arc::new(Session { session_id, client, @@ -437,11 +474,20 @@ impl Session { show_raw_agent_reasoning: config.show_raw_agent_reasoning, }); - // Gather history metadata for SessionConfiguredEvent. - let (history_log_id, history_entry_count) = - crate::message_history::history_metadata(&config).await; + // record the initial user instructions and environment context, + // regardless of whether we restored items. + let mut conversation_items = Vec::::with_capacity(2); + if let Some(user_instructions) = sess.user_instructions.as_deref() { + conversation_items.push(Prompt::format_user_instructions_message(user_instructions)); + } + conversation_items.push(ResponseItem::from(EnvironmentContext::new( + sess.get_cwd().to_path_buf(), + sess.get_approval_policy(), + sess.sandbox_policy.clone(), + ))); + sess.record_conversation_items(&conversation_items).await; - // ack + // Dispatch the SessionConfiguredEvent first and then report any errors. let events = std::iter::once(Event { id: INITIAL_SUBMIT_ID.to_owned(), msg: EventMsg::SessionConfigured(SessionConfiguredEvent { @@ -1237,15 +1283,9 @@ async fn run_turn( let prompt = Prompt { input, - user_instructions: sess.user_instructions.clone(), store: !sess.disable_response_storage, tools, base_instructions_override: sess.base_instructions.clone(), - environment_context: Some(EnvironmentContext { - cwd: sess.cwd.clone(), - approval_policy: sess.approval_policy, - sandbox_policy: sess.sandbox_policy.clone(), - }), }; let mut retries = 0; @@ -1483,9 +1523,7 @@ async fn run_compact_task( let prompt = Prompt { input: turn_input, - user_instructions: None, store: !sess.disable_response_storage, - environment_context: None, tools: Vec::new(), base_instructions_override: Some(compact_instructions.clone()), }; diff --git a/codex-rs/core/src/config_types.rs b/codex-rs/core/src/config_types.rs index 291dcb6422..cbbc6b4923 100644 --- a/codex-rs/core/src/config_types.rs +++ b/codex-rs/core/src/config_types.rs @@ -78,8 +78,9 @@ pub enum HistoryPersistence { #[derive(Deserialize, Debug, Clone, PartialEq, Default)] pub struct Tui {} -#[derive(Deserialize, Debug, Clone, Copy, PartialEq, Default, Serialize)] +#[derive(Deserialize, Debug, Clone, Copy, PartialEq, Default, Serialize, Display)] #[serde(rename_all = "kebab-case")] +#[strum(serialize_all = "kebab-case")] pub enum SandboxMode { #[serde(rename = "read-only")] #[default] diff --git a/codex-rs/core/src/environment_context.rs b/codex-rs/core/src/environment_context.rs new file mode 100644 index 0000000000..d31ddbc15e --- /dev/null +++ b/codex-rs/core/src/environment_context.rs @@ -0,0 +1,86 @@ +use serde::Deserialize; +use serde::Serialize; +use strum_macros::Display as DeriveDisplay; + +use crate::config_types::SandboxMode; +use crate::models::ContentItem; +use crate::models::ResponseItem; +use crate::protocol::AskForApproval; +use crate::protocol::SandboxPolicy; +use std::fmt::Display; +use std::path::PathBuf; + +/// wraps environment context message in a tag for the model to parse more easily. +pub(crate) const ENVIRONMENT_CONTEXT_START: &str = "\n"; +pub(crate) const ENVIRONMENT_CONTEXT_END: &str = ""; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, DeriveDisplay)] +#[serde(rename_all = "kebab-case")] +#[strum(serialize_all = "kebab-case")] +pub enum NetworkAccess { + Restricted, + Enabled, +} +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename = "environment_context", rename_all = "snake_case")] +pub(crate) struct EnvironmentContext { + pub cwd: PathBuf, + pub approval_policy: AskForApproval, + pub sandbox_mode: SandboxMode, + pub network_access: NetworkAccess, +} + +impl EnvironmentContext { + pub fn new( + cwd: PathBuf, + approval_policy: AskForApproval, + sandbox_policy: SandboxPolicy, + ) -> Self { + Self { + cwd, + approval_policy, + sandbox_mode: match sandbox_policy { + SandboxPolicy::DangerFullAccess => SandboxMode::DangerFullAccess, + SandboxPolicy::ReadOnly => SandboxMode::ReadOnly, + SandboxPolicy::WorkspaceWrite { .. } => SandboxMode::WorkspaceWrite, + }, + network_access: match sandbox_policy { + SandboxPolicy::DangerFullAccess => NetworkAccess::Enabled, + SandboxPolicy::ReadOnly => NetworkAccess::Restricted, + SandboxPolicy::WorkspaceWrite { network_access, .. } => { + if network_access { + NetworkAccess::Enabled + } else { + NetworkAccess::Restricted + } + } + }, + } + } +} + +impl Display for EnvironmentContext { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!( + f, + "Current working directory: {}", + self.cwd.to_string_lossy() + )?; + writeln!(f, "Approval policy: {}", self.approval_policy)?; + writeln!(f, "Sandbox mode: {}", self.sandbox_mode)?; + writeln!(f, "Network access: {}", self.network_access)?; + Ok(()) + } +} + +impl From for ResponseItem { + fn from(ec: EnvironmentContext) -> Self { + ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: format!("{ENVIRONMENT_CONTEXT_START}{ec}{ENVIRONMENT_CONTEXT_END}"), + }], + } + } +} diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index aeab49702c..d19fbbdb26 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -17,6 +17,7 @@ pub mod config; pub mod config_profile; pub mod config_types; mod conversation_history; +mod environment_context; pub mod error; pub mod exec; pub mod exec_env; diff --git a/codex-rs/core/tests/client.rs b/codex-rs/core/tests/client.rs index 1bcddf0796..10c6c66fb5 100644 --- a/codex-rs/core/tests/client.rs +++ b/codex-rs/core/tests/client.rs @@ -373,11 +373,11 @@ async fn includes_user_instructions_message_in_request() { .contains("be nice") ); assert_message_role(&request_body["input"][0], "user"); - assert_message_starts_with(&request_body["input"][0], "\n\n"); - assert_message_ends_with(&request_body["input"][0], ""); + assert_message_starts_with(&request_body["input"][0], ""); + assert_message_ends_with(&request_body["input"][0], ""); assert_message_role(&request_body["input"][1], "user"); - assert_message_starts_with(&request_body["input"][1], "\n\n"); - assert_message_ends_with(&request_body["input"][1], ""); + assert_message_starts_with(&request_body["input"][1], ""); + assert_message_ends_with(&request_body["input"][1], ""); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/codex-rs/core/tests/prompt_caching.rs b/codex-rs/core/tests/prompt_caching.rs index 8df7ea353d..0c2552ee05 100644 --- a/codex-rs/core/tests/prompt_caching.rs +++ b/codex-rs/core/tests/prompt_caching.rs @@ -85,7 +85,7 @@ async fn prefixes_context_and_instructions_once_and_consistently_across_requests assert_eq!(requests.len(), 2, "expected two POST requests"); let expected_env_text = format!( - "\n\nCurrent working directory: {}\nApproval policy: on-request\nSandbox policy: read-only\nNetwork access: restricted\n\n\n", + "\nCurrent working directory: {}\nApproval policy: on-request\nSandbox mode: read-only\nNetwork access: restricted\n", cwd.path().to_string_lossy() ); let expected_ui_text = @@ -113,7 +113,7 @@ async fn prefixes_context_and_instructions_once_and_consistently_across_requests let body1 = requests[0].body_json::().unwrap(); assert_eq!( body1["input"], - serde_json::json!([expected_env_msg, expected_ui_msg, expected_user_message_1]) + serde_json::json!([expected_ui_msg, expected_env_msg, expected_user_message_1]) ); let expected_user_message_2 = serde_json::json!({ diff --git a/codex-rs/tui/src/status_indicator_widget.rs b/codex-rs/tui/src/status_indicator_widget.rs index 65cba42523..1e23683171 100644 --- a/codex-rs/tui/src/status_indicator_widget.rs +++ b/codex-rs/tui/src/status_indicator_widget.rs @@ -217,17 +217,6 @@ impl WidgetRef for StatusIndicatorWidget { let mut spans: Vec> = Vec::new(); spans.push(Span::styled("▌ ", Style::default().fg(Color::Cyan))); - // Simple dim spinner to the left of the header. - let spinner_frames = ['·', '•', '●', '•']; - const SPINNER_SLOWDOWN: usize = 2; - let spinner_ch = spinner_frames[(idx / SPINNER_SLOWDOWN) % spinner_frames.len()]; - spans.push(Span::styled( - spinner_ch.to_string(), - Style::default().add_modifier(Modifier::DIM), - )); - spans.push(Span::raw(" ")); - - // Space after header // Animated header after the left bar spans.extend(animated_spans); // Space between header and bracket block @@ -336,7 +325,7 @@ mod tests { } #[test] - fn spinner_is_rendered() { + fn header_starts_at_expected_position() { let (tx_raw, _rx) = channel::(); let tx = AppEventSender::new(tx_raw); let mut w = StatusIndicatorWidget::new(tx); @@ -348,9 +337,6 @@ mod tests { w.render_ref(area, &mut buf); let ch = buf[(2, 0)].symbol().chars().next().unwrap_or(' '); - assert!( - matches!(ch, '·' | '•' | '●'), - "expected spinner char at col 2: {ch:?}" - ); + assert_eq!(ch, 'W', "expected Working header at col 2: {ch:?}"); } }