diff --git a/.github/workflows/issue-deduplicator.yml b/.github/workflows/issue-deduplicator.yml index f15c190102..fea6348c46 100644 --- a/.github/workflows/issue-deduplicator.yml +++ b/.github/workflows/issue-deduplicator.yml @@ -12,6 +12,7 @@ jobs: # Prevent runs on forks (requires OpenAI API key, wastes Actions minutes) if: github.repository == 'openai/codex' && (github.event.action == 'opened' || (github.event.action == 'labeled' && github.event.label.name == 'codex-deduplicate')) runs-on: ubuntu-latest + environment: issue-triage permissions: contents: read outputs: @@ -157,6 +158,7 @@ jobs: needs: normalize-duplicates-all if: ${{ needs.normalize-duplicates-all.result == 'success' && needs.normalize-duplicates-all.outputs.has_matches != 'true' }} runs-on: ubuntu-latest + environment: issue-triage permissions: contents: read outputs: diff --git a/.github/workflows/issue-labeler.yml b/.github/workflows/issue-labeler.yml index 77fe5d07c8..2c4eb6aa68 100644 --- a/.github/workflows/issue-labeler.yml +++ b/.github/workflows/issue-labeler.yml @@ -12,6 +12,7 @@ jobs: # Prevent runs on forks (requires OpenAI API key, wastes Actions minutes) if: github.repository == 'openai/codex' && (github.event.action == 'opened' || (github.event.action == 'labeled' && github.event.label.name == 'codex-label')) runs-on: ubuntu-latest + environment: issue-triage permissions: contents: read outputs: diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index ab279abe8d..2eeafe30b9 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2430,8 +2430,6 @@ dependencies = [ name = "codex-code-mode" version = "0.0.0" dependencies = [ - "async-channel", - "async-trait", "codex-protocol", "deno_core_icudata", "pretty_assertions", @@ -3033,7 +3031,6 @@ name = "codex-image-generation-extension" version = "0.0.0" dependencies = [ "async-trait", - "base64 0.22.1", "codex-api", "codex-core", "codex-extension-api", @@ -3048,9 +3045,6 @@ dependencies = [ "schemars 0.8.22", "serde", "serde_json", - "tempfile", - "tokio", - "tracing", ] [[package]] diff --git a/codex-rs/analytics/src/events.rs b/codex-rs/analytics/src/events.rs index 8791e5767a..3d31be12be 100644 --- a/codex-rs/analytics/src/events.rs +++ b/codex-rs/analytics/src/events.rs @@ -1069,12 +1069,9 @@ pub(crate) fn subagent_source_name(subagent_source: &SubAgentSource) -> String { } pub(crate) fn subagent_parent_thread_id(subagent_source: &SubAgentSource) -> Option { - match subagent_source { - SubAgentSource::ThreadSpawn { - parent_thread_id, .. - } => Some(parent_thread_id.to_string()), - _ => None, - } + subagent_source + .parent_thread_id() + .map(|parent_thread_id| parent_thread_id.to_string()) } fn analytics_hook_status(status: HookRunStatus) -> HookRunStatus { diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index c0d799dde4..1659fdd03d 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -2095,6 +2095,7 @@ mod tests { use codex_protocol::items::build_hook_prompt_message; use codex_protocol::models::FileSystemPermissions as CoreFileSystemPermissions; use codex_protocol::models::NetworkPermissions as CoreNetworkPermissions; + use codex_protocol::models::PermissionProfile; use codex_protocol::permissions::FileSystemAccessMode; use codex_protocol::permissions::FileSystemPath; use codex_protocol::permissions::FileSystemSandboxEntry; @@ -2110,7 +2111,6 @@ mod tests { use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::RateLimitWindow; use codex_protocol::protocol::RolloutItem; - use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; @@ -2188,7 +2188,7 @@ mod tests { agent_path: None, git_info: None, approval_mode: AskForApproval::OnRequest, - sandbox_policy: SandboxPolicy::new_read_only_policy(), + permission_profile: PermissionProfile::read_only(), token_usage: None, first_user_message: Some("before rollback".to_string()), history: Some(StoredThreadHistory { diff --git a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs index 8b23ee5183..beea7856a2 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs @@ -68,13 +68,13 @@ mod thread_processor_behavior_tests { use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_DANGER_FULL_ACCESS; use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_READ_ONLY; use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_WORKSPACE; + use codex_protocol::models::PermissionProfile; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::permissions::FileSystemAccessMode; use codex_protocol::permissions::FileSystemPath; use codex_protocol::permissions::FileSystemSandboxEntry; use codex_protocol::permissions::NetworkSandboxPolicy; use codex_protocol::protocol::AskForApproval; - use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_state::ThreadMetadataBuilder; @@ -413,7 +413,7 @@ mod thread_processor_behavior_tests { agent_path: None, git_info: None, approval_mode: AskForApproval::OnRequest, - sandbox_policy: SandboxPolicy::new_read_only_policy(), + permission_profile: PermissionProfile::read_only(), token_usage: None, first_user_message: Some("first user message".to_string()), history: None, diff --git a/codex-rs/app-server/tests/suite/v2/client_metadata.rs b/codex-rs/app-server/tests/suite/v2/client_metadata.rs index 687efefbdf..ab94fe618e 100644 --- a/codex-rs/app-server/tests/suite/v2/client_metadata.rs +++ b/codex-rs/app-server/tests/suite/v2/client_metadata.rs @@ -1,6 +1,7 @@ use anyhow::Result; use app_test_support::McpProcess; use app_test_support::create_fake_rollout; +use app_test_support::create_fake_rollout_with_source; use app_test_support::to_response; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; @@ -10,6 +11,8 @@ use codex_app_server_protocol::ReviewStartResponse; use codex_app_server_protocol::ReviewTarget; use codex_app_server_protocol::ThreadForkParams; use codex_app_server_protocol::ThreadForkResponse; +use codex_app_server_protocol::ThreadResumeParams; +use codex_app_server_protocol::ThreadResumeResponse; use codex_app_server_protocol::ThreadSource; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; @@ -18,6 +21,9 @@ use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::TurnSteerParams; use codex_app_server_protocol::TurnSteerResponse; use codex_app_server_protocol::UserInput as V2UserInput; +use codex_protocol::ThreadId as CoreThreadId; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; use core_test_support::responses; use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; @@ -290,6 +296,104 @@ async fn review_start_sends_fork_lineage_in_turn_metadata_for_thread_fork_v2() - Ok(()) } +#[tokio::test] +async fn turn_start_sends_subagent_lineage_after_cold_thread_resume_v2() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let response_mock = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_assistant_message("msg-1", "Done"), + responses::ev_completed("resp-1"), + ]), + ) + .await; + + let codex_home = TempDir::new()?; + create_config_toml( + codex_home.path(), + &server.uri(), + /*supports_websockets*/ false, + )?; + + let parent_thread_id = CoreThreadId::new(); + let parent_thread_id_str = parent_thread_id.to_string(); + let subagent_thread_id = create_fake_rollout_with_source( + codex_home.path(), + "2025-01-05T12-00-00", + "2025-01-05T12:00:00Z", + "Saved subagent message", + Some("mock_provider"), + /*git_info*/ None, + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + }), + )?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let resume_req = mcp + .send_thread_resume_request(ThreadResumeParams { + thread_id: subagent_thread_id.clone(), + ..Default::default() + }) + .await?; + let resume_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(resume_req)), + ) + .await??; + let ThreadResumeResponse { thread, .. } = to_response::(resume_resp)?; + assert_eq!(thread.id, subagent_thread_id); + + let turn_req = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id.clone(), + input: vec![V2UserInput::Text { + text: "Continue".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + let turn_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), + ) + .await??; + let TurnStartResponse { turn } = to_response::(turn_resp)?; + + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + + let request = response_mock.single_request(); + let metadata = request + .header("x-codex-turn-metadata") + .as_deref() + .map(parse_json_header) + .unwrap_or_else(|| panic!("missing x-codex-turn-metadata header")); + assert_eq!( + metadata["parent_thread_id"].as_str(), + Some(parent_thread_id_str.as_str()) + ); + assert_eq!(metadata["subagent_kind"].as_str(), Some("thread_spawn")); + assert_eq!(metadata["thread_id"].as_str(), Some(thread.id.as_str())); + assert_eq!(metadata["turn_id"].as_str(), Some(turn.id.as_str())); + assert!(metadata.get("forked_from_thread_id").is_none()); + + Ok(()) +} + #[tokio::test] async fn turn_steer_updates_client_metadata_on_follow_up_responses_request_v2() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/app-server/tests/suite/v2/web_search.rs b/codex-rs/app-server/tests/suite/v2/web_search.rs index ef8d562703..0d2d995328 100644 --- a/codex-rs/app-server/tests/suite/v2/web_search.rs +++ b/codex-rs/app-server/tests/suite/v2/web_search.rs @@ -140,6 +140,7 @@ async fn standalone_web_search_round_trips_encrypted_output() -> Result<()> { ); let search_body = search_request_body(&server).await?; + assert_eq!(search_body["model"], json!("mock-model")); assert_eq!( search_body["commands"], json!({ diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 8dfe6dc33f..e5660301b5 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -50,6 +50,8 @@ mod marketplace_cmd; mod mcp_cmd; mod plugin_cmd; mod remote_control_cmd; +#[cfg(target_os = "windows")] +mod sandbox_setup; mod state_db_recovery; #[cfg(not(windows))] mod wsl_paths; @@ -1250,6 +1252,16 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { .await?; } Some(Subcommand::Sandbox(mut sandbox_cli)) => { + #[cfg(target_os = "windows")] + if let Some(setup_cli) = sandbox_setup::parse_setup_command(&sandbox_cli.command)? { + reject_remote_mode_for_subcommand( + root_remote.as_deref(), + root_remote_auth_token_env.as_deref(), + "sandbox setup", + )?; + sandbox_setup::run(setup_cli).await?; + return Ok(()); + } reject_remote_mode_for_subcommand( root_remote.as_deref(), root_remote_auth_token_env.as_deref(), diff --git a/codex-rs/cli/src/sandbox_setup.rs b/codex-rs/cli/src/sandbox_setup.rs new file mode 100644 index 0000000000..958411c428 --- /dev/null +++ b/codex-rs/cli/src/sandbox_setup.rs @@ -0,0 +1,206 @@ +use std::path::PathBuf; + +use clap::ArgAction; +use clap::ArgGroup; +use clap::Parser; +use codex_core::config::edit::ConfigEditsBuilder; +use codex_core::config::find_codex_home; + +#[derive(Debug, Parser)] +#[command(group( + ArgGroup::new("sandbox_user") + .required(true) + .args(["user", "current_user"]) +))] +pub(crate) struct SandboxSetupCommand { + /// Set up the elevated Windows sandbox. + #[arg(long = "elevated", action = ArgAction::SetTrue)] + elevated_sandbox_level: bool, + + /// Windows user that will run Codex after managed deployment. + #[arg( + long = "user", + value_name = "USER", + conflicts_with = "current_user", + requires = "codex_home" + )] + user: Option, + + /// Use the current Windows user as the Codex user. + #[arg( + long = "current-user", + default_value_t = false, + conflicts_with = "user" + )] + current_user: bool, + + /// CODEX_HOME for the Codex user. Required with --user. + #[arg(long = "codex-home", value_name = "DIR")] + codex_home: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SandboxSetupLevel { + Elevated, +} + +impl SandboxSetupCommand { + fn setup_level(&self) -> anyhow::Result { + if self.elevated_sandbox_level { + Ok(SandboxSetupLevel::Elevated) + } else { + anyhow::bail!("`codex sandbox setup` currently requires --elevated"); + } + } +} + +pub(crate) async fn run(cmd: SandboxSetupCommand) -> anyhow::Result<()> { + match cmd.setup_level()? { + SandboxSetupLevel::Elevated => run_elevated(cmd).await, + } +} + +pub(crate) fn parse_setup_command( + sandbox_command: &[String], +) -> anyhow::Result> { + if sandbox_command + .first() + .is_none_or(|command| command != "setup") + { + return Ok(None); + } + + SandboxSetupCommand::try_parse_from(sandbox_command.iter().map(String::as_str)) + .map(Some) + .map_err(anyhow::Error::from) +} + +async fn run_elevated(cmd: SandboxSetupCommand) -> anyhow::Result<()> { + let identity = resolve_sandbox_setup_identity(&cmd)?; + + codex_core::windows_sandbox::run_elevated_provisioning_setup( + identity.codex_home.as_path(), + identity.real_user.as_str(), + )?; + ConfigEditsBuilder::new(identity.codex_home.as_path()) + .set_windows_sandbox_mode("elevated") + .apply() + .await + .map_err(|err| { + anyhow::anyhow!( + "sandbox provisioning succeeded, but failed to persist elevated sandbox config: {err}" + ) + })?; + + println!( + "Windows elevated sandbox setup completed for {} at {}.", + identity.real_user, + identity.codex_home.display() + ); + Ok(()) +} + +struct SandboxSetupIdentity { + real_user: String, + codex_home: PathBuf, +} + +fn resolve_sandbox_setup_identity( + cmd: &SandboxSetupCommand, +) -> anyhow::Result { + if cmd.current_user { + let real_user = std::env::var("USERNAME") + .or_else(|_| std::env::var("USER")) + .map_err(|err| { + anyhow::anyhow!("failed to determine current user from environment: {err}") + })?; + let codex_home = match cmd.codex_home.clone() { + Some(codex_home) => codex_home, + None => find_codex_home()?.to_path_buf(), + }; + return Ok(SandboxSetupIdentity { + real_user, + codex_home, + }); + } + + let real_user = cmd + .user + .clone() + .ok_or_else(|| anyhow::anyhow!("--user or --current-user is required"))?; + let codex_home = cmd + .codex_home + .clone() + .ok_or_else(|| anyhow::anyhow!("--codex-home is required with --user"))?; + Ok(SandboxSetupIdentity { + real_user, + codex_home, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_managed_user_identity() { + let command = SandboxSetupCommand::try_parse_from([ + "setup", + "--elevated", + "--user", + "DOMAIN\\alice", + "--codex-home", + r"C:\Users\alice\.codex", + ]) + .expect("parse"); + + assert!(command.elevated_sandbox_level); + assert_eq!(command.user.as_deref(), Some(r"DOMAIN\alice")); + assert!(!command.current_user); + assert_eq!( + command.codex_home.as_deref(), + Some(std::path::Path::new(r"C:\Users\alice\.codex")) + ); + } + + #[test] + fn requires_explicit_user_identity() { + let err = SandboxSetupCommand::try_parse_from(["setup", "--elevated"]) + .expect_err("parse should fail"); + + assert_eq!(err.kind(), clap::error::ErrorKind::MissingRequiredArgument); + } + + #[test] + fn requires_codex_home_for_managed_user() { + let err = + SandboxSetupCommand::try_parse_from(["setup", "--elevated", "--user", "DOMAIN\\alice"]) + .expect_err("parse should fail"); + + assert_eq!(err.kind(), clap::error::ErrorKind::MissingRequiredArgument); + } + + #[test] + fn parses_setup_from_sandbox_command_args() { + let command = parse_setup_command(&[ + "setup".to_string(), + "--elevated".to_string(), + "--user".to_string(), + r"DOMAIN\alice".to_string(), + "--codex-home".to_string(), + r"C:\Users\alice\.codex".to_string(), + ]) + .expect("parse") + .expect("setup command"); + + assert_eq!(command.user.as_deref(), Some(r"DOMAIN\alice")); + } + + #[test] + fn ignores_non_setup_sandbox_command_args() { + let command = + parse_setup_command(&["echo".to_string(), "hello".to_string()]).expect("parse"); + + assert!(command.is_none()); + } +} diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml index 19d6c3ab45..879404d8aa 100644 --- a/codex-rs/code-mode/Cargo.toml +++ b/codex-rs/code-mode/Cargo.toml @@ -16,8 +16,6 @@ sandbox = ["v8/v8_enable_sandbox"] workspace = true [dependencies] -async-channel = { workspace = true } -async-trait = { workspace = true } codex-protocol = { workspace = true } deno_core_icudata = { workspace = true } serde = { workspace = true, features = ["derive"] } diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index 3da8c77325..37ba1d2a91 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -29,9 +29,18 @@ pub use runtime::WaitOutcome; pub use runtime::WaitRequest; pub use runtime::WaitToPendingOutcome; pub use runtime::WaitToPendingRequest; +pub use service::CellId; pub use service::CodeModeService; -pub use service::CodeModeTurnHost; -pub use service::CodeModeTurnWorker; +pub use service::CodeModeSession; +pub use service::CodeModeSessionDelegate; +pub use service::CodeModeSessionProvider; +pub use service::CodeModeSessionProviderFuture; +pub use service::CodeModeSessionResultFuture; +pub use service::InProcessCodeModeSessionProvider; +pub use service::NoopCodeModeSessionDelegate; +pub use service::NotificationFuture; +pub use service::StartedCell; +pub use service::ToolInvocationFuture; pub const PUBLIC_TOOL_NAME: &str = "exec"; pub const WAIT_TOOL_NAME: &str = "wait"; diff --git a/codex-rs/code-mode/src/runtime/mod.rs b/codex-rs/code-mode/src/runtime/mod.rs index 89893d2b25..21757328ee 100644 --- a/codex-rs/code-mode/src/runtime/mod.rs +++ b/codex-rs/code-mode/src/runtime/mod.rs @@ -19,6 +19,7 @@ use crate::description::EnabledToolMetadata; use crate::description::ToolDefinition; use crate::description::enabled_tool_metadata; use crate::response::FunctionCallOutputContentItem; +use crate::service::CellId; pub const DEFAULT_EXEC_YIELD_TIME_MS: u64 = 10_000; pub const DEFAULT_WAIT_YIELD_TIME_MS: u64 = 10_000; @@ -27,11 +28,6 @@ const EXIT_SENTINEL: &str = "__codex_code_mode_exit__"; #[derive(Clone, Debug)] pub struct ExecuteRequest { - /// Runtime cell id for this execution. - /// - /// Callers allocate this before execution so tracing, waits, and nested tool - /// calls can refer to the cell as soon as JavaScript starts. - pub cell_id: String, pub tool_call_id: String, pub enabled_tools: Vec, pub source: String, @@ -41,14 +37,13 @@ pub struct ExecuteRequest { #[derive(Clone, Debug)] pub struct WaitRequest { - pub cell_id: String, + pub cell_id: CellId, pub yield_time_ms: u64, - pub terminate: bool, } #[derive(Clone, Debug)] pub struct WaitToPendingRequest { - pub cell_id: String, + pub cell_id: CellId, } /// Result of waiting on a code-mode cell. @@ -73,7 +68,7 @@ pub enum ExecuteToPendingOutcome { /// The cell is waiting for more runtime input after draining the runtime /// input queue that was ready at the pending boundary. Pending { - cell_id: String, + cell_id: CellId, content_items: Vec, /// Runtime tool-call ids emitted before this paused execution frontier /// sealed. Hosts can use these ids to drain their tool-call transport @@ -105,15 +100,15 @@ impl From for RuntimeResponse { #[derive(Debug, PartialEq, Serialize)] pub enum RuntimeResponse { Yielded { - cell_id: String, + cell_id: CellId, content_items: Vec, }, Terminated { - cell_id: String, + cell_id: CellId, content_items: Vec, }, Result { - cell_id: String, + cell_id: CellId, content_items: Vec, error_text: Option, }, @@ -126,23 +121,13 @@ pub enum RuntimeResponse { /// if their tool-call graph requires globally unique ids. #[derive(Debug)] pub struct CodeModeNestedToolCall { - pub cell_id: String, + pub cell_id: CellId, pub runtime_tool_call_id: String, pub tool_name: ToolName, pub tool_kind: CodeModeToolKind, pub input: Option, } -#[derive(Debug)] -pub(crate) enum TurnMessage { - ToolCall(CodeModeNestedToolCall), - Notify { - cell_id: String, - call_id: String, - text: String, - }, -} - #[derive(Debug)] pub(crate) enum RuntimeCommand { ToolResponse { id: String, result: JsonValue }, @@ -460,7 +445,6 @@ mod tests { fn execute_request(source: &str) -> ExecuteRequest { ExecuteRequest { - cell_id: "1".to_string(), tool_call_id: "call_1".to_string(), enabled_tools: Vec::new(), source: source.to_string(), diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index d153f1b416..6589e6ec5a 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -1,14 +1,20 @@ use std::collections::HashMap; +use std::fmt; +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use std::time::Duration; -use async_trait::async_trait; +use serde::Deserialize; +use serde::Serialize; use serde_json::Value as JsonValue; use tokio::sync::Mutex; use tokio::sync::mpsc; use tokio::sync::oneshot; +use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; use tracing::warn; @@ -22,35 +28,162 @@ use crate::runtime::RuntimeCommand; use crate::runtime::RuntimeControlCommand; use crate::runtime::RuntimeEvent; use crate::runtime::RuntimeResponse; -use crate::runtime::TurnMessage; use crate::runtime::WaitOutcome; use crate::runtime::WaitRequest; use crate::runtime::WaitToPendingOutcome; use crate::runtime::WaitToPendingRequest; use crate::runtime::spawn_runtime; -#[async_trait] -pub trait CodeModeTurnHost: Send + Sync { - async fn invoke_tool( - &self, +pub type CodeModeSessionResultFuture<'a, T> = + Pin> + Send + 'a>>; +pub type CodeModeSessionProviderFuture<'a> = + CodeModeSessionResultFuture<'a, Arc>; +pub type ToolInvocationFuture<'a> = + Pin> + Send + 'a>>; +pub type NotificationFuture<'a> = Pin> + Send + 'a>>; + +#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)] +pub struct CellId(String); + +impl CellId { + pub fn new(value: String) -> Self { + Self(value) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl AsRef for CellId { + fn as_ref(&self) -> &str { + self.as_str() + } +} + +impl fmt::Display for CellId { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(self.as_str()) + } +} + +pub struct StartedCell { + pub cell_id: CellId, + initial_response_rx: oneshot::Receiver, +} + +impl StartedCell { + pub async fn initial_response(self) -> Result { + self.initial_response_rx + .await + .map_err(|_| "exec runtime ended unexpectedly".to_string()) + } +} + +/// Host callbacks used by a code-mode session while cells are executing. +pub trait CodeModeSessionDelegate: Send + Sync { + fn invoke_tool<'a>( + &'a self, invocation: CodeModeNestedToolCall, cancellation_token: CancellationToken, - ) -> Result; + ) -> ToolInvocationFuture<'a>; - async fn notify(&self, call_id: String, cell_id: String, text: String) -> Result<(), String>; + fn notify<'a>( + &'a self, + call_id: String, + cell_id: CellId, + text: String, + cancellation_token: CancellationToken, + ) -> NotificationFuture<'a>; + + /// Releases delegate state associated with a cell after it reaches a terminal state. + fn cell_closed(&self, cell_id: &CellId); +} + +pub struct NoopCodeModeSessionDelegate; + +impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate { + fn invoke_tool<'a>( + &'a self, + _invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async move { + cancellation_token.cancelled().await; + Err("code mode nested tools are unavailable".to_string()) + }) + } + + fn notify<'a>( + &'a self, + _call_id: String, + _cell_id: CellId, + _text: String, + _cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async { Ok(()) }) + } + + fn cell_closed(&self, _cell_id: &CellId) {} +} + +/// A durable code-mode session owned by one Codex thread. +/// +/// Cells executed in the same session share stored values. Separate sessions +/// must keep those values isolated. Implementations may execute cells +/// in-process or remotely. +pub trait CodeModeSession: Send + Sync { + fn execute<'a>( + &'a self, + request: ExecuteRequest, + ) -> CodeModeSessionResultFuture<'a, StartedCell>; + + fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome>; + + fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome>; + + fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()>; +} + +/// Creates code-mode sessions for one Codex thread. +/// +/// Providers choose where a session executes and receive the host delegate that +/// the session should use for nested tool calls and notifications. +pub trait CodeModeSessionProvider: Send + Sync { + fn create_session<'a>( + &'a self, + delegate: Arc, + ) -> CodeModeSessionProviderFuture<'a>; +} + +#[derive(Default)] +pub struct InProcessCodeModeSessionProvider; + +impl CodeModeSessionProvider for InProcessCodeModeSessionProvider { + fn create_session<'a>( + &'a self, + delegate: Arc, + ) -> CodeModeSessionProviderFuture<'a> { + Box::pin(async move { + let session: Arc = + Arc::new(CodeModeService::with_delegate(delegate)); + Ok(session) + }) + } } #[derive(Clone)] -struct SessionHandle { - control_tx: mpsc::UnboundedSender, +struct CellHandle { + control_tx: mpsc::UnboundedSender, runtime_tx: std::sync::mpsc::Sender, + cancellation_token: CancellationToken, } struct Inner { stored_values: Mutex>, - sessions: Mutex>, - turn_message_tx: async_channel::Sender, - turn_message_rx: async_channel::Receiver, + cells: Mutex>, + delegate: Arc, + shutting_down: AtomicBool, next_cell_id: AtomicU64, } @@ -60,50 +193,50 @@ pub struct CodeModeService { impl CodeModeService { pub fn new() -> Self { - let (turn_message_tx, turn_message_rx) = async_channel::unbounded(); + Self::with_delegate(Arc::new(NoopCodeModeSessionDelegate)) + } + pub fn with_delegate(delegate: Arc) -> Self { Self { inner: Arc::new(Inner { stored_values: Mutex::new(HashMap::new()), - sessions: Mutex::new(HashMap::new()), - turn_message_tx, - turn_message_rx, + cells: Mutex::new(HashMap::new()), + delegate, + shutting_down: AtomicBool::new(false), next_cell_id: AtomicU64::new(1), }), } } - async fn stored_values(&self) -> HashMap { - self.inner.stored_values.lock().await.clone() + fn allocate_cell_id(&self) -> CellId { + CellId::new( + self.inner + .next_cell_id + .fetch_add(1, Ordering::Relaxed) + .to_string(), + ) } - /// Reserves the runtime cell id for a future `execute` request. - /// - /// The runtime can issue nested tool calls before the first `execute` - /// response is returned. Hosts that need a parent trace object for those - /// nested calls should allocate the cell id up front and pass it back on the - /// `ExecuteRequest`. - pub fn allocate_cell_id(&self) -> String { - self.inner - .next_cell_id - .fetch_add(1, Ordering::Relaxed) - .to_string() - } - - pub async fn execute(&self, request: ExecuteRequest) -> Result { + pub async fn execute(&self, request: ExecuteRequest) -> Result { + if self.inner.shutting_down.load(Ordering::Acquire) { + return Err("code mode session is shutting down".to_string()); + } let initial_yield_time_ms = request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS); let (response_tx, response_rx) = oneshot::channel(); - self.start_session( + let cell_id = self.allocate_cell_id(); + self.start_cell( + cell_id.clone(), request, - SessionResponseSender::Runtime(response_tx), + CellResponseSender::Runtime(response_tx), Some(initial_yield_time_ms), PendingRuntimeMode::Continue, ) .await?; - response_rx - .await - .map_err(|_| "exec runtime ended unexpectedly".to_string()) + Ok(StartedCell { + cell_id, + initial_response_rx: response_rx, + }) } pub async fn execute_to_pending( @@ -111,9 +244,11 @@ impl CodeModeService { request: ExecuteRequest, ) -> Result { let (response_tx, response_rx) = oneshot::channel(); - self.start_session( + let cell_id = self.allocate_cell_id(); + self.start_cell( + cell_id, request, - SessionResponseSender::ExecuteToPending(response_tx), + CellResponseSender::ExecuteToPending(response_tx), /*initial_yield_time_ms*/ None, PendingRuntimeMode::PauseUntilResumed, ) @@ -124,47 +259,50 @@ impl CodeModeService { .map_err(|_| "exec runtime ended unexpectedly".to_string()) } - async fn start_session( + async fn start_cell( &self, + cell_id: CellId, request: ExecuteRequest, - initial_response_tx: SessionResponseSender, + initial_response_tx: CellResponseSender, initial_yield_time_ms: Option, pending_mode: PendingRuntimeMode, ) -> Result<(), String> { - let cell_id = request.cell_id.clone(); let (event_tx, event_rx) = mpsc::unbounded_channel(); let (control_tx, control_rx) = mpsc::unbounded_channel(); - let stored_values = self.stored_values().await; + let stored_values = self.inner.stored_values.lock().await.clone(); + let cancellation_token = CancellationToken::new(); let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = { - let mut sessions = self.inner.sessions.lock().await; - if sessions.contains_key(&cell_id) { + let mut cells = self.inner.cells.lock().await; + if self.inner.shutting_down.load(Ordering::Acquire) { + return Err("code mode session is shutting down".to_string()); + } + if cells.contains_key(&cell_id) { return Err(format!("exec cell {cell_id} already exists")); } let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = spawn_runtime(stored_values, request, event_tx, pending_mode)?; - // Keep the session registry locked through insertion so a - // caller-owned cell id cannot race with another execute and replace - // a live runtime. - sessions.insert( + cells.insert( cell_id.clone(), - SessionHandle { + CellHandle { control_tx, runtime_tx: runtime_tx.clone(), + cancellation_token: cancellation_token.clone(), }, ); (runtime_tx, runtime_control_tx, runtime_terminate_handle) }; - tokio::spawn(run_session_control( + tokio::spawn(run_cell_control( Arc::clone(&self.inner), - SessionControlContext { - cell_id: cell_id.clone(), + CellControlContext { + cell_id, runtime_tx, runtime_control_tx, pending_mode, runtime_terminate_handle, + cancellation_token, }, event_rx, control_rx, @@ -176,34 +314,44 @@ impl CodeModeService { } pub async fn wait(&self, request: WaitRequest) -> Result { - let cell_id = request.cell_id.clone(); - let handle = self - .inner - .sessions - .lock() - .await - .get(&request.cell_id) - .cloned(); + let WaitRequest { + cell_id, + yield_time_ms, + } = request; + let handle = self.inner.cells.lock().await.get(&cell_id).cloned(); let Some(handle) = handle else { return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); }; let (response_tx, response_rx) = oneshot::channel(); - let control_message = if request.terminate { - SessionControlCommand::Terminate { response_tx } - } else { - SessionControlCommand::Poll { - yield_time_ms: request.yield_time_ms, - response_tx, - } + let control_message = CellControlCommand::Poll { + yield_time_ms, + response_tx, }; if handle.control_tx.send(control_message).is_err() { return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); } match response_rx.await { Ok(response) => Ok(WaitOutcome::LiveCell(response)), - Err(_) => Ok(WaitOutcome::MissingCell(missing_cell_response( - request.cell_id, - ))), + Err(_) => Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))), + } + } + + pub async fn terminate(&self, cell_id: CellId) -> Result { + let handle = self.inner.cells.lock().await.get(&cell_id).cloned(); + let Some(handle) = handle else { + return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); + }; + let (response_tx, response_rx) = oneshot::channel(); + if handle + .control_tx + .send(CellControlCommand::Terminate { response_tx }) + .is_err() + { + return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); + } + match response_rx.await { + Ok(response) => Ok(WaitOutcome::LiveCell(response)), + Err(_) => Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))), } } @@ -211,14 +359,8 @@ impl CodeModeService { &self, request: WaitToPendingRequest, ) -> Result { - let cell_id = request.cell_id.clone(); - let handle = self - .inner - .sessions - .lock() - .await - .get(&request.cell_id) - .cloned(); + let cell_id = request.cell_id; + let handle = self.inner.cells.lock().await.get(&cell_id).cloned(); let Some(handle) = handle else { return Ok(WaitToPendingOutcome::MissingCell(missing_cell_response( cell_id, @@ -227,7 +369,7 @@ impl CodeModeService { let (response_tx, response_rx) = oneshot::channel(); if handle .control_tx - .send(SessionControlCommand::PollToPending { response_tx }) + .send(CellControlCommand::PollToPending { response_tx }) .is_err() { return Ok(WaitToPendingOutcome::MissingCell(missing_cell_response( @@ -237,74 +379,33 @@ impl CodeModeService { match response_rx.await { Ok(response) => Ok(WaitToPendingOutcome::LiveCell(response)), Err(_) => Ok(WaitToPendingOutcome::MissingCell(missing_cell_response( - request.cell_id, + cell_id, ))), } } - pub fn start_turn_worker(&self, host: Arc) -> CodeModeTurnWorker { - let (shutdown_tx, mut shutdown_rx) = oneshot::channel(); - let inner = Arc::clone(&self.inner); - let turn_message_rx = self.inner.turn_message_rx.clone(); - - tokio::spawn(async move { - loop { - let next_message = tokio::select! { - _ = &mut shutdown_rx => break, - message = turn_message_rx.recv() => message.ok(), - }; - let Some(next_message) = next_message else { - break; - }; - match next_message { - TurnMessage::Notify { - cell_id, - call_id, - text, - } => { - if let Err(err) = host.notify(call_id, cell_id.clone(), text).await { - warn!( - "failed to deliver code mode notification for cell {cell_id}: {err}" - ); - } - } - TurnMessage::ToolCall(invocation) => { - let host = Arc::clone(&host); - let inner = Arc::clone(&inner); - tokio::spawn(async move { - let cell_id = invocation.cell_id.clone(); - let runtime_tool_call_id = invocation.runtime_tool_call_id.clone(); - let response = - host.invoke_tool(invocation, CancellationToken::new()).await; - let runtime_tx = inner - .sessions - .lock() - .await - .get(&cell_id) - .map(|handle| handle.runtime_tx.clone()); - let Some(runtime_tx) = runtime_tx else { - return; - }; - let command = match response { - Ok(result) => RuntimeCommand::ToolResponse { - id: runtime_tool_call_id, - result, - }, - Err(error_text) => RuntimeCommand::ToolError { - id: runtime_tool_call_id, - error_text, - }, - }; - let _ = runtime_tx.send(command); - }); - } - } - } - }); - - CodeModeTurnWorker { - shutdown_tx: Some(shutdown_tx), + pub async fn shutdown(&self) -> Result<(), String> { + self.inner.shutting_down.store(true, Ordering::Release); + let handles = self + .inner + .cells + .lock() + .await + .values() + .cloned() + .collect::>(); + for handle in handles { + handle.cancellation_token.cancel(); + let (response_tx, _response_rx) = oneshot::channel(); + let _ = handle + .control_tx + .send(CellControlCommand::Terminate { response_tx }); + let _ = handle.runtime_tx.send(RuntimeCommand::Terminate); } + while !self.inner.cells.lock().await.is_empty() { + tokio::task::yield_now().await; + } + Ok(()) } } @@ -314,19 +415,44 @@ impl Default for CodeModeService { } } -pub struct CodeModeTurnWorker { - shutdown_tx: Option>, -} - -impl Drop for CodeModeTurnWorker { +impl Drop for CodeModeService { fn drop(&mut self) { - if let Some(shutdown_tx) = self.shutdown_tx.take() { - let _ = shutdown_tx.send(()); + self.inner.shutting_down.store(true, Ordering::Release); + if let Ok(cells) = self.inner.cells.try_lock() { + for handle in cells.values() { + handle.cancellation_token.cancel(); + let (response_tx, _response_rx) = oneshot::channel(); + let _ = handle + .control_tx + .send(CellControlCommand::Terminate { response_tx }); + let _ = handle.runtime_tx.send(RuntimeCommand::Terminate); + } } } } -enum SessionControlCommand { +impl CodeModeSession for CodeModeService { + fn execute<'a>( + &'a self, + request: ExecuteRequest, + ) -> CodeModeSessionResultFuture<'a, StartedCell> { + Box::pin(CodeModeService::execute(self, request)) + } + + fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(CodeModeService::wait(self, request)) + } + + fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(CodeModeService::terminate(self, cell_id)) + } + + fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> { + Box::pin(CodeModeService::shutdown(self)) + } +} + +enum CellControlCommand { Poll { yield_time_ms: u64, response_tx: oneshot::Sender, @@ -339,7 +465,7 @@ enum SessionControlCommand { }, } -enum SessionResponseSender { +enum CellResponseSender { Runtime(oneshot::Sender), ExecuteToPending(oneshot::Sender), } @@ -349,15 +475,16 @@ struct PendingResult { error_text: Option, } -struct SessionControlContext { - cell_id: String, +struct CellControlContext { + cell_id: CellId, runtime_tx: std::sync::mpsc::Sender, runtime_control_tx: std::sync::mpsc::Sender, pending_mode: PendingRuntimeMode, runtime_terminate_handle: v8::IsolateHandle, + cancellation_token: CancellationToken, } -fn missing_cell_response(cell_id: String) -> RuntimeResponse { +fn missing_cell_response(cell_id: CellId) -> RuntimeResponse { RuntimeResponse::Result { error_text: Some(format!("exec cell {cell_id} not found")), cell_id, @@ -365,29 +492,29 @@ fn missing_cell_response(cell_id: String) -> RuntimeResponse { } } -fn pending_result_response(cell_id: &str, result: PendingResult) -> RuntimeResponse { +fn pending_result_response(cell_id: &CellId, result: PendingResult) -> RuntimeResponse { RuntimeResponse::Result { - cell_id: cell_id.to_string(), + cell_id: cell_id.clone(), content_items: result.content_items, error_text: result.error_text, } } -fn send_terminal_response(response_tx: SessionResponseSender, response: RuntimeResponse) { +fn send_terminal_response(response_tx: CellResponseSender, response: RuntimeResponse) { match response_tx { - SessionResponseSender::Runtime(response_tx) => { + CellResponseSender::Runtime(response_tx) => { let _ = response_tx.send(response); } - SessionResponseSender::ExecuteToPending(response_tx) => { + CellResponseSender::ExecuteToPending(response_tx) => { let _ = response_tx.send(ExecuteToPendingOutcome::Completed(response)); } } } fn send_or_buffer_result( - cell_id: &str, + cell_id: &CellId, result: PendingResult, - response_tx: &mut Option, + response_tx: &mut Option, pending_result: &mut Option, ) -> bool { if let Some(response_tx) = response_tx.take() { @@ -401,42 +528,41 @@ fn send_or_buffer_result( } fn send_yield_response( - cell_id: &str, + cell_id: &CellId, content_items: &mut Vec, - response_tx: &mut Option, + response_tx: &mut Option, ) { let Some(current_response_tx) = response_tx.take() else { return; }; match current_response_tx { - SessionResponseSender::Runtime(response_tx) => { + CellResponseSender::Runtime(response_tx) => { let _ = response_tx.send(RuntimeResponse::Yielded { - cell_id: cell_id.to_string(), + cell_id: cell_id.clone(), content_items: std::mem::take(content_items), }); } - SessionResponseSender::ExecuteToPending(execute_to_pending_tx) => { - *response_tx = Some(SessionResponseSender::ExecuteToPending( - execute_to_pending_tx, - )); + CellResponseSender::ExecuteToPending(execute_to_pending_tx) => { + *response_tx = Some(CellResponseSender::ExecuteToPending(execute_to_pending_tx)); } } } -async fn run_session_control( +async fn run_cell_control( inner: Arc, - context: SessionControlContext, + context: CellControlContext, mut event_rx: mpsc::UnboundedReceiver, - mut control_rx: mpsc::UnboundedReceiver, - initial_response_tx: SessionResponseSender, + mut control_rx: mpsc::UnboundedReceiver, + initial_response_tx: CellResponseSender, initial_yield_time_ms: Option, ) { - let SessionControlContext { + let CellControlContext { cell_id, runtime_tx, runtime_control_tx, pending_mode, runtime_terminate_handle, + cancellation_token, } = context; let mut content_items = Vec::new(); let mut pending_tool_call_ids = Vec::new(); @@ -445,6 +571,7 @@ async fn run_session_control( let mut termination_requested = false; let mut runtime_closed = false; let mut yield_timer: Option>> = None; + let mut notification_tasks = JoinSet::new(); loop { tokio::select! { @@ -492,11 +619,11 @@ async fn run_session_control( RuntimeEvent::Pending => { if let Some(current_response_tx) = response_tx.take() { match current_response_tx { - SessionResponseSender::Runtime(runtime_response_tx) => { + CellResponseSender::Runtime(runtime_response_tx) => { response_tx = - Some(SessionResponseSender::Runtime(runtime_response_tx)); + Some(CellResponseSender::Runtime(runtime_response_tx)); } - SessionResponseSender::ExecuteToPending(response_tx) => { + CellResponseSender::ExecuteToPending(response_tx) => { let _ = response_tx.send(ExecuteToPendingOutcome::Pending { cell_id: cell_id.clone(), content_items: std::mem::take(&mut content_items), @@ -516,11 +643,26 @@ async fn run_session_control( send_yield_response(&cell_id, &mut content_items, &mut response_tx); } RuntimeEvent::Notify { call_id, text } => { - let _ = inner.turn_message_tx.send(TurnMessage::Notify { - cell_id: cell_id.clone(), - call_id, - text, - }).await; + let delegate = Arc::clone(&inner.delegate); + let cell_id = cell_id.clone(); + let cancellation_token = cancellation_token.child_token(); + notification_tasks.spawn(async move { + tokio::select! { + result = delegate.notify( + call_id, + cell_id.clone(), + text, + cancellation_token.clone(), + ) => { + if let Err(err) = result { + warn!( + "failed to deliver code mode notification for cell {cell_id}: {err}" + ); + } + } + _ = cancellation_token.cancelled() => {} + } + }); } RuntimeEvent::ToolCall { id, @@ -533,15 +675,25 @@ async fn run_session_control( } let tool_call = CodeModeNestedToolCall { cell_id: cell_id.clone(), - runtime_tool_call_id: id, + runtime_tool_call_id: id.clone(), tool_name: name, tool_kind: kind, input, }; - let _ = inner - .turn_message_tx - .send(TurnMessage::ToolCall(tool_call)) - .await; + let delegate = Arc::clone(&inner.delegate); + let runtime_tx = runtime_tx.clone(); + let cancellation_token = cancellation_token.child_token(); + tokio::spawn(async move { + let response = tokio::select! { + response = delegate.invoke_tool(tool_call, cancellation_token.clone()) => response, + _ = cancellation_token.cancelled() => return, + }; + let command = match response { + Ok(result) => RuntimeCommand::ToolResponse { id, result }, + Err(error_text) => RuntimeCommand::ToolError { id, error_text }, + }; + let _ = runtime_tx.send(command); + }); } RuntimeEvent::Result { stored_value_writes, @@ -558,6 +710,7 @@ async fn run_session_control( } break; } + drain_notification_tasks(&mut notification_tasks).await; inner .stored_values .lock() @@ -578,12 +731,19 @@ async fn run_session_control( } } } + task_result = notification_tasks.join_next(), if !notification_tasks.is_empty() => { + if let Some(Err(err)) = task_result + && !err.is_cancelled() + { + warn!("code mode notification task failed: {err}"); + } + } maybe_command = control_rx.recv() => { let Some(command) = maybe_command else { break; }; match command { - SessionControlCommand::Poll { + CellControlCommand::Poll { yield_time_ms, response_tx: next_response_tx, } => { @@ -591,11 +751,11 @@ async fn run_session_control( let _ = next_response_tx.send(pending_result_response(&cell_id, result)); break; } - response_tx = Some(SessionResponseSender::Runtime(next_response_tx)); + response_tx = Some(CellResponseSender::Runtime(next_response_tx)); yield_timer = Some(Box::pin(tokio::time::sleep(Duration::from_millis(yield_time_ms)))); resume_paused_runtime(&runtime_control_tx, pending_mode); } - SessionControlCommand::PollToPending { + CellControlCommand::PollToPending { response_tx: next_response_tx, } => { if let Some(result) = pending_result.take() { @@ -605,18 +765,19 @@ async fn run_session_control( break; } response_tx = - Some(SessionResponseSender::ExecuteToPending(next_response_tx)); + Some(CellResponseSender::ExecuteToPending(next_response_tx)); yield_timer = None; resume_paused_runtime(&runtime_control_tx, pending_mode); } - SessionControlCommand::Terminate { response_tx: next_response_tx } => { + CellControlCommand::Terminate { response_tx: next_response_tx } => { if let Some(result) = pending_result.take() { let _ = next_response_tx.send(pending_result_response(&cell_id, result)); break; } - response_tx = Some(SessionResponseSender::Runtime(next_response_tx)); + response_tx = Some(CellResponseSender::Runtime(next_response_tx)); termination_requested = true; + cancellation_token.cancel(); yield_timer = None; let _ = runtime_tx.send(RuntimeCommand::Terminate); terminate_paused_runtime(&runtime_control_tx, pending_mode); @@ -650,8 +811,21 @@ async fn run_session_control( } let _ = runtime_tx.send(RuntimeCommand::Terminate); + cancellation_token.cancel(); + drain_notification_tasks(&mut notification_tasks).await; terminate_paused_runtime(&runtime_control_tx, pending_mode); - inner.sessions.lock().await.remove(&cell_id); + inner.cells.lock().await.remove(&cell_id); + inner.delegate.cell_closed(&cell_id); +} + +async fn drain_notification_tasks(notification_tasks: &mut JoinSet<()>) { + while let Some(result) = notification_tasks.join_next().await { + if let Err(err) = result + && !err.is_cancelled() + { + warn!("code mode notification task failed: {err}"); + } + } } fn resume_paused_runtime( @@ -677,6 +851,7 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::AtomicU64; + use std::sync::atomic::Ordering; use std::time::Duration; use codex_protocol::ToolName; @@ -685,19 +860,21 @@ mod tests { use tokio::sync::mpsc; use tokio::sync::oneshot; + use super::CellControlCommand; + use super::CellControlContext; + use super::CellId; + use super::CellResponseSender; use super::CodeModeService; use super::Inner; + use super::NoopCodeModeSessionDelegate; use super::PendingRuntimeMode; use super::RuntimeCommand; use super::RuntimeResponse; - use super::SessionControlCommand; - use super::SessionControlContext; - use super::SessionResponseSender; use super::WaitOutcome; use super::WaitRequest; use super::WaitToPendingOutcome; use super::WaitToPendingRequest; - use super::run_session_control; + use super::run_cell_control; use crate::CodeModeToolKind; use crate::FunctionCallOutputContentItem; use crate::ToolDefinition; @@ -708,7 +885,6 @@ mod tests { fn execute_request(source: &str) -> ExecuteRequest { ExecuteRequest { - cell_id: "1".to_string(), tool_call_id: "call_1".to_string(), enabled_tools: Vec::new(), source: source.to_string(), @@ -717,13 +893,26 @@ mod tests { } } + fn cell_id(value: &str) -> CellId { + CellId::new(value.to_string()) + } + + async fn execute(service: &CodeModeService, request: ExecuteRequest) -> RuntimeResponse { + service + .execute(request) + .await + .unwrap() + .initial_response() + .await + .unwrap() + } + fn test_inner() -> Arc { - let (turn_message_tx, turn_message_rx) = async_channel::unbounded(); Arc::new(Inner { stored_values: Mutex::new(HashMap::new()), - sessions: Mutex::new(HashMap::new()), - turn_message_tx, - turn_message_rx, + cells: Mutex::new(HashMap::new()), + delegate: Arc::new(NoopCodeModeSessionDelegate), + shutting_down: std::sync::atomic::AtomicBool::new(false), next_cell_id: AtomicU64::new(1), }) } @@ -732,19 +921,20 @@ mod tests { async fn synchronous_exit_returns_successfully() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#"text("before"); exit(); text("after");"#.to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "before".to_string(), }], @@ -753,6 +943,116 @@ mod tests { ); } + #[tokio::test] + async fn stored_values_are_shared_between_cells_but_not_sessions() { + let first_session = CodeModeService::new(); + let second_session = CodeModeService::new(); + + let write_response = execute( + &first_session, + ExecuteRequest { + source: r#"store("key", "visible");"#.to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + let same_session = execute( + &first_session, + ExecuteRequest { + source: r#"text(String(load("key")));"#.to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + let other_session = execute( + &second_session, + ExecuteRequest { + source: r#"text(String(load("key")));"#.to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + assert_eq!( + write_response, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: Vec::new(), + error_text: None, + } + ); + assert_eq!( + same_session, + RuntimeResponse::Result { + cell_id: cell_id("2"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "visible".to_string(), + }], + error_text: None, + } + ); + assert_eq!( + other_session, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "undefined".to_string(), + }], + error_text: None, + } + ); + } + + #[tokio::test] + async fn shutdown_interrupts_cpu_bound_cells() { + let service = CodeModeService::new(); + + let cell = service + .execute(ExecuteRequest { + source: "while (true) {}".to_string(), + ..execute_request("") + }) + .await + .unwrap(); + assert_eq!( + cell.initial_response().await.unwrap(), + RuntimeResponse::Yielded { + cell_id: cell_id("1"), + content_items: Vec::new(), + } + ); + + tokio::time::timeout(Duration::from_secs(1), service.shutdown()) + .await + .unwrap() + .unwrap(); + } + + #[tokio::test] + async fn start_cell_rejects_new_cell_after_shutdown_begins() { + let service = CodeModeService::new(); + service.inner.shutting_down.store(true, Ordering::Release); + let (response_tx, _response_rx) = oneshot::channel(); + + let error = service + .start_cell( + cell_id("late-cell"), + execute_request(""), + CellResponseSender::Runtime(response_tx), + Some(/*initial_yield_time_ms*/ 1), + PendingRuntimeMode::Continue, + ) + .await + .unwrap_err(); + + assert_eq!(error, "code mode session is shutting down".to_string()); + assert!(service.inner.cells.lock().await.is_empty()); + } + #[tokio::test] async fn execute_to_pending_returns_completed_for_synchronous_results() { let service = CodeModeService::new(); @@ -769,7 +1069,7 @@ mod tests { assert_eq!( response, ExecuteToPendingOutcome::Completed(RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), }], @@ -797,7 +1097,7 @@ mod tests { assert_eq!( response, ExecuteToPendingOutcome::Pending { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "before".to_string(), }], @@ -805,19 +1105,12 @@ mod tests { } ); - let termination = service - .wait(WaitRequest { - cell_id: "1".to_string(), - yield_time_ms: 1, - terminate: true, - }) - .await - .unwrap(); + let termination = service.terminate(cell_id("1")).await.unwrap(); assert_eq!( termination, WaitOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), }) ); @@ -853,25 +1146,18 @@ await Promise.all([ assert_eq!( response, ExecuteToPendingOutcome::Pending { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), pending_tool_call_ids: vec!["tool-1".to_string(), "tool-2".to_string()], } ); - let termination = service - .wait(WaitRequest { - cell_id: "1".to_string(), - yield_time_ms: 1, - terminate: true, - }) - .await - .unwrap(); + let termination = service.terminate(cell_id("1")).await.unwrap(); assert_eq!( termination, WaitOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), }) ); @@ -910,7 +1196,7 @@ await Promise.all([ assert_eq!( initial_response, ExecuteToPendingOutcome::Pending { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), pending_tool_call_ids: vec!["tool-1".to_string(), "tool-2".to_string()], } @@ -918,10 +1204,10 @@ await Promise.all([ let runtime_tx = service .inner - .sessions + .cells .lock() .await - .get("1") + .get(&cell_id("1")) .unwrap() .runtime_tx .clone(); @@ -932,7 +1218,7 @@ await Promise.all([ let resumed_response = tokio::time::timeout( Duration::from_secs(1), service.wait_to_pending(WaitToPendingRequest { - cell_id: "1".to_string(), + cell_id: cell_id("1"), }), ) .await @@ -942,25 +1228,18 @@ await Promise.all([ assert_eq!( resumed_response, WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Pending { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), pending_tool_call_ids: vec!["tool-3".to_string()], }) ); - let termination = service - .wait(WaitRequest { - cell_id: "1".to_string(), - yield_time_ms: 1, - terminate: true, - }) - .await - .unwrap(); + let termination = service.terminate(cell_id("1")).await.unwrap(); assert_eq!( termination, WaitOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), }) ); @@ -987,7 +1266,7 @@ await new Promise(() => {}); assert_eq!( initial_response, ExecuteToPendingOutcome::Pending { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), pending_tool_call_ids: Vec::new(), } @@ -995,10 +1274,10 @@ await new Promise(() => {}); let runtime_tx = service .inner - .sessions + .cells .lock() .await - .get("1") + .get(&cell_id("1")) .unwrap() .runtime_tx .clone(); @@ -1009,7 +1288,7 @@ await new Promise(() => {}); let resumed_response = tokio::time::timeout( Duration::from_secs(1), service.wait_to_pending(WaitToPendingRequest { - cell_id: "1".to_string(), + cell_id: cell_id("1"), }), ) .await @@ -1019,7 +1298,7 @@ await new Promise(() => {}); assert_eq!( resumed_response, WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Pending { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "after".to_string(), }], @@ -1027,19 +1306,12 @@ await new Promise(() => {}); }) ); - let termination = service - .wait(WaitRequest { - cell_id: "1".to_string(), - yield_time_ms: 1, - terminate: true, - }) - .await - .unwrap(); + let termination = service.terminate(cell_id("1")).await.unwrap(); assert_eq!( termination, WaitOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), }) ); @@ -1065,7 +1337,7 @@ text("done"); assert_eq!( initial_response, ExecuteToPendingOutcome::Pending { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), pending_tool_call_ids: Vec::new(), } @@ -1073,10 +1345,10 @@ text("done"); let runtime_tx = service .inner - .sessions + .cells .lock() .await - .get("1") + .get(&cell_id("1")) .unwrap() .runtime_tx .clone(); @@ -1087,7 +1359,7 @@ text("done"); let resumed_response = tokio::time::timeout( Duration::from_secs(1), service.wait_to_pending(WaitToPendingRequest { - cell_id: "1".to_string(), + cell_id: cell_id("1"), }), ) .await @@ -1098,7 +1370,7 @@ text("done"); resumed_response, WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Completed( RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), }], @@ -1112,19 +1384,20 @@ text("done"); async fn v8_console_is_not_exposed_on_global_this() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#"text(String(Object.hasOwn(globalThis, "console")));"#.to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "false".to_string(), }], @@ -1137,8 +1410,9 @@ text("done"); async fn date_locale_string_formats_with_icu_data() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" const value = new Date("2025-01-02T03:04:05Z") .toLocaleString("fr-FR", { @@ -1156,14 +1430,14 @@ text(value); .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "jeudi 2 janvier \u{e0} 03:04:05".to_string(), }], @@ -1176,8 +1450,9 @@ text(value); async fn intl_date_time_format_formats_with_icu_data() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" const formatter = new Intl.DateTimeFormat("fr-FR", { weekday: "long", @@ -1194,14 +1469,14 @@ text(formatter.format(new Date("2025-01-02T03:04:05Z"))); .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "jeudi 2 janvier \u{e0} 03:04:05".to_string(), }], @@ -1214,8 +1489,9 @@ text(formatter.format(new Date("2025-01-02T03:04:05Z"))); async fn output_helpers_return_undefined() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" const returnsUndefined = [ text("first"), @@ -1227,14 +1503,14 @@ text(JSON.stringify(returnsUndefined)); .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![ FunctionCallOutputContentItem::InputText { text: "first".to_string(), @@ -1256,8 +1532,9 @@ text(JSON.stringify(returnsUndefined)); async fn image_helper_accepts_raw_mcp_image_block_with_original_detail() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" image({ type: "image", @@ -1269,14 +1546,14 @@ image({ .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputImage { image_url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==".to_string(), detail: Some(crate::ImageDetail::Original), @@ -1290,8 +1567,9 @@ image({ async fn image_helper_second_arg_overrides_explicit_object_detail() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" image( { @@ -1304,14 +1582,14 @@ image( .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputImage { image_url: "https://example.com/image.jpg".to_string(), detail: Some(crate::ImageDetail::Original), @@ -1325,8 +1603,9 @@ image( async fn image_helper_second_arg_overrides_raw_mcp_image_detail() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" image( { @@ -1341,14 +1620,14 @@ image( .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputImage { image_url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==".to_string(), detail: Some(crate::ImageDetail::High), @@ -1362,8 +1641,9 @@ image( async fn image_helper_accepts_low_detail() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" image({ image_url: "https://example.com/image.jpg", @@ -1373,14 +1653,14 @@ image({ .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputImage { image_url: "https://example.com/image.jpg".to_string(), detail: Some(crate::ImageDetail::Low), @@ -1394,8 +1674,9 @@ image({ async fn image_helper_rejects_unsupported_detail() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" image({ image_url: "https://example.com/image.jpg", @@ -1405,14 +1686,14 @@ image({ .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), error_text: Some( "image detail must be one of: auto, low, high, original".to_string() @@ -1425,8 +1706,9 @@ image({ async fn image_helper_rejects_raw_mcp_result_container() { let service = CodeModeService::new(); - let response = service - .execute(ExecuteRequest { + let response = execute( + &service, + ExecuteRequest { source: r#" image({ content: [ @@ -1443,14 +1725,14 @@ image({ .to_string(), yield_time_ms: None, ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await; assert_eq!( response, RuntimeResponse::Result { - cell_id: "1".to_string(), + cell_id: cell_id("1"), content_items: Vec::new(), error_text: Some( "image expects a non-empty image URL string, an object with image_url and optional detail, or a raw MCP image block".to_string(), @@ -1465,9 +1747,8 @@ image({ let response = service .wait(WaitRequest { - cell_id: "missing".to_string(), + cell_id: cell_id("missing"), yield_time_ms: 1, - terminate: false, }) .await .unwrap(); @@ -1475,7 +1756,7 @@ image({ assert_eq!( response, WaitOutcome::MissingCell(RuntimeResponse::Result { - cell_id: "missing".to_string(), + cell_id: cell_id("missing"), content_items: Vec::new(), error_text: Some("exec cell missing not found".to_string()), }) @@ -1501,18 +1782,19 @@ image({ ) .unwrap(); - tokio::spawn(run_session_control( + tokio::spawn(run_cell_control( inner, - SessionControlContext { - cell_id: "cell-1".to_string(), + CellControlContext { + cell_id: cell_id("cell-1"), runtime_tx: runtime_tx.clone(), runtime_control_tx, pending_mode: PendingRuntimeMode::Continue, runtime_terminate_handle, + cancellation_token: tokio_util::sync::CancellationToken::new(), }, event_rx, control_rx, - SessionResponseSender::Runtime(initial_response_tx), + CellResponseSender::Runtime(initial_response_tx), Some(/*initial_yield_time_ms*/ 60_000), )); @@ -1521,14 +1803,14 @@ image({ assert_eq!( initial_response_rx.await.unwrap(), RuntimeResponse::Yielded { - cell_id: "cell-1".to_string(), + cell_id: cell_id("cell-1"), content_items: Vec::new(), } ); let (terminate_response_tx, terminate_response_rx) = oneshot::channel(); control_tx - .send(SessionControlCommand::Terminate { + .send(CellControlCommand::Terminate { response_tx: terminate_response_tx, }) .unwrap(); @@ -1545,7 +1827,7 @@ image({ assert_eq!( terminate_response.await, RuntimeResponse::Terminated { - cell_id: "cell-1".to_string(), + cell_id: cell_id("cell-1"), content_items: Vec::new(), } ); diff --git a/codex-rs/codex-api/src/endpoint/search.rs b/codex-rs/codex-api/src/endpoint/search.rs index 2940f872d8..d01fbfb780 100644 --- a/codex-rs/codex-api/src/endpoint/search.rs +++ b/codex-rs/codex-api/src/endpoint/search.rs @@ -145,7 +145,7 @@ mod tests { .search( &SearchRequest { id: "search-session".to_string(), - model: Some("gpt-test".to_string()), + model: "gpt-test".to_string(), reasoning: None, input: Some(SearchInput::Items(vec![ResponseItem::Message { id: None, diff --git a/codex-rs/codex-api/src/search.rs b/codex-rs/codex-api/src/search.rs index 8328415f83..061b3ac8c6 100644 --- a/codex-rs/codex-api/src/search.rs +++ b/codex-rs/codex-api/src/search.rs @@ -7,8 +7,7 @@ use serde::Serialize; #[derive(Debug, Clone, Serialize, PartialEq)] pub struct SearchRequest { pub id: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub model: Option, + pub model: String, #[serde(skip_serializing_if = "Option::is_none")] pub reasoning: Option, #[serde(skip_serializing_if = "Option::is_none")] diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 608594873d..5bc99ae072 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -1247,12 +1247,7 @@ impl AgentControl { } fn thread_spawn_parent_thread_id(session_source: &SessionSource) -> Option { - match session_source { - SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, .. - }) => Some(*parent_thread_id), - _ => None, - } + session_source.parent_thread_id() } fn agent_matches_prefix(agent_path: Option<&AgentPath>, prefix: &AgentPath) -> bool { diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index f47be864fb..620645d8f6 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1734,19 +1734,9 @@ fn subagent_header_value(session_source: &SessionSource) -> Option { } fn parent_thread_id_header_value(session_source: &SessionSource) -> Option { - match session_source { - SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, .. - }) => Some(parent_thread_id.to_string()), - SessionSource::Cli - | SessionSource::VSCode - | SessionSource::Exec - | SessionSource::Mcp - | SessionSource::Custom(_) - | SessionSource::Internal(_) - | SessionSource::SubAgent(_) - | SessionSource::Unknown => None, - } + session_source + .parent_thread_id() + .map(|parent_thread_id| parent_thread_id.to_string()) } const RESPONSE_STREAM_CHANNEL_CAPACITY: usize = 1600; diff --git a/codex-rs/core/src/config/edit.rs b/codex-rs/core/src/config/edit.rs index 290edf6fb8..862660ba2e 100644 --- a/codex-rs/core/src/config/edit.rs +++ b/codex-rs/core/src/config/edit.rs @@ -24,6 +24,8 @@ use toml_edit::value; const NOTICE_TABLE_KEY: &str = "notice"; +mod document_helpers; + /// Discrete config mutations supported by the persistence engine. #[derive(Clone, Debug)] pub enum ConfigEdit { @@ -198,319 +200,6 @@ pub fn model_availability_nux_count_edits(shown_count: &HashMap) -> edits } -// TODO(jif) move to a dedicated file -mod document_helpers { - use codex_config::types::AppToolApproval; - use codex_config::types::McpServerConfig; - use codex_config::types::McpServerEnvVar; - use codex_config::types::McpServerToolConfig; - use codex_config::types::McpServerTransportConfig; - use codex_config::types::ToolSuggestDisabledTool; - use codex_config::types::ToolSuggestDiscoverableType; - use toml_edit::Array as TomlArray; - use toml_edit::InlineTable; - use toml_edit::Item as TomlItem; - use toml_edit::Table as TomlTable; - use toml_edit::Value as TomlValue; - use toml_edit::value; - - pub(super) fn ensure_table_for_write(item: &mut TomlItem) -> Option<&mut TomlTable> { - match item { - TomlItem::Table(table) => Some(table), - TomlItem::Value(value) => { - if let Some(inline) = value.as_inline_table() { - *item = TomlItem::Table(table_from_inline(inline)); - item.as_table_mut() - } else { - *item = TomlItem::Table(new_implicit_table()); - item.as_table_mut() - } - } - TomlItem::None => { - *item = TomlItem::Table(new_implicit_table()); - item.as_table_mut() - } - _ => None, - } - } - - pub(super) fn ensure_table_for_read(item: &mut TomlItem) -> Option<&mut TomlTable> { - match item { - TomlItem::Table(table) => Some(table), - TomlItem::Value(value) => { - let inline = value.as_inline_table()?; - *item = TomlItem::Table(table_from_inline(inline)); - item.as_table_mut() - } - _ => None, - } - } - - fn serialize_mcp_server_table(config: &McpServerConfig) -> TomlTable { - let mut entry = TomlTable::new(); - entry.set_implicit(false); - - match &config.transport { - McpServerTransportConfig::Stdio { - command, - args, - env, - env_vars, - cwd, - } => { - entry["command"] = value(command.clone()); - if !args.is_empty() { - entry["args"] = array_from_iter(args.iter().cloned()); - } - if let Some(env) = env - && !env.is_empty() - { - entry["env"] = table_from_pairs(env.iter()); - } - if !env_vars.is_empty() { - entry["env_vars"] = array_from_env_vars(env_vars); - } - if let Some(cwd) = cwd { - entry["cwd"] = value(cwd.to_string_lossy().to_string()); - } - } - McpServerTransportConfig::StreamableHttp { - url, - bearer_token_env_var, - http_headers, - env_http_headers, - } => { - entry["url"] = value(url.clone()); - if let Some(env_var) = bearer_token_env_var { - entry["bearer_token_env_var"] = value(env_var.clone()); - } - if let Some(headers) = http_headers - && !headers.is_empty() - { - entry["http_headers"] = table_from_pairs(headers.iter()); - } - if let Some(headers) = env_http_headers - && !headers.is_empty() - { - entry["env_http_headers"] = table_from_pairs(headers.iter()); - } - } - } - - if !config.enabled { - entry["enabled"] = value(false); - } - if !config.is_local_environment() { - entry["environment_id"] = value(config.environment_id.clone()); - } - if config.required { - entry["required"] = value(true); - } - if config.supports_parallel_tool_calls { - entry["supports_parallel_tool_calls"] = value(true); - } - if let Some(timeout) = config.startup_timeout_sec { - entry["startup_timeout_sec"] = value(timeout.as_secs_f64()); - } - if let Some(timeout) = config.tool_timeout_sec { - entry["tool_timeout_sec"] = value(timeout.as_secs_f64()); - } - if let Some(approval_mode) = config.default_tools_approval_mode { - entry["default_tools_approval_mode"] = value(match approval_mode { - AppToolApproval::Auto => "auto", - AppToolApproval::Prompt => "prompt", - AppToolApproval::Approve => "approve", - }); - } - if let Some(enabled_tools) = &config.enabled_tools - && !enabled_tools.is_empty() - { - entry["enabled_tools"] = array_from_iter(enabled_tools.iter().cloned()); - } - if let Some(disabled_tools) = &config.disabled_tools - && !disabled_tools.is_empty() - { - entry["disabled_tools"] = array_from_iter(disabled_tools.iter().cloned()); - } - if let Some(scopes) = &config.scopes - && !scopes.is_empty() - { - entry["scopes"] = array_from_iter(scopes.iter().cloned()); - } - if let Some(oauth) = &config.oauth - && let Some(client_id) = &oauth.client_id - && !client_id.is_empty() - { - let mut oauth_table = TomlTable::new(); - oauth_table.set_implicit(false); - oauth_table["client_id"] = value(client_id.clone()); - entry["oauth"] = TomlItem::Table(oauth_table); - } - if let Some(resource) = &config.oauth_resource - && !resource.is_empty() - { - entry["oauth_resource"] = value(resource.clone()); - } - if !config.tools.is_empty() { - let mut tools = new_implicit_table(); - let mut tool_entries: Vec<_> = config.tools.iter().collect(); - tool_entries.sort_by_key(|(name, _)| *name); - for (name, tool_config) in tool_entries { - tools.insert(name, serialize_mcp_server_tool(tool_config)); - } - entry.insert("tools", TomlItem::Table(tools)); - } - - entry - } - - fn serialize_mcp_server_tool(config: &McpServerToolConfig) -> TomlItem { - let mut entry = TomlTable::new(); - entry.set_implicit(false); - if let Some(approval_mode) = config.approval_mode { - entry["approval_mode"] = value(match approval_mode { - AppToolApproval::Auto => "auto", - AppToolApproval::Prompt => "prompt", - AppToolApproval::Approve => "approve", - }); - } - TomlItem::Table(entry) - } - - pub(super) fn serialize_mcp_server(config: &McpServerConfig) -> TomlItem { - TomlItem::Table(serialize_mcp_server_table(config)) - } - - pub(super) fn serialize_mcp_server_inline(config: &McpServerConfig) -> InlineTable { - serialize_mcp_server_table(config).into_inline_table() - } - - pub(super) fn merge_inline_table(existing: &mut InlineTable, replacement: InlineTable) { - existing.retain(|key, _| replacement.get(key).is_some()); - - for (key, value) in replacement.iter() { - if let Some(existing_value) = existing.get_mut(key) { - let mut updated_value = value.clone(); - *updated_value.decor_mut() = existing_value.decor().clone(); - *existing_value = updated_value; - } else { - existing.insert(key.to_string(), value.clone()); - } - } - } - - fn table_from_inline(inline: &InlineTable) -> TomlTable { - let mut table = new_implicit_table(); - for (key, value) in inline.iter() { - let mut value = value.clone(); - let decor = value.decor_mut(); - decor.set_suffix(""); - table.insert(key, TomlItem::Value(value)); - } - table - } - - pub(super) fn new_implicit_table() -> TomlTable { - let mut table = TomlTable::new(); - table.set_implicit(true); - table - } - - pub(super) fn parse_tool_suggest_disabled_tool( - value: &TomlValue, - ) -> Option { - let table = value.as_inline_table()?; - let kind = match table.get("type").and_then(TomlValue::as_str) { - Some("connector") => ToolSuggestDiscoverableType::Connector, - Some("plugin") => ToolSuggestDiscoverableType::Plugin, - _ => return None, - }; - let id = table.get("id").and_then(TomlValue::as_str)?; - Some(ToolSuggestDisabledTool { - kind, - id: id.to_string(), - }) - } - - pub(super) fn parse_tool_suggest_disabled_tool_table( - table: &TomlTable, - ) -> Option { - let kind = match table.get("type").and_then(TomlItem::as_str) { - Some("connector") => ToolSuggestDiscoverableType::Connector, - Some("plugin") => ToolSuggestDiscoverableType::Plugin, - _ => return None, - }; - let id = table.get("id").and_then(TomlItem::as_str)?; - Some(ToolSuggestDisabledTool { - kind, - id: id.to_string(), - }) - } - - pub(super) fn tool_suggest_disabled_tools_value( - disabled_tools: &[ToolSuggestDisabledTool], - ) -> TomlItem { - let mut array = TomlArray::new(); - for disabled_tool in disabled_tools { - let mut table = InlineTable::new(); - table.insert( - "type", - match disabled_tool.kind { - ToolSuggestDiscoverableType::Connector => "connector", - ToolSuggestDiscoverableType::Plugin => "plugin", - } - .into(), - ); - table.insert("id", disabled_tool.id.clone().into()); - array.push(table); - } - TomlItem::Value(array.into()) - } - - fn array_from_iter(iter: I) -> TomlItem - where - I: Iterator, - { - let mut array = TomlArray::new(); - for value in iter { - array.push(value); - } - TomlItem::Value(array.into()) - } - - fn array_from_env_vars(env_vars: &[McpServerEnvVar]) -> TomlItem { - let mut array = TomlArray::new(); - for env_var in env_vars { - match env_var { - McpServerEnvVar::Name(name) => array.push(name.clone()), - McpServerEnvVar::Config { name, source } => { - let mut table = InlineTable::new(); - table.insert("name", name.clone().into()); - if let Some(source) = source { - table.insert("source", source.clone().into()); - } - array.push(table); - } - } - } - TomlItem::Value(array.into()) - } - - fn table_from_pairs<'a, I>(pairs: I) -> TomlItem - where - I: IntoIterator, - { - let mut entries: Vec<_> = pairs.into_iter().collect(); - entries.sort_by_key(|(key, _)| *key); - let mut table = TomlTable::new(); - table.set_implicit(false); - for (key, val) in entries { - table.insert(key, value(val.clone())); - } - TomlItem::Table(table) - } -} - struct ConfigDocument { doc: DocumentMut, } diff --git a/codex-rs/core/src/config/edit/document_helpers.rs b/codex-rs/core/src/config/edit/document_helpers.rs new file mode 100644 index 0000000000..5b99306eb9 --- /dev/null +++ b/codex-rs/core/src/config/edit/document_helpers.rs @@ -0,0 +1,309 @@ +use codex_config::types::AppToolApproval; +use codex_config::types::McpServerConfig; +use codex_config::types::McpServerEnvVar; +use codex_config::types::McpServerToolConfig; +use codex_config::types::McpServerTransportConfig; +use codex_config::types::ToolSuggestDisabledTool; +use codex_config::types::ToolSuggestDiscoverableType; +use toml_edit::Array as TomlArray; +use toml_edit::InlineTable; +use toml_edit::Item as TomlItem; +use toml_edit::Table as TomlTable; +use toml_edit::Value as TomlValue; +use toml_edit::value; + +pub(super) fn ensure_table_for_write(item: &mut TomlItem) -> Option<&mut TomlTable> { + match item { + TomlItem::Table(table) => Some(table), + TomlItem::Value(value) => { + if let Some(inline) = value.as_inline_table() { + *item = TomlItem::Table(table_from_inline(inline)); + item.as_table_mut() + } else { + *item = TomlItem::Table(new_implicit_table()); + item.as_table_mut() + } + } + TomlItem::None => { + *item = TomlItem::Table(new_implicit_table()); + item.as_table_mut() + } + _ => None, + } +} + +pub(super) fn ensure_table_for_read(item: &mut TomlItem) -> Option<&mut TomlTable> { + match item { + TomlItem::Table(table) => Some(table), + TomlItem::Value(value) => { + let inline = value.as_inline_table()?; + *item = TomlItem::Table(table_from_inline(inline)); + item.as_table_mut() + } + _ => None, + } +} + +fn serialize_mcp_server_table(config: &McpServerConfig) -> TomlTable { + let mut entry = TomlTable::new(); + entry.set_implicit(false); + + match &config.transport { + McpServerTransportConfig::Stdio { + command, + args, + env, + env_vars, + cwd, + } => { + entry["command"] = value(command.clone()); + if !args.is_empty() { + entry["args"] = array_from_iter(args.iter().cloned()); + } + if let Some(env) = env + && !env.is_empty() + { + entry["env"] = table_from_pairs(env.iter()); + } + if !env_vars.is_empty() { + entry["env_vars"] = array_from_env_vars(env_vars); + } + if let Some(cwd) = cwd { + entry["cwd"] = value(cwd.to_string_lossy().to_string()); + } + } + McpServerTransportConfig::StreamableHttp { + url, + bearer_token_env_var, + http_headers, + env_http_headers, + } => { + entry["url"] = value(url.clone()); + if let Some(env_var) = bearer_token_env_var { + entry["bearer_token_env_var"] = value(env_var.clone()); + } + if let Some(headers) = http_headers + && !headers.is_empty() + { + entry["http_headers"] = table_from_pairs(headers.iter()); + } + if let Some(headers) = env_http_headers + && !headers.is_empty() + { + entry["env_http_headers"] = table_from_pairs(headers.iter()); + } + } + } + + if !config.enabled { + entry["enabled"] = value(false); + } + if !config.is_local_environment() { + entry["environment_id"] = value(config.environment_id.clone()); + } + if config.required { + entry["required"] = value(true); + } + if config.supports_parallel_tool_calls { + entry["supports_parallel_tool_calls"] = value(true); + } + if let Some(timeout) = config.startup_timeout_sec { + entry["startup_timeout_sec"] = value(timeout.as_secs_f64()); + } + if let Some(timeout) = config.tool_timeout_sec { + entry["tool_timeout_sec"] = value(timeout.as_secs_f64()); + } + if let Some(approval_mode) = config.default_tools_approval_mode { + entry["default_tools_approval_mode"] = value(match approval_mode { + AppToolApproval::Auto => "auto", + AppToolApproval::Prompt => "prompt", + AppToolApproval::Approve => "approve", + }); + } + if let Some(enabled_tools) = &config.enabled_tools + && !enabled_tools.is_empty() + { + entry["enabled_tools"] = array_from_iter(enabled_tools.iter().cloned()); + } + if let Some(disabled_tools) = &config.disabled_tools + && !disabled_tools.is_empty() + { + entry["disabled_tools"] = array_from_iter(disabled_tools.iter().cloned()); + } + if let Some(scopes) = &config.scopes + && !scopes.is_empty() + { + entry["scopes"] = array_from_iter(scopes.iter().cloned()); + } + if let Some(oauth) = &config.oauth + && let Some(client_id) = &oauth.client_id + && !client_id.is_empty() + { + let mut oauth_table = TomlTable::new(); + oauth_table.set_implicit(false); + oauth_table["client_id"] = value(client_id.clone()); + entry["oauth"] = TomlItem::Table(oauth_table); + } + if let Some(resource) = &config.oauth_resource + && !resource.is_empty() + { + entry["oauth_resource"] = value(resource.clone()); + } + if !config.tools.is_empty() { + let mut tools = new_implicit_table(); + let mut tool_entries: Vec<_> = config.tools.iter().collect(); + tool_entries.sort_by_key(|(name, _)| *name); + for (name, tool_config) in tool_entries { + tools.insert(name, serialize_mcp_server_tool(tool_config)); + } + entry.insert("tools", TomlItem::Table(tools)); + } + + entry +} + +fn serialize_mcp_server_tool(config: &McpServerToolConfig) -> TomlItem { + let mut entry = TomlTable::new(); + entry.set_implicit(false); + if let Some(approval_mode) = config.approval_mode { + entry["approval_mode"] = value(match approval_mode { + AppToolApproval::Auto => "auto", + AppToolApproval::Prompt => "prompt", + AppToolApproval::Approve => "approve", + }); + } + TomlItem::Table(entry) +} + +pub(super) fn serialize_mcp_server(config: &McpServerConfig) -> TomlItem { + TomlItem::Table(serialize_mcp_server_table(config)) +} + +pub(super) fn serialize_mcp_server_inline(config: &McpServerConfig) -> InlineTable { + serialize_mcp_server_table(config).into_inline_table() +} + +pub(super) fn merge_inline_table(existing: &mut InlineTable, replacement: InlineTable) { + existing.retain(|key, _| replacement.get(key).is_some()); + + for (key, value) in replacement.iter() { + if let Some(existing_value) = existing.get_mut(key) { + let mut updated_value = value.clone(); + *updated_value.decor_mut() = existing_value.decor().clone(); + *existing_value = updated_value; + } else { + existing.insert(key.to_string(), value.clone()); + } + } +} + +fn table_from_inline(inline: &InlineTable) -> TomlTable { + let mut table = new_implicit_table(); + for (key, value) in inline.iter() { + let mut value = value.clone(); + let decor = value.decor_mut(); + decor.set_suffix(""); + table.insert(key, TomlItem::Value(value)); + } + table +} + +pub(super) fn new_implicit_table() -> TomlTable { + let mut table = TomlTable::new(); + table.set_implicit(true); + table +} + +pub(super) fn parse_tool_suggest_disabled_tool( + value: &TomlValue, +) -> Option { + let table = value.as_inline_table()?; + let kind = match table.get("type").and_then(TomlValue::as_str) { + Some("connector") => ToolSuggestDiscoverableType::Connector, + Some("plugin") => ToolSuggestDiscoverableType::Plugin, + _ => return None, + }; + let id = table.get("id").and_then(TomlValue::as_str)?; + Some(ToolSuggestDisabledTool { + kind, + id: id.to_string(), + }) +} + +pub(super) fn parse_tool_suggest_disabled_tool_table( + table: &TomlTable, +) -> Option { + let kind = match table.get("type").and_then(TomlItem::as_str) { + Some("connector") => ToolSuggestDiscoverableType::Connector, + Some("plugin") => ToolSuggestDiscoverableType::Plugin, + _ => return None, + }; + let id = table.get("id").and_then(TomlItem::as_str)?; + Some(ToolSuggestDisabledTool { + kind, + id: id.to_string(), + }) +} + +pub(super) fn tool_suggest_disabled_tools_value( + disabled_tools: &[ToolSuggestDisabledTool], +) -> TomlItem { + let mut array = TomlArray::new(); + for disabled_tool in disabled_tools { + let mut table = InlineTable::new(); + table.insert( + "type", + match disabled_tool.kind { + ToolSuggestDiscoverableType::Connector => "connector", + ToolSuggestDiscoverableType::Plugin => "plugin", + } + .into(), + ); + table.insert("id", disabled_tool.id.clone().into()); + array.push(table); + } + TomlItem::Value(array.into()) +} + +fn array_from_iter(iter: I) -> TomlItem +where + I: Iterator, +{ + let mut array = TomlArray::new(); + for value in iter { + array.push(value); + } + TomlItem::Value(array.into()) +} + +fn array_from_env_vars(env_vars: &[McpServerEnvVar]) -> TomlItem { + let mut array = TomlArray::new(); + for env_var in env_vars { + match env_var { + McpServerEnvVar::Name(name) => array.push(name.clone()), + McpServerEnvVar::Config { name, source } => { + let mut table = InlineTable::new(); + table.insert("name", name.clone().into()); + if let Some(source) = source { + table.insert("source", source.clone().into()); + } + array.push(table); + } + } + } + TomlItem::Value(array.into()) +} + +fn table_from_pairs<'a, I>(pairs: I) -> TomlItem +where + I: IntoIterator, +{ + let mut entries: Vec<_> = pairs.into_iter().collect(); + entries.sort_by_key(|(key, _)| *key); + let mut table = TomlTable::new(); + table.set_implicit(false); + for (key, val) in entries { + table.insert(key, value(val.clone())); + } + TomlItem::Table(table) +} diff --git a/codex-rs/core/src/realtime_context_tests.rs b/codex-rs/core/src/realtime_context_tests.rs index 955738af1b..56b0033956 100644 --- a/codex-rs/core/src/realtime_context_tests.rs +++ b/codex-rs/core/src/realtime_context_tests.rs @@ -13,10 +13,10 @@ use chrono::Utc; use codex_git_utils::GitSha; use codex_protocol::ThreadId; use codex_protocol::models::ContentItem; +use codex_protocol::models::PermissionProfile; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::GitInfo; -use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; use codex_thread_store::StoredThread; use core_test_support::PathBufExt; @@ -60,7 +60,7 @@ fn stored_thread(cwd: &str, title: &str, first_user_message: &str) -> StoredThre repository_url: None, }), approval_mode: AskForApproval::Never, - sandbox_policy: SandboxPolicy::new_read_only_policy(), + permission_profile: PermissionProfile::read_only(), token_usage: None, first_user_message: Some(first_user_message.to_string()), history: None, diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 135ea61194..7730a30ad7 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -630,6 +630,9 @@ async fn shutdown_session_runtime(sess: &Arc) { .unified_exec_manager .terminate_all_processes() .await; + if let Err(err) = sess.services.code_mode_service.shutdown().await { + warn!("failed to shutdown code mode session: {err}"); + } let mcp_shutdown = { let mut manager = sess.services.mcp_connection_manager.write().await; manager.begin_shutdown() diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index ba2496304c..2f44c8ed08 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -90,6 +90,7 @@ pub(super) async fn spawn_review_thread( sess.session_id().to_string(), sess.thread_id().to_string(), forked_from_thread_id, + &session_source, parent_turn_context.thread_source, review_turn_id.clone(), #[allow(deprecated)] diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 8219c19019..7944be57a5 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -943,16 +943,12 @@ async fn run_sampling_request( 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 _code_mode_worker = sess.services.code_mode_service.start_turn_worker( + &sess, + &turn_context, + Arc::clone(&router), + Arc::clone(&turn_diff_tracker), + ); let max_retries = turn_context.provider.info().stream_max_retries(); let mut retries = 0; let mut initial_input = Some(input); diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 910de99d33..5d09f0e611 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -543,6 +543,7 @@ impl Session { session_id.to_string(), thread_id.to_string(), session_configuration.forked_from_thread_id, + &session_configuration.session_source, session_configuration.thread_source, sub_id.clone(), cwd.clone(), diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 1ee4b4b58f..bb1f5fba65 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -5,6 +5,7 @@ use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_extension_api::ExtensionData; use codex_protocol::config_types::ModeKind; +use codex_protocol::items::ImageGenerationItem; use codex_protocol::items::TurnItem; use codex_utils_stream_parser::strip_citations; use tokio_util::sync::CancellationToken; @@ -125,6 +126,68 @@ async fn save_image_generation_result( Ok(path) } +pub(crate) async fn persist_image_generation_item( + sess: &Session, + turn_context: &TurnContext, + image_item: &mut ImageGenerationItem, +) -> Option { + let session_id = sess.conversation_id.to_string(); + match save_image_generation_result( + &turn_context.config.codex_home, + &session_id, + &image_item.id, + &image_item.result, + ) + .await + { + Ok(path) => { + image_item.saved_path = Some(path.clone()); + Some(path) + } + Err(err) => { + let output_path = image_generation_artifact_path( + &turn_context.config.codex_home, + &session_id, + &image_item.id, + ); + let output_dir = output_path + .parent() + .unwrap_or_else(|| turn_context.config.codex_home.clone()); + tracing::warn!( + call_id = %image_item.id, + output_dir = %output_dir.display(), + "failed to save generated image: {err}" + ); + None + } + } +} + +pub(crate) async fn finalize_image_generation_item( + sess: &Session, + turn_context: &TurnContext, + image_item: &mut ImageGenerationItem, +) { + if persist_image_generation_item(sess, turn_context, image_item) + .await + .is_none() + { + return; + } + let session_id = sess.conversation_id.to_string(); + let image_output_path = + image_generation_artifact_path(&turn_context.config.codex_home, &session_id, ""); + let image_output_dir = image_output_path + .parent() + .unwrap_or_else(|| turn_context.config.codex_home.clone()); + let message: ResponseItem = ContextualUserFragment::into(ImageGenerationInstructions::new( + image_output_dir.display(), + image_output_path.display(), + )); + sess.record_conversation_items(turn_context, &[message]) + .await; +} + /// Persist a completed model response item and record any cited memory usage. pub(crate) async fn record_completed_response_item( sess: &Session, @@ -487,49 +550,7 @@ pub(crate) async fn handle_non_tool_response_item( } } if let TurnItem::ImageGeneration(image_item) = &mut turn_item { - let session_id = sess.conversation_id.to_string(); - match save_image_generation_result( - &turn_context.config.codex_home, - &session_id, - &image_item.id, - &image_item.result, - ) - .await - { - Ok(path) => { - image_item.saved_path = Some(path); - let image_output_path = image_generation_artifact_path( - &turn_context.config.codex_home, - &session_id, - "", - ); - let image_output_dir = image_output_path - .parent() - .unwrap_or_else(|| turn_context.config.codex_home.clone()); - let message: ResponseItem = - ContextualUserFragment::into(ImageGenerationInstructions::new( - image_output_dir.display(), - image_output_path.display(), - )); - sess.record_conversation_items(turn_context, &[message]) - .await; - } - Err(err) => { - let output_path = image_generation_artifact_path( - &turn_context.config.codex_home, - &session_id, - &image_item.id, - ); - let output_dir = output_path - .parent() - .unwrap_or_else(|| turn_context.config.codex_home.clone()); - tracing::warn!( - call_id = %image_item.id, - output_dir = %output_dir.display(), - "failed to save generated image: {err}" - ); - } - } + finalize_image_generation_item(sess, turn_context, image_item).await; } Some(turn_item) } diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index c3650cbee2..75be4de508 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -590,12 +590,12 @@ impl ThreadManager { options: StartThreadOptions, forked_from_thread_id: Option, ) -> CodexResult { - let session_source = options - .session_source - .unwrap_or_else(|| self.state.session_source.clone()); - let thread_source = options - .thread_source - .or_else(|| options.initial_history.get_resumed_thread_source()); + let (resumed_session_source, resumed_thread_source) = options + .initial_history + .get_resumed_session_sources() + .unwrap_or_else(|| (self.state.session_source.clone(), None)); + let session_source = options.session_source.unwrap_or(resumed_session_source); + let thread_source = options.thread_source.or(resumed_thread_source); let parent_multi_agent_version = options.initial_history.get_multi_agent_version(); Box::pin(self.state.spawn_thread_with_source( options.config, @@ -682,17 +682,22 @@ impl ThreadManager { self.state.environment_manager.as_ref(), &config.cwd, ); - let thread_source = initial_history.get_resumed_thread_source(); - Box::pin(self.state.spawn_thread( + let (session_source, thread_source) = initial_history + .get_resumed_session_sources() + .unwrap_or_else(|| (self.state.session_source.clone(), None)); + Box::pin(self.state.spawn_thread_with_source( config, initial_history, auth_manager, self.agent_control(), + session_source, /*forked_from_thread_id*/ None, thread_source, Vec::new(), persist_extended_history, /*metrics_service_name*/ None, + /*inherited_shell_snapshot*/ None, + /*inherited_exec_policy*/ None, parent_trace, environments, /*user_shell_override*/ None, @@ -738,17 +743,22 @@ impl ThreadManager { self.state.environment_manager.as_ref(), &config.cwd, ); - let thread_source = initial_history.get_resumed_thread_source(); - Box::pin(self.state.spawn_thread( + let (session_source, thread_source) = initial_history + .get_resumed_session_sources() + .unwrap_or_else(|| (self.state.session_source.clone(), None)); + Box::pin(self.state.spawn_thread_with_source( config, initial_history, auth_manager, self.agent_control(), + session_source, /*forked_from_thread_id*/ None, thread_source, Vec::new(), /*persist_extended_history*/ false, /*metrics_service_name*/ None, + /*inherited_shell_snapshot*/ None, + /*inherited_exec_policy*/ None, /*parent_trace*/ None, environments, /*user_shell_override*/ Some(user_shell_override), diff --git a/codex-rs/core/src/tools/code_mode/delegate.rs b/codex-rs/core/src/tools/code_mode/delegate.rs new file mode 100644 index 0000000000..06bb71a280 --- /dev/null +++ b/codex-rs/core/src/tools/code_mode/delegate.rs @@ -0,0 +1,311 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex; + +use codex_code_mode::CellId; +use codex_code_mode::CodeModeNestedToolCall; +use codex_code_mode::CodeModeSessionDelegate; +use codex_code_mode::NotificationFuture; +use codex_code_mode::ToolInvocationFuture; +use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::ResponseItem; +use serde_json::Value as JsonValue; +use tokio::sync::oneshot; +use tokio::sync::watch; +use tokio_util::sync::CancellationToken; + +use super::ExecContext; +use super::PUBLIC_TOOL_NAME; +use super::call_nested_tool; +use crate::tools::ToolRouter; +use crate::tools::context::SharedTurnDiffTracker; +use crate::tools::parallel::ToolCallRuntime; + +pub(super) struct CodeModeDispatchBroker { + dispatch_tx: async_channel::Sender, + dispatch_rx: async_channel::Receiver, + dispatch_gates: Arc>>>, +} + +impl CodeModeDispatchBroker { + pub(super) fn new() -> Self { + let (dispatch_tx, dispatch_rx) = async_channel::unbounded(); + Self { + dispatch_tx, + dispatch_rx, + dispatch_gates: Arc::new(Mutex::new(HashMap::new())), + } + } + + pub(super) fn mark_cell_ready_for_dispatch(&self, cell_id: &CellId) { + dispatch_gate(&self.dispatch_gates, cell_id).send_replace(true); + } + + pub(super) fn close_cell(&self, cell_id: &CellId) { + remove_dispatch_gate(&self.dispatch_gates, cell_id); + } + + pub(super) fn start_turn_worker( + &self, + exec: ExecContext, + router: Arc, + tracker: SharedTurnDiffTracker, + ) -> CodeModeDispatchWorker { + let tool_runtime = ToolCallRuntime::new( + router, + Arc::clone(&exec.session), + Arc::clone(&exec.turn), + tracker, + ); + let host = Arc::new(CoreTurnHost { exec, tool_runtime }); + let dispatch_rx = self.dispatch_rx.clone(); + let dispatch_gates = Arc::clone(&self.dispatch_gates); + let (shutdown_tx, mut shutdown_rx) = oneshot::channel(); + tokio::spawn(async move { + loop { + let message = tokio::select! { + _ = &mut shutdown_rx => break, + message = dispatch_rx.recv() => message.ok(), + }; + let Some(message) = message else { + break; + }; + match message { + DispatchMessage::Notify { + call_id, + cell_id, + text, + cancellation_token, + response_tx, + } => { + let response = if wait_until_cell_ready_for_dispatch( + &dispatch_gates, + &cell_id, + &cancellation_token, + ) + .await + { + host.notify(call_id, cell_id, text).await + } else { + remove_dispatch_gate(&dispatch_gates, &cell_id); + Err("code mode notification cancelled".to_string()) + }; + let _ = response_tx.send(response); + } + DispatchMessage::InvokeTool { + invocation, + cancellation_token, + response_tx, + } => { + let cell_id = invocation.cell_id.clone(); + if !wait_until_cell_ready_for_dispatch( + &dispatch_gates, + &cell_id, + &cancellation_token, + ) + .await + { + remove_dispatch_gate(&dispatch_gates, &cell_id); + continue; + } + let host = Arc::clone(&host); + tokio::spawn(async move { + let response = tokio::select! { + response = host.invoke_tool( + invocation, + cancellation_token.clone(), + ) => response, + _ = cancellation_token.cancelled() => return, + }; + let _ = response_tx.send(response); + }); + } + } + } + }); + CodeModeDispatchWorker { + shutdown_tx: Some(shutdown_tx), + } + } +} + +fn dispatch_gate( + dispatch_gates: &Mutex>>, + cell_id: &CellId, +) -> watch::Sender { + let mut dispatch_gates = match dispatch_gates.lock() { + Ok(dispatch_gates) => dispatch_gates, + Err(poisoned) => poisoned.into_inner(), + }; + dispatch_gates + .entry(cell_id.clone()) + .or_insert_with(|| watch::channel(false).0) + .clone() +} + +fn remove_dispatch_gate( + dispatch_gates: &Mutex>>, + cell_id: &CellId, +) { + let mut dispatch_gates = match dispatch_gates.lock() { + Ok(dispatch_gates) => dispatch_gates, + Err(poisoned) => poisoned.into_inner(), + }; + dispatch_gates.remove(cell_id); +} + +async fn wait_until_cell_ready_for_dispatch( + dispatch_gates: &Mutex>>, + cell_id: &CellId, + cancellation_token: &CancellationToken, +) -> bool { + if cancellation_token.is_cancelled() { + return false; + } + let mut ready_rx = dispatch_gate(dispatch_gates, cell_id).subscribe(); + loop { + if *ready_rx.borrow_and_update() { + return true; + } + tokio::select! { + changed = ready_rx.changed() => { + if changed.is_err() { + return false; + } + } + _ = cancellation_token.cancelled() => return false, + } + } +} + +impl CodeModeSessionDelegate for CodeModeDispatchBroker { + fn invoke_tool<'a>( + &'a self, + invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async move { + if cancellation_token.is_cancelled() { + return Err("code mode nested tool call cancelled".to_string()); + } + let (response_tx, response_rx) = oneshot::channel(); + self.dispatch_tx + .send(DispatchMessage::InvokeTool { + invocation, + cancellation_token: cancellation_token.clone(), + response_tx, + }) + .await + .map_err(|_| "code mode nested tool dispatcher is unavailable".to_string())?; + tokio::select! { + response = response_rx => response + .map_err(|_| "code mode nested tool dispatcher stopped".to_string())?, + _ = cancellation_token.cancelled() => { + Err("code mode nested tool call cancelled".to_string()) + } + } + }) + } + + fn notify<'a>( + &'a self, + call_id: String, + cell_id: CellId, + text: String, + cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async move { + if cancellation_token.is_cancelled() { + return Err("code mode notification cancelled".to_string()); + } + let (response_tx, response_rx) = oneshot::channel(); + self.dispatch_tx + .send(DispatchMessage::Notify { + call_id, + cell_id, + text, + cancellation_token: cancellation_token.clone(), + response_tx, + }) + .await + .map_err(|_| "code mode notification dispatcher is unavailable".to_string())?; + tokio::select! { + response = response_rx => response + .map_err(|_| "code mode notification dispatcher stopped".to_string())?, + _ = cancellation_token.cancelled() => { + Err("code mode notification cancelled".to_string()) + } + } + }) + } + + fn cell_closed(&self, cell_id: &CellId) { + self.close_cell(cell_id); + } +} + +enum DispatchMessage { + InvokeTool { + invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + response_tx: oneshot::Sender>, + }, + Notify { + call_id: String, + cell_id: CellId, + text: String, + cancellation_token: CancellationToken, + response_tx: oneshot::Sender>, + }, +} + +pub(crate) struct CodeModeDispatchWorker { + shutdown_tx: Option>, +} + +impl Drop for CodeModeDispatchWorker { + fn drop(&mut self) { + if let Some(shutdown_tx) = self.shutdown_tx.take() { + let _ = shutdown_tx.send(()); + } + } +} + +struct CoreTurnHost { + exec: ExecContext, + tool_runtime: ToolCallRuntime, +} + +impl CoreTurnHost { + async fn invoke_tool( + &self, + invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> Result { + call_nested_tool( + self.exec.clone(), + self.tool_runtime.clone(), + invocation, + cancellation_token, + ) + .await + .map_err(|error| error.to_string()) + } + + async fn notify(&self, call_id: String, cell_id: CellId, text: String) -> Result<(), String> { + if text.trim().is_empty() { + return Ok(()); + } + self.exec + .session + .inject_if_running(vec![ResponseItem::CustomToolCallOutput { + call_id, + name: Some(PUBLIC_TOOL_NAME.to_string()), + output: FunctionCallOutputPayload::from_text(text), + }]) + .await + .map_err(|_| { + format!("failed to inject exec notify message for cell {cell_id}: no active turn") + }) + } +} diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 29f5623849..1d69c8f4cc 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -38,9 +38,22 @@ impl CodeModeExecuteHandler { let exec = ExecContext { session, turn }; let enabled_tools = codex_tools::collect_code_mode_tool_definitions(&self.nested_tool_specs); - // Allocate before starting V8 so the trace can create the parent - // CodeCell before model-authored JavaScript issues nested tool calls. - let runtime_cell_id = exec.session.services.code_mode_service.allocate_cell_id(); + let started_at = std::time::Instant::now(); + let started_cell = exec + .session + .services + .code_mode_service + .execute(codex_code_mode::ExecuteRequest { + tool_call_id: call_id.clone(), + enabled_tools, + source: args.code.clone(), + yield_time_ms: args.yield_time_ms, + max_output_tokens: args.max_output_tokens, + }) + .await + .map_err(FunctionCallError::RespondToModel)?; + let cell_id = started_cell.cell_id.clone(); + let runtime_cell_id = cell_id.to_string(); let code_cell_trace = exec .session .services @@ -51,19 +64,12 @@ impl CodeModeExecuteHandler { call_id.as_str(), args.code.as_str(), ); - let started_at = std::time::Instant::now(); - let response = exec - .session + exec.session .services .code_mode_service - .execute(codex_code_mode::ExecuteRequest { - cell_id: runtime_cell_id, - tool_call_id: call_id, - enabled_tools, - source: args.code, - yield_time_ms: args.yield_time_ms, - max_output_tokens: args.max_output_tokens, - }) + .mark_cell_ready_for_dispatch(&cell_id); + let response = started_cell + .initial_response() .await .map_err(FunctionCallError::RespondToModel)?; // Record the raw runtime boundary. The model-visible custom-tool output @@ -74,6 +80,10 @@ impl CodeModeExecuteHandler { // here when the first response also ended the runtime. if !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) { code_cell_trace.record_ended(&response); + exec.session + .services + .code_mode_service + .finish_cell_dispatch(&cell_id); } handle_runtime_response(&exec, response, args.max_output_tokens, started_at) .await diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index ff005646c2..ade23cf3b7 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -1,3 +1,4 @@ +mod delegate; mod execute_handler; pub(crate) mod execute_spec; mod response_adapter; @@ -7,13 +8,12 @@ pub(crate) mod wait_spec; use std::sync::Arc; use std::time::Duration; +use codex_code_mode::CellId; use codex_code_mode::CodeModeNestedToolCall; +use codex_code_mode::CodeModeSession; use codex_code_mode::CodeModeToolKind; -use codex_code_mode::CodeModeTurnHost; use codex_code_mode::RuntimeResponse; use codex_protocol::models::FunctionCallOutputContentItem; -use codex_protocol::models::FunctionCallOutputPayload; -use codex_protocol::models::ResponseItem; use serde_json::Value as JsonValue; use tokio_util::sync::CancellationToken; @@ -36,6 +36,8 @@ use codex_utils_output_truncation::TruncationPolicy; use codex_utils_output_truncation::formatted_truncate_text_content_items_with_policy; use codex_utils_output_truncation::truncate_function_output_items_with_policy; +use delegate::CodeModeDispatchBroker; +use delegate::CodeModeDispatchWorker; pub(crate) use execute_handler::CodeModeExecuteHandler; use response_adapter::into_function_call_output_content_items; pub(crate) use wait_handler::CodeModeWaitHandler; @@ -56,42 +58,67 @@ pub(crate) struct ExecContext { } pub(crate) struct CodeModeService { - inner: codex_code_mode::CodeModeService, + session: Option>, + dispatch_broker: Arc, } impl CodeModeService { pub(crate) fn new() -> Self { + let dispatch_broker = Arc::new(CodeModeDispatchBroker::new()); Self { - inner: codex_code_mode::CodeModeService::new(), + session: Some(Arc::new(codex_code_mode::CodeModeService::with_delegate( + dispatch_broker.clone(), + ))), + dispatch_broker, } } - pub(crate) fn allocate_cell_id(&self) -> String { - self.inner.allocate_cell_id() - } - pub(crate) async fn execute( &self, request: codex_code_mode::ExecuteRequest, - ) -> Result { - self.inner.execute(request).await + ) -> Result { + self.session()?.execute(request).await } pub(crate) async fn wait( &self, request: codex_code_mode::WaitRequest, ) -> Result { - self.inner.wait(request).await + self.session()?.wait(request).await } - pub(crate) async fn start_turn_worker( + pub(crate) async fn terminate( + &self, + cell_id: CellId, + ) -> Result { + self.session()?.terminate(cell_id).await + } + + pub(crate) async fn shutdown(&self) -> Result<(), String> { + match &self.session { + Some(session) => session.shutdown().await, + None => Ok(()), + } + } + + pub(crate) fn mark_cell_ready_for_dispatch(&self, cell_id: &codex_code_mode::CellId) { + self.dispatch_broker.mark_cell_ready_for_dispatch(cell_id); + } + + pub(crate) fn finish_cell_dispatch(&self, cell_id: &CellId) { + self.dispatch_broker.close_cell(cell_id); + } + + pub(crate) fn start_turn_worker( &self, session: &Arc, turn: &Arc, router: Arc, tracker: SharedTurnDiffTracker, - ) -> Option { - if !matches!(turn.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) { + ) -> Option { + if !matches!(turn.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) + || self.session.is_none() + { return None; } @@ -99,50 +126,16 @@ impl CodeModeService { session: Arc::clone(session), turn: Arc::clone(turn), }; - let tool_runtime = - ToolCallRuntime::new(router, Arc::clone(session), Arc::clone(turn), tracker); - let host = Arc::new(CoreTurnHost { exec, tool_runtime }); - Some(self.inner.start_turn_worker(host)) - } -} - -struct CoreTurnHost { - exec: ExecContext, - tool_runtime: ToolCallRuntime, -} - -#[async_trait::async_trait] -impl CodeModeTurnHost for CoreTurnHost { - async fn invoke_tool( - &self, - invocation: CodeModeNestedToolCall, - cancellation_token: CancellationToken, - ) -> Result { - call_nested_tool( - self.exec.clone(), - self.tool_runtime.clone(), - invocation, - cancellation_token, + Some( + self.dispatch_broker + .start_turn_worker(exec, router, tracker), ) - .await - .map_err(|error| error.to_string()) } - async fn notify(&self, call_id: String, cell_id: String, text: String) -> Result<(), String> { - if text.trim().is_empty() { - return Ok(()); - } - self.exec - .session - .inject_if_running(vec![ResponseItem::CustomToolCallOutput { - call_id, - name: Some(PUBLIC_TOOL_NAME.to_string()), - output: FunctionCallOutputPayload::from_text(text), - }]) - .await - .map_err(|_| { - format!("failed to inject exec notify message for cell {cell_id}: no active turn") - }) + fn session(&self) -> Result<&Arc, String> { + self.session + .as_ref() + .ok_or_else(|| "code mode is unavailable".to_string()) } } @@ -273,7 +266,7 @@ async fn call_nested_tool( .handle_tool_call_with_source( call, ToolCallSource::CodeMode { - cell_id, + cell_id: cell_id.to_string(), runtime_tool_call_id, }, cancellation_token, diff --git a/codex-rs/core/src/tools/code_mode/wait_handler.rs b/codex-rs/core/src/tools/code_mode/wait_handler.rs index d0c0453df6..74cc314a7b 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -73,17 +73,24 @@ impl ToolExecutor for CodeModeWaitHandler { let args: ExecWaitArgs = parse_arguments(&arguments)?; let exec = ExecContext { session, turn }; let started_at = std::time::Instant::now(); - let wait_response = exec - .session - .services - .code_mode_service - .wait(codex_code_mode::WaitRequest { - cell_id: args.cell_id, - yield_time_ms: args.yield_time_ms, - terminate: args.terminate, - }) - .await - .map_err(FunctionCallError::RespondToModel)?; + let cell_id = codex_code_mode::CellId::new(args.cell_id); + let wait_response = if args.terminate { + exec.session + .services + .code_mode_service + .terminate(cell_id) + .await + } else { + exec.session + .services + .code_mode_service + .wait(codex_code_mode::WaitRequest { + cell_id, + yield_time_ms: args.yield_time_ms, + }) + .await + } + .map_err(FunctionCallError::RespondToModel)?; if let codex_code_mode::WaitOutcome::LiveCell(response) = &wait_response && !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) { @@ -98,8 +105,15 @@ impl ToolExecutor for CodeModeWaitHandler { exec.session .services .rollout_thread_trace - .code_cell_trace_context(exec.turn.sub_id.as_str(), runtime_cell_id) + .code_cell_trace_context( + exec.turn.sub_id.as_str(), + runtime_cell_id.as_str(), + ) .record_ended(response); + exec.session + .services + .code_mode_service + .finish_cell_dispatch(runtime_cell_id); } handle_runtime_response(&exec, wait_response.into(), args.max_tokens, started_at) .await diff --git a/codex-rs/core/src/tools/handlers/extension_tools.rs b/codex-rs/core/src/tools/handlers/extension_tools.rs index d04a6c520e..95b272bbf6 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -4,15 +4,19 @@ use std::sync::Weak; use codex_protocol::items::TurnItem; use codex_tools::ConversationHistory; use codex_tools::ExtensionTurnItem; +use codex_tools::ImageGenerationCompletionFuture; use codex_tools::ToolCall as ExtensionToolCall; use codex_tools::ToolName; use codex_tools::ToolSpec; use codex_tools::TurnItemEmissionFuture; use codex_tools::TurnItemEmitter; +use crate::context::ContextualUserFragment; +use crate::context::ImageGenerationInstructions; use crate::function_tool::FunctionCallError; use crate::session::session::Session; use crate::session::turn_context::TurnContext; +use crate::stream_events_utils::persist_image_generation_item; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; @@ -90,6 +94,50 @@ impl TurnItemEmitter for CoreTurnItemEmitter { session.emit_turn_item_completed(turn.as_ref(), item).await; }) } + + fn image_generation_completed<'a>( + &'a self, + call_id: String, + prompt: String, + result: String, + ) -> ImageGenerationCompletionFuture<'a> { + Box::pin(async move { + let (Some(session), Some(turn)) = (self.session.upgrade(), self.turn.upgrade()) else { + return None; + }; + let mut item = codex_protocol::items::ImageGenerationItem { + id: call_id, + status: "completed".to_string(), + revised_prompt: Some(prompt), + result, + saved_path: None, + }; + let output_hint = + persist_image_generation_item(session.as_ref(), turn.as_ref(), &mut item) + .await + .map(|saved_path| { + let output_dir = saved_path + .parent() + .unwrap_or_else(|| turn.config.codex_home.clone()); + ImageGenerationInstructions::new(output_dir.display(), saved_path.display()) + .body() + }); + let started_item = codex_protocol::items::ImageGenerationItem { + id: item.id.clone(), + status: "in_progress".to_string(), + revised_prompt: None, + result: String::new(), + saved_path: None, + }; + session + .emit_turn_item_started(turn.as_ref(), &TurnItem::ImageGeneration(started_item)) + .await; + session + .emit_turn_item_completed(turn.as_ref(), TurnItem::ImageGeneration(item)) + .await; + output_hint + }) + } } async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall { @@ -99,6 +147,7 @@ async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall { turn_id: invocation.turn.sub_id.clone(), call_id: invocation.call_id.clone(), tool_name: invocation.tool_name.clone(), + model: invocation.turn.model_info.slug.clone(), truncation_policy: invocation.turn.truncation_policy, conversation_history, turn_item_emitter: Arc::new(CoreTurnItemEmitter { @@ -259,6 +308,7 @@ mod tests { let weak_session = Arc::downgrade(&session); let weak_turn = Arc::downgrade(&turn); let turn_id = turn.sub_id.clone(); + let model = turn.model_info.slug.clone(); let truncation_policy = turn.truncation_policy; let history_item = ResponseItem::Message { id: None, @@ -302,6 +352,7 @@ mod tests { captured_call.tool_name, codex_tools::ToolName::plain("extension_echo") ); + assert_eq!(captured_call.model, model); assert_eq!(captured_call.truncation_policy, truncation_policy); assert_eq!( captured_call.conversation_history.items(), @@ -352,4 +403,130 @@ mod tests { assert_eq!(end.query, expected.query); assert_eq!(end.action, expected.action); } + + struct ImageGenerationExtensionExecutor { + output_hint: Arc>>, + } + + #[async_trait::async_trait] + impl codex_extension_api::ToolExecutor for ImageGenerationExtensionExecutor { + fn tool_name(&self) -> codex_tools::ToolName { + codex_tools::ToolName::namespaced("image_gen", "imagegen") + } + + fn spec(&self) -> codex_tools::ToolSpec { + codex_tools::ToolSpec::Function(codex_tools::ResponsesApiTool { + name: "imagegen".to_string(), + description: "Generates an image.".to_string(), + strict: false, + parameters: codex_tools::JsonSchema::default(), + output_schema: None, + defer_loading: None, + }) + } + + async fn handle( + &self, + call: codex_tools::ToolCall, + ) -> Result, codex_tools::FunctionCallError> { + let output_hint = call + .turn_item_emitter + .image_generation_completed( + call.call_id, + "A tiny blue square".to_string(), + "cG5n".to_string(), + ) + .await; + *self.output_hint.lock().await = output_hint; + Ok(Box::new(codex_tools::JsonToolOutput::new( + json!({ "ok": true }), + ))) + } + } + + #[tokio::test] + async fn image_generation_publication_is_finalized_by_core() { + let output_hint = Arc::new(Mutex::new(None)); + let handler = ExtensionToolAdapter::new(Arc::new(ImageGenerationExtensionExecutor { + output_hint: Arc::clone(&output_hint), + })); + let (session, turn, rx) = crate::session::tests::make_session_and_context_with_rx().await; + let expected_path = crate::stream_events_utils::image_generation_artifact_path( + &turn.config.codex_home, + &session.conversation_id.to_string(), + "call-image", + ); + let invocation = ToolInvocation { + session, + turn, + cancellation_token: tokio_util::sync::CancellationToken::new(), + tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())), + call_id: "call-image".to_string(), + tool_name: codex_tools::ToolName::namespaced("image_gen", "imagegen"), + source: ToolCallSource::Direct, + payload: ToolPayload::Function { + arguments: "{}".to_string(), + }, + }; + + crate::tools::registry::ToolExecutor::handle(&handler, invocation) + .await + .expect("extension call should succeed"); + + let started = rx.recv().await.expect("item started event"); + let EventMsg::ItemStarted(started) = started.msg else { + panic!("expected item started event"); + }; + let TurnItem::ImageGeneration(started_item) = started.item else { + panic!("expected image generation item"); + }; + let begin = rx.recv().await.expect("legacy image start event"); + assert!(matches!(begin.msg, EventMsg::ImageGenerationBegin(_))); + let completed = rx.recv().await.expect("item completed event"); + let EventMsg::ItemCompleted(completed) = completed.msg else { + panic!("expected item completed event"); + }; + let TurnItem::ImageGeneration(completed_item) = completed.item else { + panic!("expected image generation item"); + }; + let end = rx.recv().await.expect("legacy image end event"); + assert!(matches!(end.msg, EventMsg::ImageGenerationEnd(_))); + + assert_eq!( + started_item, + codex_protocol::items::ImageGenerationItem { + id: "call-image".to_string(), + status: "in_progress".to_string(), + revised_prompt: None, + result: String::new(), + saved_path: None, + } + ); + assert_eq!( + completed_item, + codex_protocol::items::ImageGenerationItem { + id: "call-image".to_string(), + status: "completed".to_string(), + revised_prompt: Some("A tiny blue square".to_string()), + result: "cG5n".to_string(), + saved_path: Some(expected_path.clone()), + } + ); + assert_eq!( + std::fs::read(&expected_path).expect("generated artifact should be saved"), + b"png" + ); + assert_eq!( + *output_hint.lock().await, + Some(format!( + "Generated images are saved to {} as {} by default.\n\ + If you need to use a generated image at another path, copy it and leave the original in place unless the user explicitly asks you to delete it.", + expected_path + .parent() + .expect("generated image path should have a parent") + .display(), + expected_path.display(), + )) + ); + } } diff --git a/codex-rs/core/src/turn_metadata.rs b/codex-rs/core/src/turn_metadata.rs index 4f9406d4b0..af4d23a7c1 100644 --- a/codex-rs/core/src/turn_metadata.rs +++ b/codex-rs/core/src/turn_metadata.rs @@ -25,6 +25,7 @@ use codex_protocol::ThreadId; use codex_protocol::config_types::WindowsSandboxLevel; use codex_protocol::models::PermissionProfile; use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; +use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadSource; use codex_utils_absolute_path::AbsolutePathBuf; @@ -117,9 +118,10 @@ impl From for TurnMetadataWorkspace { /// Base payload for the outbound model request `x-codex-turn-metadata` header. /// -/// Turn-owned state populates identity fields, including optional fork lineage. A concrete -/// request kind is added at outbound model dispatch so turns, startup prewarm, and compaction -/// remain distinguishable. Detached memory requests are constructed as `memory` directly. +/// Turn-owned state populates identity fields, including optional fork and subagent lineage. A +/// concrete request kind is added at outbound model dispatch so turns, startup prewarm, and +/// compaction remain distinguishable. Detached memory requests are constructed as `memory` +/// directly. #[derive(Clone, Debug, Serialize)] pub(crate) struct TurnMetadataBag { #[serde(default, skip_serializing_if = "Option::is_none")] @@ -131,6 +133,10 @@ pub(crate) struct TurnMetadataBag { #[serde(default, skip_serializing_if = "Option::is_none")] forked_from_thread_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] + parent_thread_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + subagent_kind: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] thread_source: Option, #[serde(default, skip_serializing_if = "Option::is_none")] turn_id: Option, @@ -141,27 +147,6 @@ pub(crate) struct TurnMetadataBag { } impl TurnMetadataBag { - fn new( - request_kind: Option, - session_id: Option, - thread_id: Option, - forked_from_thread_id: Option, - thread_source: Option, - turn_id: Option, - sandbox: Option, - ) -> Self { - Self { - request_kind, - session_id, - thread_id, - forked_from_thread_id, - thread_source, - turn_id, - workspaces: BTreeMap::new(), - sandbox, - } - } - fn with_workspace_git_metadata( mut self, repo_root: Option, @@ -206,6 +191,8 @@ fn merge_turn_metadata( | "turn_id" | TURN_STARTED_AT_UNIX_MS_KEY | "forked_from_thread_id" + | "parent_thread_id" + | "subagent_kind" | REQUEST_KIND_KEY | COMPACTION_KEY | WINDOW_ID_KEY @@ -232,15 +219,18 @@ pub async fn build_turn_metadata_header( get_has_changes(cwd), ); let latest_git_commit_hash = head_commit_hash.map(|sha| sha.0); - TurnMetadataBag::new( - Some(TurnMetadataRequestKind::Memory), - /*session_id*/ None, - /*thread_id*/ None, - /*forked_from_thread_id*/ None, - /*thread_source*/ None, - /*turn_id*/ None, - sandbox.map(ToString::to_string), - ) + TurnMetadataBag { + request_kind: Some(TurnMetadataRequestKind::Memory), + session_id: None, + thread_id: None, + forked_from_thread_id: None, + parent_thread_id: None, + subagent_kind: None, + thread_source: None, + turn_id: None, + workspaces: BTreeMap::new(), + sandbox: sandbox.map(ToString::to_string), + } .with_workspace_git_metadata( repo_root, Some(WorkspaceGitMetadata { @@ -271,6 +261,7 @@ impl TurnMetadataState { session_id: String, thread_id: String, forked_from_thread_id: Option, + session_source: &SessionSource, thread_source: Option, turn_id: String, cwd: AbsolutePathBuf, @@ -287,15 +278,31 @@ impl TurnMetadataState { ) .to_string(), ); - let base_metadata = TurnMetadataBag::new( - /*request_kind*/ None, - Some(session_id), - Some(thread_id), + let (parent_thread_id, subagent_kind) = match session_source { + SessionSource::SubAgent(subagent_source) => ( + subagent_source.parent_thread_id().or(forked_from_thread_id), + Some(subagent_source.kind().to_string()), + ), + SessionSource::Cli + | SessionSource::VSCode + | SessionSource::Exec + | SessionSource::Mcp + | SessionSource::Custom(_) + | SessionSource::Internal(_) + | SessionSource::Unknown => (None, None), + }; + let base_metadata = TurnMetadataBag { + request_kind: None, + session_id: Some(session_id), + thread_id: Some(thread_id), forked_from_thread_id, + parent_thread_id, + subagent_kind, thread_source, - Some(turn_id), + turn_id: Some(turn_id), + workspaces: BTreeMap::new(), sandbox, - ); + }; let base_header = base_metadata.to_header_value(); Self { diff --git a/codex-rs/core/src/turn_metadata_tests.rs b/codex-rs/core/src/turn_metadata_tests.rs index a482392726..f6338f1c63 100644 --- a/codex-rs/core/src/turn_metadata_tests.rs +++ b/codex-rs/core/src/turn_metadata_tests.rs @@ -3,12 +3,16 @@ use super::*; use crate::sandbox_tags::permission_profile_sandbox_tag; use codex_protocol::models::PermissionProfile; use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::ThreadSource; +use codex_utils_absolute_path::AbsolutePathBuf; use core_test_support::PathBufExt; use core_test_support::PathExt; use pretty_assertions::assert_eq; use serde_json::Value; use std::collections::HashMap; +use std::time::Duration; use tempfile::TempDir; use tokio::process::Command; @@ -19,10 +23,9 @@ fn test_mcp_turn_metadata_context() -> McpTurnMetadataContext<'static> { } } -#[tokio::test] -async fn build_turn_metadata_header_marks_detached_memory_without_turn_identity() { +async fn create_clean_git_repo(repo_name: &str) -> (TempDir, AbsolutePathBuf) { let temp_dir = TempDir::new().expect("temp dir"); - let repo_path = temp_dir.path().join("repo-東京").abs(); + let repo_path = temp_dir.path().join(repo_name).abs(); std::fs::create_dir_all(&repo_path).expect("create repo"); Command::new("git") @@ -43,7 +46,6 @@ async fn build_turn_metadata_header_marks_detached_memory_without_turn_identity( .output() .await .expect("git config user.email"); - std::fs::write(repo_path.join("README.md"), "hello").expect("write file"); Command::new("git") .args(["add", "."]) @@ -58,6 +60,13 @@ async fn build_turn_metadata_header_marks_detached_memory_without_turn_identity( .await .expect("git commit"); + (temp_dir, repo_path) +} + +#[tokio::test] +async fn build_turn_metadata_header_marks_detached_memory_without_turn_identity() { + let (_temp_dir, repo_path) = create_clean_git_repo("repo-東京").await; + let header = build_turn_metadata_header(&repo_path, Some("none")) .await .expect("header"); @@ -113,6 +122,7 @@ fn turn_metadata_state_uses_platform_sandbox_tag() { "session-a".to_string(), "thread-a".to_string(), /*forked_from_thread_id*/ None, + &SessionSource::Exec, Some(ThreadSource::User), "turn-a".to_string(), cwd, @@ -138,6 +148,9 @@ fn turn_metadata_state_uses_platform_sandbox_tag() { assert_eq!(session_id, Some("session-a")); assert_eq!(thread_id, Some("thread-a")); assert_eq!(thread_source, Some("user")); + assert!(json.get("forked_from_thread_id").is_none()); + assert!(json.get("parent_thread_id").is_none()); + assert!(json.get("subagent_kind").is_none()); assert!(json.get("session_source").is_none()); } @@ -150,6 +163,7 @@ fn turn_metadata_state_uses_explicit_subagent_thread_source() { "session-a".to_string(), "thread-a".to_string(), /*forked_from_thread_id*/ None, + &SessionSource::Exec, Some(ThreadSource::Subagent), "turn-a".to_string(), cwd, @@ -177,6 +191,7 @@ fn turn_metadata_state_includes_root_fork_lineage() { "session-a".to_string(), "thread-a".to_string(), Some(source_thread_id), + &SessionSource::Exec, Some(ThreadSource::User), "turn-a".to_string(), cwd, @@ -192,6 +207,122 @@ fn turn_metadata_state_includes_root_fork_lineage() { json["forked_from_thread_id"].as_str(), Some("11111111-1111-4111-8111-111111111111") ); + assert!(json.get("parent_thread_id").is_none()); + assert!(json.get("subagent_kind").is_none()); +} + +#[test] +fn turn_metadata_state_includes_thread_spawn_subagent_parent_without_fork() { + let temp_dir = TempDir::new().expect("temp dir"); + let cwd = temp_dir.path().abs(); + let permission_profile = PermissionProfile::read_only(); + let parent_thread_id = + ThreadId::from_string("22222222-2222-4222-8222-222222222222").expect("thread id"); + + let state = TurnMetadataState::new( + "session-a".to_string(), + "thread-a".to_string(), + /*forked_from_thread_id*/ None, + &SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + }), + Some(ThreadSource::Subagent), + "turn-a".to_string(), + cwd, + &permission_profile, + WindowsSandboxLevel::Disabled, + /*enforce_managed_network*/ false, + ); + + let header = state.current_header_value().expect("header"); + let json: Value = serde_json::from_str(&header).expect("json"); + + assert!(json.get("forked_from_thread_id").is_none()); + assert_eq!( + json["parent_thread_id"].as_str(), + Some("22222222-2222-4222-8222-222222222222") + ); + assert_eq!(json["subagent_kind"].as_str(), Some("thread_spawn")); +} + +#[test] +fn turn_metadata_state_includes_forked_thread_spawn_subagent_lineage() { + let temp_dir = TempDir::new().expect("temp dir"); + let cwd = temp_dir.path().abs(); + let permission_profile = PermissionProfile::read_only(); + let parent_thread_id = + ThreadId::from_string("33333333-3333-4333-8333-333333333333").expect("thread id"); + + let state = TurnMetadataState::new( + "session-a".to_string(), + "thread-a".to_string(), + Some(parent_thread_id), + &SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + }), + Some(ThreadSource::Subagent), + "turn-a".to_string(), + cwd, + &permission_profile, + WindowsSandboxLevel::Disabled, + /*enforce_managed_network*/ false, + ); + + let header = state.current_header_value().expect("header"); + let json: Value = serde_json::from_str(&header).expect("json"); + + assert_eq!( + json["forked_from_thread_id"].as_str(), + Some("33333333-3333-4333-8333-333333333333") + ); + assert_eq!( + json["parent_thread_id"].as_str(), + Some("33333333-3333-4333-8333-333333333333") + ); + assert_eq!(json["subagent_kind"].as_str(), Some("thread_spawn")); +} + +#[test] +fn turn_metadata_state_includes_known_parent_for_other_subagent() { + let temp_dir = TempDir::new().expect("temp dir"); + let cwd = temp_dir.path().abs(); + let permission_profile = PermissionProfile::read_only(); + let parent_thread_id = + ThreadId::from_string("44444444-4444-4444-8444-444444444444").expect("thread id"); + + let state = TurnMetadataState::new( + "session-a".to_string(), + "thread-a".to_string(), + Some(parent_thread_id), + &SessionSource::SubAgent(SubAgentSource::Other("guardian".to_string())), + Some(ThreadSource::Subagent), + "turn-a".to_string(), + cwd, + &permission_profile, + WindowsSandboxLevel::Disabled, + /*enforce_managed_network*/ false, + ); + + let header = state.current_header_value().expect("header"); + let json: Value = serde_json::from_str(&header).expect("json"); + + assert_eq!( + json["forked_from_thread_id"].as_str(), + Some("44444444-4444-4444-8444-444444444444") + ); + assert_eq!( + json["parent_thread_id"].as_str(), + Some("44444444-4444-4444-8444-444444444444") + ); + assert_eq!(json["subagent_kind"].as_str(), Some("guardian")); } #[test] @@ -204,6 +335,7 @@ fn turn_metadata_state_includes_turn_started_at_unix_ms_after_start() { "session-a".to_string(), "thread-a".to_string(), /*forked_from_thread_id*/ None, + &SessionSource::Exec, Some(ThreadSource::User), "turn-a".to_string(), cwd, @@ -232,6 +364,7 @@ fn turn_metadata_state_includes_model_and_reasoning_effort_only_in_request_meta( "session-a".to_string(), "thread-a".to_string(), /*forked_from_thread_id*/ None, + &SessionSource::Exec, /*thread_source*/ None, "turn-a".to_string(), cwd, @@ -279,6 +412,7 @@ fn turn_metadata_state_marks_user_input_requested_during_turn_only_for_mcp_reque "session-a".to_string(), "thread-a".to_string(), /*forked_from_thread_id*/ None, + &SessionSource::Exec, /*thread_source*/ None, "turn-a".to_string(), cwd, @@ -330,6 +464,7 @@ fn turn_metadata_state_ignores_client_reserved_metadata_before_start() { "session-a".to_string(), "thread-a".to_string(), /*forked_from_thread_id*/ None, + &SessionSource::Exec, Some(ThreadSource::User), "turn-a".to_string(), cwd, @@ -346,6 +481,11 @@ fn turn_metadata_state_ignores_client_reserved_metadata_before_start() { "forked_from_thread_id".to_string(), "client-supplied".to_string(), ), + ( + "parent_thread_id".to_string(), + "client-supplied".to_string(), + ), + ("subagent_kind".to_string(), "client-supplied".to_string()), ])); let header = state.current_header_value().expect("header"); @@ -353,6 +493,8 @@ fn turn_metadata_state_ignores_client_reserved_metadata_before_start() { assert!(json.get("turn_started_at_unix_ms").is_none()); assert!(json.get("forked_from_thread_id").is_none()); + assert!(json.get("parent_thread_id").is_none()); + assert!(json.get("subagent_kind").is_none()); } #[test] @@ -362,11 +504,20 @@ fn turn_metadata_state_merges_client_metadata_without_replacing_reserved_fields( let permission_profile = PermissionProfile::read_only(); let source_thread_id = ThreadId::from_string("44444444-4444-4444-8444-444444444444").expect("thread id"); + let parent_thread_id = + ThreadId::from_string("55555555-5555-4555-8555-555555555555").expect("thread id"); let state = TurnMetadataState::new( "session-a".to_string(), "thread-a".to_string(), Some(source_thread_id), + &SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + }), Some(ThreadSource::User), "turn-a".to_string(), cwd, @@ -388,6 +539,11 @@ fn turn_metadata_state_merges_client_metadata_without_replacing_reserved_fields( "forked_from_thread_id".to_string(), "client-supplied".to_string(), ), + ( + "parent_thread_id".to_string(), + "client-supplied".to_string(), + ), + ("subagent_kind".to_string(), "client-supplied".to_string()), ("turn_id".to_string(), "client-supplied".to_string()), (WINDOW_ID_KEY.to_string(), "client-supplied".to_string()), ("thread_source".to_string(), "client-supplied".to_string()), @@ -414,6 +570,11 @@ fn turn_metadata_state_merges_client_metadata_without_replacing_reserved_fields( json["forked_from_thread_id"].as_str(), Some("44444444-4444-4444-8444-444444444444") ); + assert_eq!( + json["parent_thread_id"].as_str(), + Some("55555555-5555-4555-8555-555555555555") + ); + assert_eq!(json["subagent_kind"].as_str(), Some("thread_spawn")); assert_eq!(json["thread_source"].as_str(), Some("user")); assert_eq!(json["turn_id"].as_str(), Some("turn-a")); assert!(json.get("request_kind").is_none()); @@ -451,6 +612,7 @@ fn turn_metadata_state_overlays_compaction_only_on_compaction_requests() { "session-a".to_string(), "thread-a".to_string(), /*forked_from_thread_id*/ None, + &SessionSource::Exec, Some(ThreadSource::User), "turn-a".to_string(), cwd, @@ -497,3 +659,59 @@ fn turn_metadata_state_overlays_compaction_only_on_compaction_requests() { assert_eq!(regular_json[WINDOW_ID_KEY].as_str(), Some("thread-a:3")); assert!(regular_json.get("compaction").is_none()); } + +#[tokio::test] +async fn turn_metadata_state_preserves_lineage_after_git_enrichment() { + let (_temp_dir, repo_path) = create_clean_git_repo("repo").await; + + let permission_profile = PermissionProfile::read_only(); + let parent_thread_id = + ThreadId::from_string("66666666-6666-4666-8666-666666666666").expect("thread id"); + let state = TurnMetadataState::new( + "session-a".to_string(), + "thread-a".to_string(), + Some(parent_thread_id), + &SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + }), + Some(ThreadSource::Subagent), + "turn-a".to_string(), + repo_path, + &permission_profile, + WindowsSandboxLevel::Disabled, + /*enforce_managed_network*/ false, + ); + + state.spawn_git_enrichment_task(); + + let json = tokio::time::timeout(Duration::from_secs(2), async { + loop { + let header = state.current_header_value().expect("header"); + let json: Value = serde_json::from_str(&header).expect("json"); + if json + .get("workspaces") + .and_then(Value::as_object) + .is_some_and(|workspaces| !workspaces.is_empty()) + { + return json; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("git enrichment should complete"); + + assert_eq!( + json["forked_from_thread_id"].as_str(), + Some("66666666-6666-4666-8666-666666666666") + ); + assert_eq!( + json["parent_thread_id"].as_str(), + Some("66666666-6666-4666-8666-666666666666") + ); + assert_eq!(json["subagent_kind"].as_str(), Some("thread_spawn")); +} diff --git a/codex-rs/core/src/windows_sandbox.rs b/codex-rs/core/src/windows_sandbox.rs index 19b404fb71..a1ff5dda09 100644 --- a/codex-rs/core/src/windows_sandbox.rs +++ b/codex-rs/core/src/windows_sandbox.rs @@ -169,6 +169,11 @@ pub fn run_elevated_setup( ) } +#[cfg(target_os = "windows")] +pub fn run_elevated_provisioning_setup(codex_home: &Path, real_user: &str) -> anyhow::Result<()> { + codex_windows_sandbox::run_elevated_provisioning_setup(codex_home, real_user) +} + #[cfg(not(target_os = "windows"))] pub fn run_elevated_setup( _permission_profile: &PermissionProfile, @@ -180,6 +185,11 @@ pub fn run_elevated_setup( anyhow::bail!("elevated Windows sandbox setup is only supported on Windows") } +#[cfg(not(target_os = "windows"))] +pub fn run_elevated_provisioning_setup(_codex_home: &Path, _real_user: &str) -> anyhow::Result<()> { + anyhow::bail!("elevated Windows sandbox setup is only supported on Windows") +} + #[cfg(target_os = "windows")] pub fn run_legacy_setup_preflight( permission_profile: &PermissionProfile, diff --git a/codex-rs/core/tests/common/lib.rs b/codex-rs/core/tests/common/lib.rs index 194fb2ae0a..3d7fd21228 100644 --- a/codex-rs/core/tests/common/lib.rs +++ b/codex-rs/core/tests/common/lib.rs @@ -248,6 +248,39 @@ where wait_for_event_with_timeout(codex, predicate, Duration::from_secs(1)).await } +/// Waits for a configured MCP server to finish startup and requires it to be ready. +pub async fn wait_for_mcp_server(codex: &CodexThread, server_name: &str) -> anyhow::Result<()> { + use codex_protocol::protocol::EventMsg; + + // Wait for the startup summary regardless of outcome, then interpret the + // requested server's ready, failed, or cancelled entry below. + let summary = loop { + let event = codex + .next_event() + .await + .expect("stream ended unexpectedly while waiting for MCP startup"); + if let EventMsg::McpStartupComplete(summary) = event.msg { + break summary; + } + }; + if let Some(failure) = summary + .failed + .iter() + .find(|failure| failure.server == server_name) + { + let error = &failure.error; + anyhow::bail!("MCP server {server_name} failed to start: {error}"); + } + if summary.cancelled.iter().any(|server| server == server_name) { + anyhow::bail!("MCP server {server_name} startup was cancelled"); + } + assert!( + summary.ready.iter().any(|server| server == server_name), + "expected MCP server {server_name} to be ready; startup summary: {summary:?}" + ); + Ok(()) +} + pub async fn submit_thread_settings( codex: &CodexThread, thread_settings: codex_protocol::protocol::ThreadSettingsOverrides, diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index cbf6993bc9..9881869313 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -38,6 +38,7 @@ use core_test_support::test_codex::test_codex; use core_test_support::test_codex::turn_permission_fields; use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; +use core_test_support::wait_for_mcp_server; use pretty_assertions::assert_eq; use serde_json::Value; use std::collections::HashMap; @@ -283,6 +284,7 @@ async fn run_code_mode_turn_with_rmcp_config( .expect("test mcp servers should accept any configuration"); }); let test = builder.build(server).await?; + wait_for_mcp_server(&test.codex, "rmcp").await?; responses::mount_sse_once( server, diff --git a/codex-rs/core/tests/suite/hooks_mcp.rs b/codex-rs/core/tests/suite/hooks_mcp.rs index 940db7cf4c..00f1546a7c 100644 --- a/codex-rs/core/tests/suite/hooks_mcp.rs +++ b/codex-rs/core/tests/suite/hooks_mcp.rs @@ -22,6 +22,7 @@ use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use core_test_support::stdio_server_bin; use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_mcp_server; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; @@ -284,6 +285,7 @@ async fn pre_tool_use_blocks_mcp_tool_before_execution( }) .build(&server) .await?; + wait_for_mcp_server(&test.codex, RMCP_SERVER).await?; test.submit_turn("call the rmcp echo tool with the MCP pre hook") .await?; @@ -375,6 +377,7 @@ async fn pre_tool_use_rewrites_mcp_tool_before_execution() -> Result<()> { }) .build(&server) .await?; + wait_for_mcp_server(&test.codex, RMCP_SERVER).await?; test.submit_turn("call the rmcp echo tool with the MCP pre hook rewrite") .await?; @@ -471,6 +474,7 @@ async fn post_tool_use_records_mcp_tool_payload_and_context( }) .build(&server) .await?; + wait_for_mcp_server(&test.codex, RMCP_SERVER).await?; test.submit_turn("call the rmcp echo tool with the MCP post hook") .await?; diff --git a/codex-rs/core/tests/suite/plugins.rs b/codex-rs/core/tests/suite/plugins.rs index b89b611a83..d7d7c7e8f5 100644 --- a/codex-rs/core/tests/suite/plugins.rs +++ b/codex-rs/core/tests/suite/plugins.rs @@ -6,7 +6,6 @@ use std::time::Duration; use std::time::Instant; use anyhow::Result; -use anyhow::bail; use codex_features::Feature; use codex_login::CodexAuth; use codex_protocol::protocol::EventMsg; @@ -21,7 +20,7 @@ use core_test_support::skip_if_no_network; use core_test_support::stdio_server_bin; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; -use core_test_support::wait_for_event_with_timeout; +use core_test_support::wait_for_mcp_server; use tempfile::TempDir; use wiremock::MockServer; @@ -140,45 +139,6 @@ async fn build_apps_enabled_plugin_test_codex( .codex) } -async fn wait_for_sample_mcp_ready(codex: &codex_core::CodexThread) -> Result<()> { - let startup_event = wait_for_event_with_timeout( - codex, - |ev| match ev { - EventMsg::McpStartupComplete(summary) => { - summary.ready.iter().any(|server| server == "sample") - || summary - .failed - .iter() - .any(|failure| failure.server == "sample") - || summary.cancelled.iter().any(|server| server == "sample") - } - _ => false, - }, - Duration::from_secs(70), - ) - .await; - let EventMsg::McpStartupComplete(startup) = startup_event else { - unreachable!("event guard guarantees McpStartupComplete"); - }; - if let Some(failure) = startup - .failed - .iter() - .find(|failure| failure.server == "sample") - { - let error = &failure.error; - bail!("plugin MCP server failed to start: {error}"); - } - if startup.cancelled.iter().any(|server| server == "sample") { - bail!("plugin MCP server startup was cancelled"); - } - assert!( - startup.ready.iter().any(|server| server == "sample"), - "expected plugin MCP server to be ready; startup summary: {startup:?}" - ); - - Ok(()) -} - fn tool_names(body: &serde_json::Value) -> Vec { body.get("tools") .and_then(serde_json::Value::as_array) @@ -296,7 +256,7 @@ async fn explicit_plugin_mentions_inject_plugin_guidance() -> Result<()> { let codex = build_apps_enabled_plugin_test_codex(&server, codex_home, apps_server.chatgpt_base_url) .await?; - wait_for_sample_mcp_ready(&codex).await?; + wait_for_mcp_server(&codex, "sample").await?; codex .submit(Op::UserInput { diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 95658fb826..eb37bfcf7a 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -55,7 +55,7 @@ use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::test_codex; use core_test_support::test_codex::turn_permission_fields; use core_test_support::wait_for_event; -use core_test_support::wait_for_event_with_timeout; +use core_test_support::wait_for_mcp_server; use reqwest::Client; use reqwest::StatusCode; use serde_json::Value; @@ -268,44 +268,6 @@ fn copy_binary_to_remote_env( Ok(remote_path) } -async fn wait_for_mcp_server(fixture: &TestCodex, server_name: &str) -> anyhow::Result<()> { - let startup_event = wait_for_event_with_timeout( - &fixture.codex, - |ev| match ev { - EventMsg::McpStartupComplete(summary) => { - summary.ready.iter().any(|server| server == server_name) - || summary - .failed - .iter() - .any(|failure| failure.server == server_name) - || summary.cancelled.iter().any(|server| server == server_name) - } - _ => false, - }, - Duration::from_secs(70), - ) - .await; - let EventMsg::McpStartupComplete(summary) = startup_event else { - unreachable!("event guard guarantees McpStartupComplete"); - }; - if let Some(failure) = summary - .failed - .iter() - .find(|failure| failure.server == server_name) - { - let error = &failure.error; - anyhow::bail!("MCP server {server_name} failed to start: {error}"); - } - if summary.cancelled.iter().any(|server| server == server_name) { - anyhow::bail!("MCP server {server_name} startup was cancelled"); - } - ensure!( - summary.ready.iter().any(|server| server == server_name), - "expected MCP server {server_name} to be ready; startup summary: {summary:?}" - ); - Ok(()) -} - struct TestMcpServerOptions { environment_id: String, supports_parallel_tool_calls: bool, @@ -517,6 +479,8 @@ async fn stdio_server_round_trip() -> anyhow::Result<()> { }) .build_with_remote_env(&server) .await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; + fixture .codex .submit(read_only_user_turn(&fixture, "call the rmcp echo tool")) @@ -640,6 +604,7 @@ async fn stdio_server_uses_configured_cwd_before_runtime_fallback() -> anyhow::R }) .build_with_remote_env(&server) .await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; let expected_cwd = expected_cwd .lock() @@ -702,6 +667,7 @@ async fn local_stdio_server_uses_runtime_fallback_cwd_when_config_omits_cwd() -> }) .build(&server) .await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; let expected_cwd = expected_cwd .lock() @@ -760,7 +726,7 @@ async fn stdio_mcp_tool_call_includes_sandbox_state_meta() -> anyhow::Result<()> .build_with_remote_env(&server) .await?; - wait_for_mcp_server(&fixture, server_name).await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; fixture .submit_turn_with_permission_profile( @@ -857,6 +823,8 @@ async fn stdio_mcp_parallel_tool_calls_default_false_runs_serially() -> anyhow:: }) .build_with_remote_env(&server) .await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; + fixture .codex // Keep this baseline on the mutable sync tool so read-only hints do not @@ -991,6 +959,8 @@ async fn stdio_mcp_read_only_tool_calls_run_concurrently_without_server_opt_in() }) .build_with_remote_env(&server) .await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; + fixture .codex .submit(read_only_user_turn( @@ -1073,6 +1043,8 @@ async fn stdio_mcp_parallel_tool_calls_opt_in_runs_concurrently() -> anyhow::Res }) .build_with_remote_env(&server) .await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; + fixture .codex // Exercise the server opt-in with the mutable sync tool rather than the @@ -1157,7 +1129,7 @@ async fn stdio_image_responses_round_trip() -> anyhow::Result<()> { }) .build_with_remote_env(&server) .await?; - wait_for_mcp_server(&fixture, server_name).await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; fixture .codex @@ -1290,7 +1262,7 @@ async fn stdio_image_responses_preserve_original_detail_metadata() -> anyhow::Re }) .build_with_remote_env(&server) .await?; - wait_for_mcp_server(&fixture, server_name).await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; fixture .codex @@ -1428,6 +1400,7 @@ async fn stdio_image_responses_are_sanitized_for_text_only_model() -> anyhow::Re }) .build_with_remote_env(&server) .await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; fixture .thread_manager @@ -1530,6 +1503,8 @@ async fn stdio_server_propagates_whitelisted_env_vars() -> anyhow::Result<()> { }) .build_with_remote_env(&server) .await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; + fixture .codex .submit(read_only_user_turn(&fixture, "call the rmcp echo tool")) @@ -1648,6 +1623,7 @@ async fn stdio_server_propagates_explicit_local_env_var_source() -> anyhow::Resu }) .build_with_remote_env(&server) .await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; fixture .codex @@ -1740,6 +1716,7 @@ async fn remote_stdio_env_var_source_does_not_copy_local_env() -> anyhow::Result }) .build_with_remote_env(&server) .await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; fixture .codex @@ -1923,6 +1900,8 @@ async fn streamable_http_tool_call_round_trip() -> anyhow::Result<()> { }) .build_with_remote_env(&server) .await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; + // Phase 4: submit the user turn that should trigger the MCP tool call. fixture .codex @@ -2111,7 +2090,7 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { .await?; // Phase 5: wait for MCP startup before the turn is submitted, which keeps // failures tied to server startup/discovery. - wait_for_mcp_server(&fixture, server_name).await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; // Phase 6: submit the user turn that should invoke the OAuth-backed tool. fixture diff --git a/codex-rs/core/tests/suite/search_tool.rs b/codex-rs/core/tests/suite/search_tool.rs index 0dc956d951..42b1b6bb3f 100644 --- a/codex-rs/core/tests/suite/search_tool.rs +++ b/codex-rs/core/tests/suite/search_tool.rs @@ -47,6 +47,7 @@ use core_test_support::skip_if_no_network; use core_test_support::stdio_server_bin; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; +use core_test_support::wait_for_mcp_server; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; @@ -1114,6 +1115,7 @@ async fn tool_search_indexes_only_enabled_non_app_mcp_tools() -> Result<()> { .expect("test mcp servers should accept any configuration"); }); let test = builder.build(&server).await?; + wait_for_mcp_server(&test.codex, "rmcp").await?; test.submit_turn_with_approval_and_permission_profile( "Find the rmcp echo and image tools.", @@ -1243,6 +1245,7 @@ async fn tool_search_surfaced_mcp_tool_errors_are_returned_to_model() -> Result< .expect("test mcp servers should accept any configuration"); }); let test = builder.build(&server).await?; + wait_for_mcp_server(&test.codex, "rmcp").await?; test.codex .submit(Op::UserInput { @@ -1391,6 +1394,7 @@ async fn tool_search_uses_non_app_mcp_server_instructions_as_namespace_descripti .expect("test mcp servers should accept any configuration"); }); let test = builder.build(&server).await?; + wait_for_mcp_server(&test.codex, "rmcp").await?; test.submit_turn_with_approval_and_permission_profile( "Find the rmcp echo tool.", diff --git a/codex-rs/core/tests/suite/sqlite_state.rs b/codex-rs/core/tests/suite/sqlite_state.rs index 74557f56e1..8e4ef48236 100644 --- a/codex-rs/core/tests/suite/sqlite_state.rs +++ b/codex-rs/core/tests/suite/sqlite_state.rs @@ -29,6 +29,7 @@ use core_test_support::test_codex::test_codex; use core_test_support::test_codex::turn_permission_fields; use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; +use core_test_support::wait_for_mcp_server; use pretty_assertions::assert_eq; use serde_json::json; use std::collections::HashMap; @@ -461,6 +462,7 @@ async fn mcp_call_marks_thread_memory_mode_polluted_when_configured() -> Result< .expect("test mcp servers should accept any configuration"); }); let test = builder.build(&server).await?; + wait_for_mcp_server(&test.codex, server_name).await?; let db = test.codex.state_db().expect("state db enabled"); let thread_id = test.session_configured.thread_id; let cwd = test.cwd_path().to_path_buf(); diff --git a/codex-rs/core/tests/suite/truncation.rs b/codex-rs/core/tests/suite/truncation.rs index fe078977ff..db90b663e0 100644 --- a/codex-rs/core/tests/suite/truncation.rs +++ b/codex-rs/core/tests/suite/truncation.rs @@ -23,6 +23,7 @@ use core_test_support::skip_if_no_network; use core_test_support::stdio_server_bin; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; +use core_test_support::wait_for_mcp_server; use serde_json::Value; use serde_json::json; use std::collections::HashMap; @@ -410,6 +411,7 @@ async fn mcp_tool_call_output_exceeds_limit_truncated_for_model() -> Result<()> config.tool_output_token_limit = Some(500); }); let fixture = builder.build(&server).await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; fixture .submit_turn_with_permission_profile( @@ -509,6 +511,7 @@ async fn mcp_image_output_preserves_image_and_no_text_summary() -> Result<()> { .expect("test mcp servers should accept any configuration"); }); let fixture = builder.build(&server).await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; let session_model = fixture.session_configured.model.clone(); let permission_profile = PermissionProfile::read_only(); let sandbox_policy = permission_profile.to_legacy_sandbox_policy(fixture.cwd.path())?; @@ -798,6 +801,7 @@ async fn mcp_tool_call_output_not_truncated_with_custom_limit() -> Result<()> { .expect("test mcp servers should accept any configuration"); }); let fixture = builder.build(&server).await?; + wait_for_mcp_server(&fixture.codex, server_name).await?; fixture .submit_turn_with_permission_profile( diff --git a/codex-rs/exec-server/src/client/reqwest_http_client.rs b/codex-rs/exec-server/src/client/reqwest_http_client.rs index 44305e466a..9e22fc01d3 100644 --- a/codex-rs/exec-server/src/client/reqwest_http_client.rs +++ b/codex-rs/exec-server/src/client/reqwest_http_client.rs @@ -5,6 +5,7 @@ //! - in a remote environment, that means the remote runtime after the //! orchestrator has forwarded `http/request` over JSON-RPC +use std::error::Error as StdError; use std::time::Duration; use codex_app_server_protocol::JSONRPCErrorError; @@ -135,15 +136,21 @@ impl ReqwestHttpRequestRunner { } let headers = Self::build_headers(params.headers)?; - let mut request = self.client.request(method, url).headers(headers); + let mut request = self.client.request(method.clone(), url).headers(headers); if let Some(body) = params.body { request = request.body(body.into_inner()); } - let response = request - .send() - .await - .map_err(|error| internal_error(format!("http/request failed: {error}")))?; + let response = match request.send().await { + Ok(response) => response, + Err(error) => { + let error_message = error.to_string(); + log_send_error(&method, error); + return Err(internal_error(format!( + "http/request failed: {error_message}" + ))); + } + }; let status = response.status().as_u16(); let headers = Self::response_headers(response.headers()); @@ -265,3 +272,26 @@ impl ReqwestHttpRequestRunner { .collect() } } + +fn log_send_error(method: &Method, error: reqwest::Error) { + let error = error.without_url(); + let source_chain = error_source_chain(&error); + tracing::warn!( + http_method = method.as_str(), + error_is_timeout = error.is_timeout(), + error_is_connect = error.is_connect(), + error = %error, + error_sources = ?source_chain, + "http/request send failed" + ); +} + +fn error_source_chain(error: &reqwest::Error) -> Option { + let mut sources = Vec::new(); + let mut source = error.source(); + while let Some(error) = source { + sources.push(error.to_string()); + source = error.source(); + } + (!sources.is_empty()).then(|| sources.join(": ")) +} diff --git a/codex-rs/exec-server/src/fs_sandbox.rs b/codex-rs/exec-server/src/fs_sandbox.rs index 3356791d44..8ef20b2d1d 100644 --- a/codex-rs/exec-server/src/fs_sandbox.rs +++ b/codex-rs/exec-server/src/fs_sandbox.rs @@ -230,6 +230,8 @@ fn helper_env_from_vars( fn helper_env_key_is_allowed(key: &str) -> bool { FS_HELPER_ENV_ALLOWLIST.contains(&key) + // CoreFoundation consults this before falling back to user lookup during helper startup. + || (cfg!(target_os = "macos") && key == "__CF_USER_TEXT_ENCODING") || bazel_bwrap_env_key_is_allowed(key) || (cfg!(windows) && key.eq_ignore_ascii_case("PATH")) } @@ -298,6 +300,7 @@ fn spawn_command( command.stdin(std::process::Stdio::piped()); command.stdout(std::process::Stdio::piped()); command.stderr(std::process::Stdio::piped()); + command.kill_on_drop(true); command.spawn().map_err(io_error) } @@ -433,6 +436,26 @@ mod tests { ); } + #[cfg(target_os = "macos")] + #[test] + fn helper_env_preserves_corefoundation_text_encoding() { + let env = helper_env_from_vars( + [ + ("__CF_USER_TEXT_ENCODING", "0x1F6:0x0:0x0"), + ("HOME", "/Users/test"), + ] + .map(|(key, value)| (OsString::from(key), OsString::from(value))), + ); + + assert_eq!( + env, + HashMap::from([( + "__CF_USER_TEXT_ENCODING".to_string(), + "0x1F6:0x0:0x0".to_string(), + )]) + ); + } + #[cfg(windows)] #[test] fn helper_env_preserves_windows_path_key_for_system_bwrap_discovery() { diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 4705581c8a..2f00a8dd98 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -13,6 +13,7 @@ pub use capabilities::ResponseItemInjector; pub use codex_tools::ConversationHistory; pub use codex_tools::ExtensionTurnItem; pub use codex_tools::FunctionCallError; +pub use codex_tools::ImageGenerationCompletionFuture; pub use codex_tools::JsonToolOutput; pub use codex_tools::NoopTurnItemEmitter; pub use codex_tools::ResponsesApiTool; diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index 92f4313fc7..f325c5f90f 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -1129,6 +1129,7 @@ fn tool_call(tool_name: &str, call_id: &str, arguments: serde_json::Value) -> To turn_id: "turn-1".to_string(), call_id: call_id.to_string(), tool_name: codex_extension_api::ToolName::plain(tool_name), + model: "gpt-test".to_string(), truncation_policy: TruncationPolicy::Bytes(1024), conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), diff --git a/codex-rs/ext/image-generation/Cargo.toml b/codex-rs/ext/image-generation/Cargo.toml index d4f2cb50de..995f1786bc 100644 --- a/codex-rs/ext/image-generation/Cargo.toml +++ b/codex-rs/ext/image-generation/Cargo.toml @@ -14,7 +14,6 @@ workspace = true [dependencies] async-trait = { workspace = true } -base64 = { workspace = true } codex-api = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } @@ -28,10 +27,6 @@ http = { workspace = true } schemars = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } -tokio = { workspace = true, features = ["fs"] } -tracing = { workspace = true } [dev-dependencies] pretty_assertions = { workspace = true } -tempfile = { workspace = true } -tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/codex-rs/ext/image-generation/src/extension.rs b/codex-rs/ext/image-generation/src/extension.rs index 8f0b09f8cb..f7a9f2406f 100644 --- a/codex-rs/ext/image-generation/src/extension.rs +++ b/codex-rs/ext/image-generation/src/extension.rs @@ -1,4 +1,3 @@ -use std::path::PathBuf; use std::sync::Arc; use codex_core::config::Config; @@ -17,7 +16,6 @@ use codex_model_provider_info::ModelProviderInfo; use crate::backend::CodexImagesBackend; use crate::tool::ImageGenerationTool; -use crate::tool::generated_image_output_dir; #[derive(Clone)] struct ImageGenerationExtension { @@ -28,7 +26,6 @@ struct ImageGenerationExtension { struct ImageGenerationExtensionConfig { enabled: bool, provider: ModelProviderInfo, - codex_home: PathBuf, } impl From<&Config> for ImageGenerationExtensionConfig { @@ -38,7 +35,6 @@ impl From<&Config> for ImageGenerationExtensionConfig { enabled: config.features.enabled(Feature::ImageGenExt) && config.model_provider.is_openai(), provider: config.model_provider.clone(), - codex_home: config.codex_home.to_path_buf(), } } } @@ -80,13 +76,9 @@ impl ToolContributor for ImageGenerationExtension { return Vec::new(); } - vec![Arc::new(ImageGenerationTool::new( - CodexImagesBackend::new(create_model_provider( - config.provider.clone(), - Some(self.auth_manager.clone()), - )), - generated_image_output_dir(&config.codex_home, thread_store.level_id()), - ))] + vec![Arc::new(ImageGenerationTool::new(CodexImagesBackend::new( + create_model_provider(config.provider.clone(), Some(self.auth_manager.clone())), + )))] } } diff --git a/codex-rs/ext/image-generation/src/tests.rs b/codex-rs/ext/image-generation/src/tests.rs index 1b56270ed3..f866a9018b 100644 --- a/codex-rs/ext/image-generation/src/tests.rs +++ b/codex-rs/ext/image-generation/src/tests.rs @@ -20,14 +20,13 @@ use super::GeneratedImageOutput; use super::ImageRequest; use super::ImagegenAction; use super::ImagegenArgs; -use super::generated_image_output_dir; use super::imagegen_tool_spec; -use super::persist_generated_image; use super::request_for_action; use crate::IMAGE_GEN_NAMESPACE; use crate::IMAGEGEN_TOOL_NAME; const RESULT: &str = "cG5n"; +const OUTPUT_HINT: &str = "Generated images are saved to /tmp as /tmp/call-1.png by default."; #[test] fn uses_reserved_image_gen_namespace() { @@ -55,15 +54,11 @@ fn generate_uses_fixed_request_defaults() { ); } -#[tokio::test] -async fn generated_output_returns_image_input_and_persists_artifact() { - let tempdir = tempfile::tempdir().expect("tempdir"); - let output_hint = persist_generated_image(tempdir.path(), "call-1", RESULT) - .await - .expect("generated image should persist"); +#[test] +fn generated_output_returns_image_input_and_output_hint() { let output = GeneratedImageOutput { result: RESULT.to_string(), - output_hint: Some(output_hint), + output_hint: Some(OUTPUT_HINT.to_string()), }; let ResponseInputItem::FunctionCallOutput { @@ -84,19 +79,10 @@ async fn generated_output_returns_image_input_and_persists_artifact() { detail: Some(DEFAULT_IMAGE_DETAIL), }, FunctionCallOutputContentItem::InputText { - text: format!( - "Generated images are saved to {} as {} by default.\n\ - If you need to use a generated image at another path, copy it and leave the original in place unless the user explicitly asks you to delete it.", - tempdir.path().display(), - tempdir.path().join("call-1.png").display(), - ), + text: OUTPUT_HINT.to_string(), }, ] ); - assert_eq!( - std::fs::read(tempdir.path().join("call-1.png")).expect("saved generated image"), - b"png" - ); } #[test] @@ -265,14 +251,6 @@ fn edit_without_image_history_returns_tool_error() { ); } -#[test] -fn generated_image_output_dir_is_scoped_to_sanitized_thread_id() { - assert_eq!( - generated_image_output_dir(std::path::Path::new("/tmp/codex-home"), "thread/1"), - std::path::PathBuf::from("/tmp/codex-home/generated_images/thread_1") - ); -} - fn args(action: ImagegenAction, prompt: &str) -> ImagegenArgs { ImagegenArgs { prompt: prompt.to_string(), diff --git a/codex-rs/ext/image-generation/src/tool.rs b/codex-rs/ext/image-generation/src/tool.rs index fa1614cd47..75a79db2bf 100644 --- a/codex-rs/ext/image-generation/src/tool.rs +++ b/codex-rs/ext/image-generation/src/tool.rs @@ -1,8 +1,3 @@ -use std::path::Path; -use std::path::PathBuf; - -use base64::Engine; -use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_api::ImageBackground; use codex_api::ImageEditRequest; use codex_api::ImageGenerationRequest; @@ -41,21 +36,16 @@ use crate::backend::CodexImagesBackend; const IMAGE_MODEL: &str = "gpt-image-2"; const MAX_EDIT_IMAGES: usize = 5; const IMAGEGEN_DESCRIPTION: &str = include_str!("../imagegen_description.md"); -const GENERATED_IMAGE_ARTIFACTS_DIR: &str = "generated_images"; #[derive(Clone)] pub(crate) struct ImageGenerationTool { backend: CodexImagesBackend, - output_dir: PathBuf, } impl ImageGenerationTool { /// Creates an image-generation tool backed by an image API executor. - pub(crate) fn new(backend: CodexImagesBackend, output_dir: PathBuf) -> Self { - Self { - backend, - output_dir, - } + pub(crate) fn new(backend: CodexImagesBackend) -> Self { + Self { backend } } } @@ -94,7 +84,6 @@ impl ToolExecutor for ImageGenerationTool { async fn handle(&self, call: ToolCall) -> Result, FunctionCallError> { let args = parse_args(&call)?; let request = request_for_action(&args, call.conversation_history.items())?; - let response = match request { ImageRequest::Generate(request) => self.backend.generate(request).await, ImageRequest::Edit(request) => self.backend.edit(request).await, @@ -107,18 +96,10 @@ impl ToolExecutor for ImageGenerationTool { "image generation returned no image data".to_string(), )); }; - let output_hint = - match persist_generated_image(&self.output_dir, &call.call_id, &result).await { - Ok(output_hint) => Some(output_hint), - Err(err) => { - tracing::warn!( - call_id = %call.call_id, - output_dir = %self.output_dir.display(), - "failed to save generated image: {err}" - ); - None - } - }; + let output_hint = call + .turn_item_emitter + .image_generation_completed(call.call_id.clone(), args.prompt, result.clone()) + .await; Ok(Box::new(GeneratedImageOutput { result, output_hint, @@ -268,58 +249,6 @@ fn parse_args(call: &ToolCall) -> Result { .map_err(|err| FunctionCallError::RespondToModel(err.to_string())) } -/// Resolves where generated images for one thread are persisted by the extension. -pub(crate) fn generated_image_output_dir(codex_home: &Path, thread_id: &str) -> PathBuf { - codex_home - .join(GENERATED_IMAGE_ARTIFACTS_DIR) - .join(sanitize_path_component(thread_id)) -} - -fn generated_image_output_path(output_dir: &Path, call_id: &str) -> PathBuf { - output_dir.join(format!("{}.png", sanitize_path_component(call_id))) -} - -fn sanitize_path_component(value: &str) -> String { - let sanitized: String = value - .chars() - .map(|ch| { - if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' { - ch - } else { - '_' - } - }) - .collect(); - if sanitized.is_empty() { - "generated_image".to_string() - } else { - sanitized - } -} - -async fn persist_generated_image( - output_dir: &Path, - call_id: &str, - result: &str, -) -> Result { - let bytes = BASE64_STANDARD - .decode(result.trim().as_bytes()) - .map_err(|err| format!("invalid image generation payload: {err}"))?; - tokio::fs::create_dir_all(output_dir) - .await - .map_err(|err| err.to_string())?; - tokio::fs::write(generated_image_output_path(output_dir, call_id), bytes) - .await - .map_err(|err| err.to_string())?; - - Ok(format!( - "Generated images are saved to {} as {} by default.\n\ - If you need to use a generated image at another path, copy it and leave the original in place unless the user explicitly asks you to delete it.", - output_dir.display(), - generated_image_output_path(output_dir, call_id).display(), - )) -} - /// Builds the namespace function schema exposed to the model. fn imagegen_tool_spec() -> ToolSpec { let mut schema_value = serde_json::to_value( @@ -369,7 +298,7 @@ impl ToolOutput for GeneratedImageOutput { true } - /// Returns generated bytes and persisted-artifact context for the model's follow-up response. + /// Returns generated bytes and persisted-artifact context for model follow-up. fn to_response_item(&self, call_id: &str, _payload: &ToolPayload) -> ResponseInputItem { let mut content = vec![FunctionCallOutputContentItem::InputImage { image_url: format!("data:image/png;base64,{}", self.result), diff --git a/codex-rs/ext/memories/src/tests.rs b/codex-rs/ext/memories/src/tests.rs index f2d79c850c..ce03d74965 100644 --- a/codex-rs/ext/memories/src/tests.rs +++ b/codex-rs/ext/memories/src/tests.rs @@ -211,6 +211,7 @@ async fn add_ad_hoc_note_tool_creates_note_file() { turn_id: "turn-1".to_string(), call_id: "call-1".to_string(), tool_name: memory_tool_name(crate::ADD_AD_HOC_NOTE_TOOL_NAME), + model: "gpt-test".to_string(), truncation_policy: TruncationPolicy::Bytes(1024), conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), @@ -253,6 +254,7 @@ async fn add_ad_hoc_note_tool_rejects_paths_as_filenames() { turn_id: "turn-1".to_string(), call_id: "call-1".to_string(), tool_name: memory_tool_name(crate::ADD_AD_HOC_NOTE_TOOL_NAME), + model: "gpt-test".to_string(), truncation_policy: TruncationPolicy::Bytes(1024), conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), @@ -296,6 +298,7 @@ async fn read_tool_reads_memory_file() { turn_id: "turn-1".to_string(), call_id: "call-1".to_string(), tool_name: memory_tool_name(crate::READ_TOOL_NAME), + model: "gpt-test".to_string(), truncation_policy: TruncationPolicy::Bytes(1024), conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), @@ -342,6 +345,7 @@ async fn search_tool_accepts_multiple_queries() { turn_id: "turn-1".to_string(), call_id: "call-1".to_string(), tool_name: memory_tool_name(crate::SEARCH_TOOL_NAME), + model: "gpt-test".to_string(), truncation_policy: TruncationPolicy::Bytes(1024), conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), @@ -414,6 +418,7 @@ async fn search_tool_accepts_windowed_all_match_mode() { turn_id: "turn-1".to_string(), call_id: "call-1".to_string(), tool_name: memory_tool_name(crate::SEARCH_TOOL_NAME), + model: "gpt-test".to_string(), truncation_policy: TruncationPolicy::Bytes(1024), conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), @@ -466,6 +471,7 @@ async fn search_tool_rejects_legacy_single_query() { turn_id: "turn-1".to_string(), call_id: "call-1".to_string(), tool_name: memory_tool_name(crate::SEARCH_TOOL_NAME), + model: "gpt-test".to_string(), truncation_policy: TruncationPolicy::Bytes(1024), conversation_history: codex_extension_api::ConversationHistory::default(), turn_item_emitter: Arc::new(NoopTurnItemEmitter), diff --git a/codex-rs/ext/web-search/src/tool.rs b/codex-rs/ext/web-search/src/tool.rs index 2cf25b007c..35b855c5c6 100644 --- a/codex-rs/ext/web-search/src/tool.rs +++ b/codex-rs/ext/web-search/src/tool.rs @@ -90,7 +90,7 @@ impl ToolExecutor for WebSearchTool { ); let request = SearchRequest { id: self.session_id.clone(), - model: None, + model: call.model.clone(), reasoning: None, input: recent_input(call.conversation_history.items()), commands: Some(commands), diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index cfa712c19b..077cd17a35 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2442,12 +2442,22 @@ impl InitialHistory { } } + pub fn get_resumed_session_sources(&self) -> Option<(SessionSource, Option)> { + let meta = self.get_resumed_session_meta()?; + Some((meta.source.clone(), meta.thread_source)) + } + pub fn get_resumed_thread_source(&self) -> Option { + self.get_resumed_session_meta() + .and_then(|meta| meta.thread_source) + } + + fn get_resumed_session_meta(&self) -> Option<&SessionMeta> { match self { InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => None, InitialHistory::Resumed(resumed) => { resumed.history.iter().find_map(|item| match item { - RolloutItem::SessionMeta(meta_line) => meta_line.meta.thread_source, + RolloutItem::SessionMeta(meta_line) => Some(&meta_line.meta), _ => None, }) } @@ -2658,6 +2668,19 @@ impl SessionSource { .restriction_product() .is_some_and(|product| product.matches_product_restriction(products)) } + + pub fn parent_thread_id(&self) -> Option { + match self { + SessionSource::SubAgent(subagent_source) => subagent_source.parent_thread_id(), + SessionSource::Cli + | SessionSource::VSCode + | SessionSource::Exec + | SessionSource::Mcp + | SessionSource::Custom(_) + | SessionSource::Internal(_) + | SessionSource::Unknown => None, + } + } } impl fmt::Display for SubAgentSource { @@ -2678,6 +2701,30 @@ impl fmt::Display for SubAgentSource { } } +impl SubAgentSource { + pub fn kind(&self) -> &str { + match self { + SubAgentSource::Review => "review", + SubAgentSource::Compact => "compact", + SubAgentSource::ThreadSpawn { .. } => "thread_spawn", + SubAgentSource::MemoryConsolidation => "memory_consolidation", + SubAgentSource::Other(other) => other, + } + } + + pub fn parent_thread_id(&self) -> Option { + match self { + SubAgentSource::ThreadSpawn { + parent_thread_id, .. + } => Some(*parent_thread_id), + SubAgentSource::Review + | SubAgentSource::Compact + | SubAgentSource::MemoryConsolidation + | SubAgentSource::Other(_) => None, + } + } +} + impl fmt::Display for InternalSessionSource { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index 0a5643aabf..6f98f78920 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -28,6 +28,7 @@ use reqwest::header::CONTENT_TYPE; use reqwest::header::HeaderMap; use reqwest::header::HeaderName; use rmcp::model::ClientJsonRpcMessage; +use rmcp::model::JsonRpcMessage; use rmcp::model::ServerJsonRpcMessage; use rmcp::transport::streamable_http_client::AuthRequiredError; use rmcp::transport::streamable_http_client::InsufficientScopeError; @@ -88,6 +89,8 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { auth_token: Option, custom_headers: HashMap, ) -> std::result::Result> { + let (mcp_method, mcp_request_id) = client_jsonrpc_message_fields(&message); + let has_session_id = session_id.is_some(); let mut headers = self.default_headers.clone(); headers.extend(custom_headers); self.add_auth_headers(&mut headers); @@ -121,7 +124,8 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { } let body = serde_json::to_vec(&message).map_err(StreamableHttpError::Deserialize)?; - let (response, mut body_stream) = self + let has_authorization_header = headers.contains_key(AUTHORIZATION); + let response = self .http_client .http_request_stream(HttpRequestParams { method: "POST".to_string(), @@ -132,9 +136,22 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { request_id: "buffered-request".to_string(), stream_response: true, }) - .await - .map_err(StreamableHttpClientAdapterError::from) - .map_err(StreamableHttpError::Client)?; + .await; + let (response, mut body_stream) = match response { + Ok(response) => response, + Err(error) => { + log_post_message_http_error( + &uri, + mcp_method.as_deref(), + mcp_request_id.as_deref(), + has_session_id, + has_authorization_header, + ); + return Err(StreamableHttpError::Client( + StreamableHttpClientAdapterError::from(error), + )); + } + }; if response.status == StatusCode::NOT_FOUND.as_u16() && session_id.is_some() { return Err(StreamableHttpError::Client( @@ -354,6 +371,50 @@ fn body_preview(body: impl Into) -> String { body_preview } +fn client_jsonrpc_message_fields( + message: &ClientJsonRpcMessage, +) -> (Option, Option) { + match message { + JsonRpcMessage::Request(request) => ( + Some(request.request.method().to_string()), + Some(request.id.to_string()), + ), + JsonRpcMessage::Response(response) => (None, Some(response.id.to_string())), + JsonRpcMessage::Notification(_) => (None, None), + JsonRpcMessage::Error(error) => (None, error.id.as_ref().map(ToString::to_string)), + } +} + +fn log_post_message_http_error( + uri: &str, + mcp_method: Option<&str>, + mcp_request_id: Option<&str>, + has_session_id: bool, + has_authorization_header: bool, +) { + let parsed_url = reqwest::Url::parse(uri).ok(); + tracing::warn!( + endpoint_scheme = parsed_url + .as_ref() + .map(reqwest::Url::scheme) + .unwrap_or(""), + endpoint_host = parsed_url + .as_ref() + .and_then(reqwest::Url::host_str) + .unwrap_or(""), + endpoint_path = parsed_url + .as_ref() + .map(reqwest::Url::path) + .unwrap_or(""), + endpoint_has_query = parsed_url.as_ref().is_some_and(|url| url.query().is_some()), + mcp_method = mcp_method.unwrap_or(""), + mcp_request_id = mcp_request_id.unwrap_or(""), + has_session_id = has_session_id, + has_authorization_header = has_authorization_header, + "streamable HTTP post_message failed" + ); +} + fn insert_header( headers: &mut HeaderMap, name: HeaderName, diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index 8384c00d6b..1e0bfc9279 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -77,7 +77,8 @@ fn apply_turn_context(metadata: &mut ThreadMetadata, turn_ctx: &TurnContextItem) } metadata.model = Some(turn_ctx.model.clone()); metadata.reasoning_effort = turn_ctx.effort; - metadata.sandbox_policy = enum_to_string(&turn_ctx.sandbox_policy); + metadata.sandbox_policy = + serde_json::to_string(&turn_ctx.permission_profile()).unwrap_or_default(); metadata.approval_mode = enum_to_string(&turn_ctx.approval_policy); } @@ -158,6 +159,7 @@ mod tests { use chrono::Utc; use codex_protocol::ThreadId; use codex_protocol::models::ContentItem; + use codex_protocol::models::PermissionProfile; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::AskForApproval; @@ -366,11 +368,46 @@ mod tests { assert_eq!(metadata.cwd, PathBuf::from("/child/worktree")); assert_eq!( metadata.sandbox_policy, - super::enum_to_string(&SandboxPolicy::DangerFullAccess) + serde_json::to_string(&PermissionProfile::Disabled) + .expect("serialize permission profile") ); assert_eq!(metadata.approval_mode, "never"); } + #[test] + fn turn_context_sets_permission_profile_metadata() { + let mut metadata = metadata_for_test(); + let permission_profile = PermissionProfile::workspace_write(); + + apply_rollout_item( + &mut metadata, + &RolloutItem::TurnContext(TurnContextItem { + turn_id: Some("turn-1".to_string()), + cwd: PathBuf::from("/workspace"), + workspace_roots: None, + current_date: None, + timezone: None, + approval_policy: AskForApproval::OnRequest, + sandbox_policy: SandboxPolicy::DangerFullAccess, + permission_profile: Some(permission_profile.clone()), + network: None, + file_system_sandbox_policy: None, + model: "gpt-5".to_string(), + personality: None, + collaboration_mode: None, + realtime_active: None, + effort: None, + summary: codex_protocol::config_types::ReasoningSummary::Auto, + }), + "test-provider", + ); + + assert_eq!( + metadata.sandbox_policy, + serde_json::to_string(&permission_profile).expect("serialize permission profile") + ); + } + #[test] fn turn_context_sets_cwd_when_session_cwd_missing() { let mut metadata = metadata_for_test(); diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 46cceadf87..8d287bce61 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -977,13 +977,7 @@ pub(super) fn extract_memory_mode(items: &[RolloutItem]) -> Option { fn thread_spawn_parent_thread_id_from_source_str(source: &str) -> Option { let parsed_source = serde_json::from_str(source) .or_else(|_| serde_json::from_value::(Value::String(source.to_string()))); - match parsed_source.ok() { - Some(SessionSource::SubAgent(codex_protocol::protocol::SubAgentSource::ThreadSpawn { - parent_thread_id, - .. - })) => Some(parent_thread_id), - _ => None, - } + parsed_source.ok()?.parent_thread_id() } #[derive(Clone, Copy)] diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index c08d810427..fe6ef88c31 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -8,9 +8,9 @@ use std::sync::OnceLock; use async_trait::async_trait; use chrono::Utc; use codex_protocol::ThreadId; +use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::RolloutItem; -use codex_protocol::protocol::SandboxPolicy; use crate::AppendThreadItemsParams; use crate::ArchiveThreadParams; @@ -366,9 +366,9 @@ fn stored_thread_from_state( approval_mode: metadata .and_then(|metadata| metadata.approval_mode) .unwrap_or(AskForApproval::Never), - sandbox_policy: metadata - .and_then(|metadata| metadata.sandbox_policy.clone()) - .unwrap_or_else(SandboxPolicy::new_read_only_policy), + permission_profile: metadata + .and_then(|metadata| metadata.permission_profile.clone()) + .unwrap_or_else(PermissionProfile::read_only), token_usage: metadata.and_then(|metadata| metadata.token_usage.clone()), first_user_message: metadata.and_then(|metadata| metadata.first_user_message.clone()), history, diff --git a/codex-rs/thread-store/src/local/helpers.rs b/codex-rs/thread-store/src/local/helpers.rs index 7d9db73532..b7abb8954f 100644 --- a/codex-rs/thread-store/src/local/helpers.rs +++ b/codex-rs/thread-store/src/local/helpers.rs @@ -9,8 +9,10 @@ use chrono::DateTime; use chrono::Utc; use codex_git_utils::GitSha; use codex_protocol::ThreadId; +use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::GitInfo; +use codex_protocol::protocol::NetworkAccess; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; use codex_rollout::ARCHIVED_SESSIONS_SUBDIR; @@ -141,13 +143,34 @@ pub(super) fn stored_thread_from_rollout_item( agent_path: None, git_info, approval_mode: AskForApproval::OnRequest, - sandbox_policy: SandboxPolicy::new_read_only_policy(), + permission_profile: PermissionProfile::read_only(), token_usage: None, first_user_message: item.first_user_message, history: None, }) } +pub(super) fn permission_profile_from_metadata_value(value: &str, cwd: &Path) -> PermissionProfile { + serde_json::from_str::(value) + .or_else(|_| { + parse_legacy_sandbox_policy(value) + .map(|policy| PermissionProfile::from_legacy_sandbox_policy_for_cwd(&policy, cwd)) + }) + .unwrap_or_else(|_| PermissionProfile::read_only()) +} + +pub(super) fn permission_profile_to_metadata_value( + permission_profile: &PermissionProfile, +) -> String { + match serde_json::to_string(permission_profile) { + Ok(value) => value, + Err(err) => { + tracing::warn!("failed to serialize permission profile metadata: {err}"); + String::new() + } + } +} + pub(super) fn distinct_thread_metadata_title(metadata: &ThreadMetadata) -> Option { let title = metadata.title.trim(); if title.is_empty() || metadata.first_user_message.as_deref().map(str::trim) == Some(title) { @@ -170,6 +193,20 @@ fn parse_rfc3339(value: Option<&str>) -> Option> { .map(|dt| dt.with_timezone(&Utc)) } +fn parse_legacy_sandbox_policy(value: &str) -> serde_json::Result { + serde_json::from_str(value) + .or_else(|_| serde_json::from_value(serde_json::Value::String(value.to_string()))) + .or_else(|_| match value { + "danger-full-access" => Ok(SandboxPolicy::DangerFullAccess), + "read-only" => Ok(SandboxPolicy::new_read_only_policy()), + "workspace-write" => Ok(SandboxPolicy::new_workspace_write_policy()), + "external-sandbox" => Ok(SandboxPolicy::ExternalSandbox { + network_access: NetworkAccess::Restricted, + }), + _ => serde_json::from_value(serde_json::Value::String(value.to_string())), + }) +} + pub(super) fn git_info_from_parts( sha: Option, branch: Option, diff --git a/codex-rs/thread-store/src/local/read_thread.rs b/codex-rs/thread-store/src/local/read_thread.rs index a9d305e9dc..9f59c72ee3 100644 --- a/codex-rs/thread-store/src/local/read_thread.rs +++ b/codex-rs/thread-store/src/local/read_thread.rs @@ -1,7 +1,7 @@ use chrono::DateTime; use chrono::Utc; +use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; -use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_rollout::RolloutRecorder; @@ -15,6 +15,7 @@ use codex_state::ThreadMetadata; use super::LocalThreadStore; use super::helpers::distinct_thread_metadata_title; use super::helpers::git_info_from_parts; +use super::helpers::permission_profile_from_metadata_value; use super::helpers::rollout_path_is_archived; use super::helpers::set_thread_name_from_title; use super::helpers::stored_thread_from_rollout_item; @@ -45,6 +46,7 @@ pub(super) async fn read_thread( ) .await) { + let metadata_sandbox_policy = metadata.sandbox_policy.clone(); let mut thread = stored_thread_from_sqlite_metadata(store, metadata).await; if !params.include_history && let Some(rollout_path) = thread.rollout_path.clone() @@ -57,6 +59,10 @@ pub(super) async fn read_thread( rollout_thread.name = thread.name; } rollout_thread.git_info = thread.git_info; + rollout_thread.permission_profile = permission_profile_from_metadata_value( + &metadata_sandbox_policy, + rollout_thread.cwd.as_path(), + ); thread = rollout_thread; } attach_history_if_requested(&mut thread, params.include_history).await?; @@ -290,6 +296,8 @@ async fn stored_thread_from_sqlite_metadata( .clone() .or_else(|| metadata.first_user_message.clone()) .unwrap_or_default(); + let permission_profile = + permission_profile_from_metadata_value(&metadata.sandbox_policy, metadata.cwd.as_path()); StoredThread { thread_id: metadata.id, rollout_path: Some(metadata.rollout_path), @@ -320,10 +328,7 @@ async fn stored_thread_from_sqlite_metadata( metadata.git_origin_url, ), approval_mode: parse_or_default(&metadata.approval_mode, AskForApproval::OnRequest), - sandbox_policy: parse_or_default( - &metadata.sandbox_policy, - SandboxPolicy::new_read_only_policy(), - ), + permission_profile, token_usage: None, first_user_message: metadata.first_user_message, history: None, @@ -383,7 +388,7 @@ fn stored_thread_from_meta_line( agent_path: meta_line.meta.agent_path, git_info: meta_line.git, approval_mode: AskForApproval::OnRequest, - sandbox_policy: SandboxPolicy::new_read_only_policy(), + permission_profile: PermissionProfile::read_only(), token_usage: None, first_user_message: None, history: None, @@ -418,6 +423,7 @@ mod tests { use chrono::Utc; use codex_protocol::ThreadId; + use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; use codex_state::ThreadMetadataBuilder; use pretty_assertions::assert_eq; @@ -677,6 +683,84 @@ mod tests { assert_eq!(thread.name, Some("Saved title".to_string())); } + #[tokio::test] + async fn read_thread_returns_permission_profile_from_sqlite_metadata() { + let home = TempDir::new().expect("temp dir"); + let config = test_config(home.path()); + let uuid = Uuid::from_u128(225); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); + let rollout_path = + write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); + let runtime = codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.default_model_provider_id.clone(), + ) + .await + .expect("state db should initialize"); + let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let mut builder = + ThreadMetadataBuilder::new(thread_id, rollout_path, Utc::now(), SessionSource::Cli); + builder.model_provider = Some(config.default_model_provider_id.clone()); + builder.cwd = home.path().to_path_buf(); + let mut metadata = builder.build(config.default_model_provider_id.as_str()); + metadata.sandbox_policy = + serde_json::to_string(&PermissionProfile::Disabled).expect("serialize profile"); + runtime + .upsert_thread(&metadata) + .await + .expect("state db upsert should succeed"); + + let thread = store + .read_thread(ReadThreadParams { + thread_id, + include_archived: false, + include_history: false, + }) + .await + .expect("read thread"); + + assert_eq!(thread.preview, "Hello from user"); + assert_eq!(thread.permission_profile, PermissionProfile::Disabled); + } + + #[tokio::test] + async fn read_thread_accepts_legacy_sandbox_policy_metadata() { + let home = TempDir::new().expect("temp dir"); + let config = test_config(home.path()); + let uuid = Uuid::from_u128(226); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); + let rollout_path = + write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file"); + let runtime = codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.default_model_provider_id.clone(), + ) + .await + .expect("state db should initialize"); + let store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); + let mut builder = + ThreadMetadataBuilder::new(thread_id, rollout_path, Utc::now(), SessionSource::Cli); + builder.model_provider = Some(config.default_model_provider_id.clone()); + builder.cwd = home.path().to_path_buf(); + let mut metadata = builder.build(config.default_model_provider_id.as_str()); + metadata.sandbox_policy = "danger-full-access".to_string(); + runtime + .upsert_thread(&metadata) + .await + .expect("state db upsert should succeed"); + + let thread = store + .read_thread(ReadThreadParams { + thread_id, + include_archived: false, + include_history: true, + }) + .await + .expect("read thread"); + + assert_eq!(thread.permission_profile, PermissionProfile::Disabled); + } + #[tokio::test] async fn read_thread_preserves_rollout_cwd_when_sqlite_metadata_exists() { let home = TempDir::new().expect("temp dir"); @@ -731,6 +815,7 @@ mod tests { let mut metadata = builder.build(config.default_model_provider_id.as_str()); metadata.title = "Saved title".to_string(); metadata.first_user_message = Some("Hello from sqlite".to_string()); + metadata.sandbox_policy = "workspace-write".to_string(); runtime .upsert_thread(&metadata) .await @@ -751,6 +836,19 @@ mod tests { assert_eq!(thread.name, Some("Saved title".to_string())); assert_eq!(thread.model_provider, "rollout-provider"); assert_eq!(thread.cwd, rollout_cwd); + let legacy_policy = SandboxPolicy::WorkspaceWrite { + writable_roots: Vec::new(), + network_access: false, + exclude_tmpdir_env_var: false, + exclude_slash_tmp: false, + }; + assert_eq!( + thread.permission_profile, + PermissionProfile::from_legacy_sandbox_policy_for_cwd( + &legacy_policy, + rollout_cwd.as_path() + ) + ); } #[tokio::test] diff --git a/codex-rs/thread-store/src/local/update_thread_metadata.rs b/codex-rs/thread-store/src/local/update_thread_metadata.rs index ef5bfd7c5d..0b242c0c72 100644 --- a/codex-rs/thread-store/src/local/update_thread_metadata.rs +++ b/codex-rs/thread-store/src/local/update_thread_metadata.rs @@ -18,6 +18,7 @@ use tracing::warn; use super::LocalThreadStore; use super::helpers::git_info_from_parts; +use super::helpers::permission_profile_to_metadata_value; use super::live_writer; use crate::GitInfoPatch; use crate::ReadThreadParams; @@ -279,8 +280,8 @@ async fn apply_metadata_update( if let Some(approval_mode) = patch.approval_mode { metadata.approval_mode = enum_to_string(&approval_mode); } - if let Some(sandbox_policy) = patch.sandbox_policy { - metadata.sandbox_policy = enum_to_string(&sandbox_policy); + if let Some(permission_profile) = patch.permission_profile { + metadata.sandbox_policy = permission_profile_to_metadata_value(&permission_profile); } if let Some(token_usage) = patch.token_usage { metadata.tokens_used = token_usage.total_tokens.max(0); @@ -387,7 +388,7 @@ fn has_observed_metadata_facts(patch: &ThreadMetadataPatch) -> bool { || patch.cwd.is_some() || patch.cli_version.is_some() || patch.approval_mode.is_some() - || patch.sandbox_policy.is_some() + || patch.permission_profile.is_some() || patch.token_usage.is_some() || patch.first_user_message.is_some() } @@ -609,6 +610,7 @@ fn rollout_path_is_archived(store: &LocalThreadStore, path: &Path) -> bool { #[cfg(test)] mod tests { + use codex_protocol::models::PermissionProfile; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; @@ -850,6 +852,45 @@ mod tests { ); } + #[tokio::test] + async fn update_thread_metadata_sets_permission_profile() { + let home = TempDir::new().expect("temp dir"); + let config = test_config(home.path()); + let runtime = codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.default_model_provider_id.clone(), + ) + .await + .expect("state db should initialize"); + let store = LocalThreadStore::new(config, Some(runtime.clone())); + let uuid = Uuid::from_u128(317); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); + write_session_file(home.path(), "2025-01-03T20-30-00", uuid).expect("session file"); + + let thread = store + .update_thread_metadata(UpdateThreadMetadataParams { + thread_id, + patch: ThreadMetadataPatch { + permission_profile: Some(PermissionProfile::Disabled), + ..Default::default() + }, + include_archived: false, + }) + .await + .expect("set permission profile"); + + assert_eq!(thread.permission_profile, PermissionProfile::Disabled); + let metadata = runtime + .get_thread(thread_id) + .await + .expect("sqlite metadata read") + .expect("sqlite metadata"); + assert_eq!( + metadata.sandbox_policy, + serde_json::to_string(&PermissionProfile::Disabled).expect("serialize profile") + ); + } + #[tokio::test] async fn update_thread_metadata_partially_updates_git_info() { let home = TempDir::new().expect("temp dir"); diff --git a/codex-rs/thread-store/src/thread_metadata_sync.rs b/codex-rs/thread-store/src/thread_metadata_sync.rs index cf8dc65320..fcea56fc36 100644 --- a/codex-rs/thread-store/src/thread_metadata_sync.rs +++ b/codex-rs/thread-store/src/thread_metadata_sync.rs @@ -236,7 +236,7 @@ impl ThreadMetadataSync { update.model = Some(turn_ctx.model.clone()); update.reasoning_effort = turn_ctx.effort; update.approval_mode = Some(turn_ctx.approval_policy); - update.sandbox_policy = Some(turn_ctx.sandbox_policy.clone()); + update.permission_profile = Some(turn_ctx.permission_profile()); } RolloutItem::EventMsg(EventMsg::UserMessage(user)) => { if let Some(preview) = user_message_preview(user) { @@ -356,7 +356,7 @@ fn update_has_metadata_facts(update: &ThreadMetadataPatch) -> bool { || update.cwd.is_some() || update.cli_version.is_some() || update.approval_mode.is_some() - || update.sandbox_policy.is_some() + || update.permission_profile.is_some() || update.token_usage.is_some() || update.first_user_message.is_some() || update.git_info.is_some() diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index d9706e64d1..ac29234086 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -5,12 +5,12 @@ use chrono::Utc; use codex_protocol::ThreadId; use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::BaseInstructions; +use codex_protocol::models::PermissionProfile; use codex_protocol::openai_models::MultiAgentVersion; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::GitInfo; use codex_protocol::protocol::RolloutItem; -use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadMemoryMode as MemoryMode; use codex_protocol::protocol::ThreadSource; @@ -401,8 +401,8 @@ pub struct StoredThread { pub git_info: Option, /// Approval mode captured for the thread. pub approval_mode: AskForApproval, - /// Sandbox policy captured for the thread. - pub sandbox_policy: SandboxPolicy, + /// Canonical runtime permissions captured for the thread. + pub permission_profile: PermissionProfile, /// Last observed token usage. pub token_usage: Option, /// First user message observed for this thread, if any. @@ -531,8 +531,8 @@ pub struct ThreadMetadataPatch { pub cli_version: Option, /// Approval mode. pub approval_mode: Option, - /// Sandbox policy. - pub sandbox_policy: Option, + /// Canonical runtime permissions. + pub permission_profile: Option, /// Last observed token usage. pub token_usage: Option, /// First user message observed for this thread. @@ -604,8 +604,8 @@ impl ThreadMetadataPatch { if next.approval_mode.is_some() { self.approval_mode = next.approval_mode; } - if next.sandbox_policy.is_some() { - self.sandbox_policy = next.sandbox_policy; + if next.permission_profile.is_some() { + self.permission_profile = next.permission_profile; } if next.token_usage.is_some() { self.token_usage = next.token_usage; @@ -642,7 +642,7 @@ impl ThreadMetadataPatch { && self.cwd.is_none() && self.cli_version.is_none() && self.approval_mode.is_none() - && self.sandbox_policy.is_none() + && self.permission_profile.is_none() && self.token_usage.is_none() && self.first_user_message.is_none() && self.git_info.is_none() diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index 91fa5830a3..e09526048d 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -63,6 +63,7 @@ pub use responses_api::mcp_tool_to_responses_api_tool; pub use responses_api::tool_definition_to_responses_api_tool; pub use tool_call::ConversationHistory; pub use tool_call::ExtensionTurnItem; +pub use tool_call::ImageGenerationCompletionFuture; pub use tool_call::NoopTurnItemEmitter; pub use tool_call::ToolCall; pub use tool_call::TurnItemEmissionFuture; diff --git a/codex-rs/tools/src/tool_call.rs b/codex-rs/tools/src/tool_call.rs index 3c04e216a6..e2434ccbe0 100644 --- a/codex-rs/tools/src/tool_call.rs +++ b/codex-rs/tools/src/tool_call.rs @@ -29,6 +29,10 @@ impl ConversationHistory { /// Future returned when an extension tool emits a visible turn-item lifecycle event. pub type TurnItemEmissionFuture<'a> = Pin + Send + 'a>>; +/// Future returned when an image-generation extension publishes completed image bytes. +pub type ImageGenerationCompletionFuture<'a> = + Pin> + Send + 'a>>; + /// Visible turn items that an extension fully owns and may emit as-is. /// /// Add only item kinds that require no additional host finalization before @@ -48,6 +52,19 @@ pub trait TurnItemEmitter: Send + Sync { /// Emits the completion of one visible turn item. fn emit_completed<'a>(&'a self, item: ExtensionTurnItem) -> TurnItemEmissionFuture<'a>; + + /// Publishes image bytes for host persistence and visible completion. + /// + /// Returns persisted-artifact context for the extension's model-facing + /// function output when the host saves the generated image successfully. + fn image_generation_completed<'a>( + &'a self, + _call_id: String, + _prompt: String, + _result: String, + ) -> ImageGenerationCompletionFuture<'a> { + Box::pin(std::future::ready(None)) + } } /// Turn-item emitter used when a caller does not expose visible item emission. @@ -70,6 +87,7 @@ pub struct ToolCall { pub turn_id: String, pub call_id: String, pub tool_name: ToolName, + pub model: String, pub truncation_policy: TruncationPolicy, pub conversation_history: ConversationHistory, pub turn_item_emitter: Arc, @@ -82,6 +100,7 @@ impl std::fmt::Debug for ToolCall { .field("turn_id", &self.turn_id) .field("call_id", &self.call_id) .field("tool_name", &self.tool_name) + .field("model", &self.model) .field("truncation_policy", &self.truncation_policy) .field("conversation_history", &self.conversation_history) .field("turn_item_emitter", &"") diff --git a/codex-rs/tui/src/app/config_persistence.rs b/codex-rs/tui/src/app/config_persistence.rs index 8c80230cdf..92ed5d470c 100644 --- a/codex-rs/tui/src/app/config_persistence.rs +++ b/codex-rs/tui/src/app/config_persistence.rs @@ -398,7 +398,7 @@ impl App { serde_json::json!(auto_review_preset.approvals_reviewer.to_string()), )); if previous_approvals_reviewer != auto_review_preset.approvals_reviewer { - permissions_history_label = Some("Auto-review"); + permissions_history_label = Some("Approve for me"); } } else if !effective_enabled { feature_edits.push(crate::config_update::clear_config_value( @@ -406,7 +406,7 @@ impl App { )); feature_config.approvals_reviewer = ApprovalsReviewer::User; if previous_approvals_reviewer != ApprovalsReviewer::User { - permissions_history_label = Some("Default"); + permissions_history_label = Some("Ask for approval"); } } approvals_reviewer_override = Some(feature_config.approvals_reviewer); @@ -419,7 +419,7 @@ impl App { if !self.try_set_approval_policy_on_config( &mut feature_config, auto_review_preset.approval_policy, - "Failed to enable Auto-review", + "Failed to enable Approve for me", "failed to set auto-review approval policy on staged config", ) { continue; @@ -428,7 +428,7 @@ impl App { .try_set_builtin_active_permission_profile_on_config( &mut feature_config, auto_review_preset.active_permission_profile.clone(), - "Failed to enable Auto-review", + "Failed to enable Approve for me", "failed to set auto-review permission profile on staged config", ) else { @@ -546,7 +546,7 @@ impl App { "failed to set auto-review permission profile on chat config" ); self.chat_widget - .add_error_message(format!("Failed to enable Auto-review: {err}")); + .add_error_message(format!("Failed to enable Approve for me: {err}")); } if permission_profile_override.is_some() { self.runtime_permission_profile_override = @@ -818,7 +818,7 @@ impl App { "failed to sync effective approval policy after an overridden write" ); self.chat_widget.add_error_message(format!( - "Failed to refresh overridden Auto-review settings: {err}" + "Failed to refresh overridden Approve for me settings: {err}" )); } else { self.chat_widget.set_approval_policy(policy); @@ -846,7 +846,7 @@ impl App { let Some(permission_profile) = self.try_set_builtin_active_permission_profile_on_config( &mut config, auto_review_preset.active_permission_profile.clone(), - "Failed to refresh overridden Auto-review settings", + "Failed to refresh overridden Approve for me settings", "failed to sync overridden Auto-review permission profile", ) else { return; @@ -864,7 +864,7 @@ impl App { "failed to sync overridden Auto-review permission profile on chat config" ); self.chat_widget.add_error_message(format!( - "Failed to refresh overridden Auto-review settings: {err}" + "Failed to refresh overridden Approve for me settings: {err}" )); return; } diff --git a/codex-rs/tui/src/app/event_dispatch.rs b/codex-rs/tui/src/app/event_dispatch.rs index 417844acc6..f9683e26ea 100644 --- a/codex-rs/tui/src/app/event_dispatch.rs +++ b/codex-rs/tui/src/app/event_dispatch.rs @@ -121,6 +121,9 @@ impl App { } } } + AppEvent::ArchiveCurrentThread => { + return Ok(self.archive_current_thread(app_server).await); + } AppEvent::ForkCurrentSession => { self.session_telemetry.counter( "codex.thread.fork", @@ -2146,4 +2149,31 @@ impl App { } } } + + pub(super) async fn archive_current_thread( + &mut self, + app_server: &mut AppServerSession, + ) -> AppRunControl { + let Some(thread_id) = self.active_thread_id.or(self.chat_widget.thread_id()) else { + self.chat_widget + .add_error_message("A thread must start before it can be archived.".to_string()); + return AppRunControl::Continue; + }; + if self.side_threads.contains_key(&thread_id) { + self.chat_widget.add_error_message( + "'/archive' is unavailable in side conversations. Press Ctrl+C to return to the main thread first." + .to_string(), + ); + return AppRunControl::Continue; + } + + match app_server.thread_archive(thread_id).await { + Ok(()) => AppRunControl::Exit(ExitReason::UserRequested), + Err(err) => { + self.chat_widget + .add_error_message(format!("Failed to archive current thread: {err}")); + AppRunControl::Continue + } + } + } } diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index 3fa71d7b45..993772ab77 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -1820,7 +1820,7 @@ async fn update_feature_flags_enabling_guardian_selects_auto_review() -> Result< .map(|line| line.to_string()) .collect::>() .join("\n"); - assert!(rendered.contains("Permissions updated to Auto-review")); + assert!(rendered.contains("Permissions updated to Approve for me")); let config = std::fs::read_to_string(codex_home.path().join("config.toml"))?; assert!(config.contains("guardian_approval = true")); @@ -1915,7 +1915,7 @@ async fn update_feature_flags_disabling_guardian_clears_review_policy_and_restor .map(|line| line.to_string()) .collect::>() .join("\n"); - assert!(rendered.contains("Permissions updated to Default")); + assert!(rendered.contains("Permissions updated to Ask for approval")); let config = std::fs::read_to_string(codex_home.path().join("config.toml"))?; assert!(!config.contains("guardian_approval = true")); diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index b37aee64e6..f1e612dd58 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -209,6 +209,9 @@ pub(crate) enum AppEvent { /// Resume a thread by UUID or thread name inside the running TUI session. ResumeSessionByIdOrName(String), + /// Archive the current active main thread and exit after it succeeds. + ArchiveCurrentThread, + /// Fork the current session into a new thread. ForkCurrentSession, diff --git a/codex-rs/tui/src/app_server_session.rs b/codex-rs/tui/src/app_server_session.rs index 2b7b1cd885..e66b5ee8a2 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -46,6 +46,8 @@ use codex_app_server_protocol::SkillsListResponse; use codex_app_server_protocol::Thread; use codex_app_server_protocol::ThreadApproveGuardianDeniedActionParams; use codex_app_server_protocol::ThreadApproveGuardianDeniedActionResponse; +use codex_app_server_protocol::ThreadArchiveParams; +use codex_app_server_protocol::ThreadArchiveResponse; use codex_app_server_protocol::ThreadBackgroundTerminalsCleanParams; use codex_app_server_protocol::ThreadBackgroundTerminalsCleanResponse; use codex_app_server_protocol::ThreadCompactStartParams; @@ -564,6 +566,21 @@ impl AppServerSession { Ok(response.thread) } + pub(crate) async fn thread_archive(&mut self, thread_id: ThreadId) -> Result<()> { + let request_id = self.next_request_id(); + let _: ThreadArchiveResponse = self + .client + .request_typed(ClientRequest::ThreadArchive { + request_id, + params: ThreadArchiveParams { + thread_id: thread_id.to_string(), + }, + }) + .await + .wrap_err("thread/archive failed in TUI")?; + Ok(()) + } + pub(crate) async fn thread_metadata_update_branch( &mut self, thread_id: ThreadId, diff --git a/codex-rs/tui/src/bottom_pane/chat_composer.rs b/codex-rs/tui/src/bottom_pane/chat_composer.rs index 721824ed61..ec3284866d 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer.rs @@ -7885,6 +7885,32 @@ mod tests { insta::assert_snapshot!("slash_popup_res", terminal.backend()); } + #[test] + fn slash_popup_archive_for_ar_ui() { + use ratatui::Terminal; + use ratatui::backend::TestBackend; + + let (tx, _rx) = unbounded_channel::(); + let sender = AppEventSender::new(tx); + + let mut composer = ChatComposer::new( + /*has_input_focus*/ true, + sender, + /*enhanced_keys_supported*/ false, + "Ask Codex to do anything".to_string(), + /*disable_paste_burst*/ false, + ); + + type_chars_humanlike(&mut composer, &['/', 'a', 'r']); + + let mut terminal = Terminal::new(TestBackend::new(60, 5)).expect("terminal"); + terminal + .draw(|f| composer.render(f.area(), f.buffer_mut())) + .expect("draw composer"); + + insta::assert_snapshot!("slash_popup_ar", terminal.backend()); + } + #[test] fn slash_popup_resume_for_res_logic() { use super::super::command_popup::CommandItem; diff --git a/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__chat_composer__tests__slash_popup_ar.snap b/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__chat_composer__tests__slash_popup_ar.snap new file mode 100644 index 0000000000..5014c35d28 --- /dev/null +++ b/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__chat_composer__tests__slash_popup_ar.snap @@ -0,0 +1,9 @@ +--- +source: tui/src/bottom_pane/chat_composer.rs +expression: terminal.backend() +--- +" " +"› /ar " +" " +" " +" /archive archive this session and exit " diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 7d879a09c5..8604229552 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -467,7 +467,9 @@ use unicode_segmentation::UnicodeSegmentation; const USER_SHELL_COMMAND_HELP_TITLE: &str = "Prefix a command with ! to run it locally"; const USER_SHELL_COMMAND_HELP_HINT: &str = "Example: !ls"; -const AUTO_REVIEW_DESCRIPTION: &str = "Same workspace-write permissions as Default, but eligible `on-request` approvals are routed through the auto-reviewer subagent."; +const ASK_FOR_APPROVAL_LABEL: &str = "Ask for approval"; +const APPROVE_FOR_ME_LABEL: &str = "Approve for me"; +const AUTO_REVIEW_DESCRIPTION: &str = "Only ask for actions detected as potentially unsafe."; const DEFAULT_OPENAI_BASE_URL: &str = "https://api.openai.com/v1"; const DEFAULT_STATUS_LINE_ITEMS: [&str; 2] = ["model-with-reasoning", "current-dir"]; const MAX_AGENT_COPY_HISTORY: usize = 32; @@ -1525,11 +1527,11 @@ impl ChatWidget { fn rename_confirmation_cell(name: &str, thread_id: Option) -> PlainHistoryCell { let mut line = vec![ "• ".into(), - "Thread renamed to ".into(), + "Session renamed to ".into(), name.to_string().cyan(), ]; if let Some(hint) = resume_hint(Some(name), thread_id) { - line.extend([". To resume this thread run ".into(), hint.cyan()]); + line.extend([". To resume this session run ".into(), hint.cyan()]); } PlainHistoryCell::new(vec![line.into()]) } diff --git a/codex-rs/tui/src/chatwidget/permission_popups.rs b/codex-rs/tui/src/chatwidget/permission_popups.rs index b271422788..c0d492d7aa 100644 --- a/codex-rs/tui/src/chatwidget/permission_popups.rs +++ b/codex-rs/tui/src/chatwidget/permission_popups.rs @@ -56,7 +56,9 @@ impl ChatWidget { continue; } let base_name = if preset.id == "auto" && windows_degraded_sandbox_enabled { - "Default (non-admin sandbox)".to_string() + format!("{ASK_FOR_APPROVAL_LABEL} (non-admin sandbox)") + } else if preset.id == "auto" { + ASK_FOR_APPROVAL_LABEL.to_string() } else { preset.label.to_string() }; @@ -100,7 +102,7 @@ impl ChatWidget { if guardian_approval_enabled { items.push(SelectionItem { - name: "Auto-review".to_string(), + name: APPROVE_FOR_ME_LABEL.to_string(), description: Some(AUTO_REVIEW_DESCRIPTION.to_string()), is_current: current_review_policy == ApprovalsReviewer::AutoReview && Self::preset_matches_current( @@ -111,7 +113,7 @@ impl ChatWidget { ), actions: self.permission_mode_actions( &preset, - "Auto-review".to_string(), + APPROVE_FOR_ME_LABEL.to_string(), ApprovalsReviewer::AutoReview, /*profile_selection*/ None, /*return_to_permissions*/ !include_read_only, diff --git a/codex-rs/tui/src/chatwidget/permissions_menu.rs b/codex-rs/tui/src/chatwidget/permissions_menu.rs index 7c0477116c..c9cb3b108f 100644 --- a/codex-rs/tui/src/chatwidget/permissions_menu.rs +++ b/codex-rs/tui/src/chatwidget/permissions_menu.rs @@ -94,10 +94,10 @@ impl ChatWidget { approval_policy: AskForApproval, approvals_reviewer: ApprovalsReviewer, ) -> SelectionItem { - let label = if approvals_reviewer == ApprovalsReviewer::AutoReview { - "Auto-review" - } else { - preset.label + let label = match (preset.id, approvals_reviewer) { + ("auto", ApprovalsReviewer::AutoReview) => APPROVE_FOR_ME_LABEL, + ("auto", ApprovalsReviewer::User) => ASK_FOR_APPROVAL_LABEL, + _ => preset.label, }; let active_profile_id = self .config diff --git a/codex-rs/tui/src/chatwidget/slash_dispatch.rs b/codex-rs/tui/src/chatwidget/slash_dispatch.rs index b479718304..99833da8a8 100644 --- a/codex-rs/tui/src/chatwidget/slash_dispatch.rs +++ b/codex-rs/tui/src/chatwidget/slash_dispatch.rs @@ -164,6 +164,35 @@ impl ChatWidget { SlashCommand::New => { self.app_event_tx.send(AppEvent::NewSession); } + SlashCommand::Archive => { + self.bottom_pane.show_selection_view(SelectionViewParams { + title: Some("Archive this session?".to_string()), + subtitle: Some( + "Are you sure? This will archive the current session and exit Codex" + .to_string(), + ), + footer_hint: Some(standard_popup_hint_line()), + items: vec![ + SelectionItem { + name: "No, don't archive".to_string(), + description: Some("Return to the current session".to_string()), + dismiss_on_select: true, + ..Default::default() + }, + SelectionItem { + name: "Yes, archive and exit".to_string(), + description: Some("Archive this session now".to_string()), + actions: vec![Box::new(|tx| { + tx.send(AppEvent::ArchiveCurrentThread); + })], + dismiss_on_select: true, + ..Default::default() + }, + ], + ..Default::default() + }); + self.request_redraw(); + } SlashCommand::Clear => { self.app_event_tx.send(AppEvent::ClearUi); } @@ -949,6 +978,7 @@ impl ChatWidget { | SlashCommand::TestApproval => QueueDrain::Continue, SlashCommand::Feedback | SlashCommand::New + | SlashCommand::Archive | SlashCommand::Clear | SlashCommand::Resume | SlashCommand::Fork diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__approvals_selection_popup.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__approvals_selection_popup.snap index 52ac576149..bfa5173474 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__approvals_selection_popup.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__approvals_selection_popup.snap @@ -4,11 +4,11 @@ expression: popup --- Update Model Permissions -› 1. Default Codex can read and edit files in the current workspace, and - run commands. Approval is required to access the internet or - edit other files. - 2. Full Access Codex can edit files outside this workspace and access the - internet without asking for approval. Exercise caution when - using. +› 1. Ask for approval Codex can read and edit files in the current workspace, + and run commands. Approval is required to access the + internet or edit other files. + 2. Full Access Codex can edit files outside this workspace and access + the internet without asking for approval. Exercise + caution when using. Press enter to confirm or esc to go back diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__approvals_selection_popup@windows.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__approvals_selection_popup@windows.snap index 9c512c7781..abd2450bf7 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__approvals_selection_popup@windows.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__approvals_selection_popup@windows.snap @@ -7,7 +7,7 @@ expression: popup › 1. Read Only (current) Codex can read files in the current workspace. Approval is required to edit files or access the internet. - 2. Default Codex can read and edit files in the current + 2. Ask for approval Codex can read and edit files in the current workspace, and run commands. Approval is required to access the internet or edit other files. 3. Full Access Codex can edit files outside this workspace and diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__approvals_selection_popup@windows_degraded.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__approvals_selection_popup@windows_degraded.snap index ecbe5de157..368e732855 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__approvals_selection_popup@windows_degraded.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__approvals_selection_popup@windows_degraded.snap @@ -5,17 +5,17 @@ expression: popup --- Update Model Permissions -› 1. Read Only (current) Codex can read files in the current - workspace. Approval is required to edit - files or access the internet. - 2. Default (non-admin sandbox) Codex can read and edit files in the - current workspace, and run commands. - Approval is required to access the - internet or edit other files. - 3. Full Access Codex can edit files outside this - workspace and access the internet without - asking for approval. Exercise caution - when using. +› 1. Read Only (current) Codex can read files in the current + workspace. Approval is required to + edit files or access the internet. + 2. Ask for approval (non-admin sandbox) Codex can read and edit files in the + current workspace, and run commands. + Approval is required to access the + internet or edit other files. + 3. Full Access Codex can edit files outside this + workspace and access the internet + without asking for approval. Exercise + caution when using. The non-admin sandbox protects your files and prevents network access under most circumstances. However, it carries greater risk if prompt injected. To diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__permissions_selection_history_full_access_to_default.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__permissions_selection_history_full_access_to_default.snap index 135e5b1bfa..3352c2435a 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__permissions_selection_history_full_access_to_default.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__permissions_selection_history_full_access_to_default.snap @@ -1,5 +1,5 @@ --- -source: tui/src/chatwidget/tests.rs +source: tui/src/chatwidget/tests/permissions.rs expression: "lines_to_single_string(&cells[0])" --- -• Permissions updated to Default +• Permissions updated to Ask for approval diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__permissions_selection_history_full_access_to_default@windows.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__permissions_selection_history_full_access_to_default@windows.snap index eb08108564..6e0269a521 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__permissions_selection_history_full_access_to_default@windows.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__permissions_selection_history_full_access_to_default@windows.snap @@ -2,4 +2,4 @@ source: tui/src/chatwidget/tests.rs expression: "lines_to_single_string(&cells[0])" --- -• Permissions updated to Default (non-admin sandbox) +• Permissions updated to Ask for approval (non-admin sandbox) diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__profile_permissions_selection_popup.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__profile_permissions_selection_popup.snap index c0a23815a0..5ad5481590 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__profile_permissions_selection_popup.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__profile_permissions_selection_popup.snap @@ -4,17 +4,17 @@ expression: "render_bottom_popup(&chat, 80)" --- Update Model Permissions -› 1. Default (current) Codex can read and edit files in the current - workspace, and run commands. Approval is required to - access the internet or edit other files. - 2. Auto-review Same workspace-write permissions as Default, but - eligible `on-request` approvals are routed through the - auto-reviewer subagent. - 3. Full Access Codex can edit files outside this workspace and access - the internet without asking for approval. Exercise - caution when using. - 4. Read Only Codex can read files in the current workspace. - Approval is required to edit files or access the - internet. +› 1. Ask for approval (current) Codex can read and edit files in the current + workspace, and run commands. Approval is + required to access the internet or edit other + files. + 2. Approve for me Only ask for actions detected as potentially + unsafe. + 3. Full Access Codex can edit files outside this workspace + and access the internet without asking for + approval. Exercise caution when using. + 4. Read Only Codex can read files in the current + workspace. Approval is required to edit files + or access the internet. Press enter to confirm or esc to go back diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__profile_permissions_selection_popup_with_custom_profiles.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__profile_permissions_selection_popup_with_custom_profiles.snap index 3ca5753795..dac63d4e32 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__profile_permissions_selection_popup_with_custom_profiles.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__profile_permissions_selection_popup_with_custom_profiles.snap @@ -4,12 +4,11 @@ expression: "render_bottom_popup(&chat, 80)" --- Update Model Permissions - 1. Default Codex can read and edit files in the current + 1. Ask for approval Codex can read and edit files in the current workspace, and run commands. Approval is required to access the internet or edit other files. - 2. Auto-review Same workspace-write permissions as Default, but - eligible `on-request` approvals are routed through - the auto-reviewer subagent. + 2. Approve for me Only ask for actions detected as potentially + unsafe. 3. Full Access Codex can edit files outside this workspace and access the internet without asking for approval. Exercise caution when using. diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__slash_archive_confirmation_popup.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__slash_archive_confirmation_popup.snap new file mode 100644 index 0000000000..cdc895d7be --- /dev/null +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__slash_archive_confirmation_popup.snap @@ -0,0 +1,11 @@ +--- +source: tui/src/chatwidget/tests/slash_commands.rs +expression: popup +--- + Archive this session? + Are you sure? This will archive the current session and exit Codex + +› 1. No, don't archive Return to the current session + 2. Yes, archive and exit Archive this session now + + Press enter to confirm or esc to go back diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__status_surface_previews_hardcoded_only.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__status_surface_previews_hardcoded_only.snap index 250d56ff48..28a7863858 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__status_surface_previews_hardcoded_only.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__status_surface_previews_hardcoded_only.snap @@ -2,5 +2,5 @@ source: tui/src/chatwidget/tests/status_surface_previews.rs expression: snapshot --- -status line: my-project · feat/awesome-feature · thread title · Read Only · on-request +status line: my-project · feat/awesome-feature · thread title · Read Only · Ask for approval terminal title: thread title | feat/awesome-feature | Tasks 0/0 diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__thread_name_update_resume_hint.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__thread_name_update_resume_hint.snap index b3cdf89da3..f6f827fe53 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__thread_name_update_resume_hint.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__thread_name_update_resume_hint.snap @@ -2,4 +2,4 @@ source: tui/src/chatwidget/tests/app_server.rs expression: rendered --- -• Thread renamed to review-fix. To resume this thread run codex resume, then select review-fix (123e4567-e89b-12d3-a456-426614174000) +• Session renamed to review-fix. To resume this session run codex resume, then select review-fix (123e4567-e89b-12d3-a456-426614174000) diff --git a/codex-rs/tui/src/chatwidget/status_surfaces.rs b/codex-rs/tui/src/chatwidget/status_surfaces.rs index 2b2b7d52df..6a3bbbc07b 100644 --- a/codex-rs/tui/src/chatwidget/status_surfaces.rs +++ b/codex-rs/tui/src/chatwidget/status_surfaces.rs @@ -1016,13 +1016,14 @@ fn permissions_display(config: &Config) -> String { fn approval_mode_display(config: &Config) -> String { let approval_policy = AskForApproval::from(config.permissions.approval_policy.value()); - if approval_policy == AskForApproval::OnRequest - && config.approvals_reviewer == ApprovalsReviewer::AutoReview - { - "auto-review".to_string() - } else { - config.permissions.approval_policy.value().to_string() + if approval_policy == AskForApproval::OnRequest { + return match config.approvals_reviewer { + ApprovalsReviewer::AutoReview => "Approve for me".to_string(), + ApprovalsReviewer::User => "Ask for approval".to_string(), + }; } + + config.permissions.approval_policy.value().to_string() } fn parse_items_with_invalids(ids: impl IntoIterator) -> (Vec, Vec) diff --git a/codex-rs/tui/src/chatwidget/tests/permissions.rs b/codex-rs/tui/src/chatwidget/tests/permissions.rs index f79fe3defc..20d600b303 100644 --- a/codex-rs/tui/src/chatwidget/tests/permissions.rs +++ b/codex-rs/tui/src/chatwidget/tests/permissions.rs @@ -143,7 +143,7 @@ async fn profile_permissions_selection_emits_named_profile_event_only() { approval_policy: Some(AskForApproval::OnRequest), approvals_reviewer: Some(ApprovalsReviewer::User), display_label, - }) if profile_id == ":workspace" && display_label == "Default" + }) if profile_id == ":workspace" && display_label == "Ask for approval" )); } @@ -208,7 +208,7 @@ async fn profile_permissions_selection_emits_auto_review_mode_event() { approval_policy: Some(AskForApproval::OnRequest), approvals_reviewer: Some(ApprovalsReviewer::AutoReview), display_label, - }) if profile_id == ":workspace" && display_label == "Auto-review" + }) if profile_id == ":workspace" && display_label == "Approve for me" )); } @@ -256,7 +256,7 @@ async fn approvals_selection_popup_snapshot_windows_degraded_sandbox() { let popup = render_bottom_popup(&chat, /*width*/ 80); assert!( - popup.contains("Default (non-admin sandbox)"), + popup.contains("Ask for approval (non-admin sandbox)"), "expected degraded sandbox label in approvals popup: {popup}" ); assert!( @@ -285,7 +285,7 @@ async fn preset_matching_accepts_workspace_write_with_extra_roots() { cwd.as_path(), &preset ), - "WorkspaceWrite with extra roots should still match the Default preset" + "WorkspaceWrite with extra roots should still match the Ask for approval preset" ); assert!( !ChatWidget::preset_matches_current( @@ -629,7 +629,7 @@ async fn permissions_selection_history_snapshot_full_access_to_default() { chat.open_permissions_popup(); let popup = render_bottom_popup(&chat, /*width*/ 120); chat.handle_key_event(KeyEvent::from(KeyCode::Up)); - if popup.contains("Auto-review") { + if popup.contains("Approve for me") { chat.handle_key_event(KeyEvent::from(KeyCode::Up)); } chat.handle_key_event(KeyEvent::from(KeyCode::Enter)); @@ -699,8 +699,8 @@ async fn permissions_selection_hides_auto_review_when_feature_disabled() { let popup = render_bottom_popup(&chat, /*width*/ 120); assert!( - !popup.contains("Auto-review"), - "expected Auto-review to stay hidden until the feature is enabled: {popup}" + !popup.contains("Approve for me"), + "expected Approve for me to stay hidden until the feature is enabled: {popup}" ); } @@ -730,8 +730,8 @@ async fn permissions_selection_hides_auto_review_when_feature_disabled_even_if_a let popup = render_bottom_popup(&chat, /*width*/ 120); assert!( - !popup.contains("Auto-review"), - "expected Auto-review to stay hidden when the feature is disabled: {popup}" + !popup.contains("Approve for me"), + "expected Approve for me to stay hidden when the feature is disabled: {popup}" ); } @@ -776,8 +776,8 @@ async fn permissions_selection_marks_auto_review_current_after_session_configure let popup = render_bottom_popup(&chat, /*width*/ 120); assert!( - popup.contains("Auto-review (current)"), - "expected Auto-review to be current after SessionConfigured sync: {popup}" + popup.contains("Approve for me (current)"), + "expected Approve for me to be current after SessionConfigured sync: {popup}" ); } @@ -826,8 +826,8 @@ async fn permissions_selection_marks_auto_review_current_with_custom_workspace_w let popup = render_bottom_popup(&chat, /*width*/ 120); assert!( - popup.contains("Auto-review (current)"), - "expected Auto-review to be current even with custom workspace-write details: {popup}" + popup.contains("Approve for me (current)"), + "expected Approve for me to be current even with custom workspace-write details: {popup}" ); } @@ -861,7 +861,7 @@ async fn permissions_selection_can_disable_auto_review() { event, AppEvent::UpdateApprovalsReviewer(ApprovalsReviewer::User) )), - "expected selecting Default from Auto-review to switch back to manual approval review: {events:?}" + "expected selecting Ask for approval from Approve for me to switch back to manual approval review: {events:?}" ); assert!( !events @@ -906,8 +906,8 @@ async fn permissions_selection_sends_approvals_reviewer_in_override_turn_context assert!( popup .lines() - .any(|line| line.contains("Auto-review") && line.contains('›')), - "expected one Down from Default to select Auto-review: {popup}" + .any(|line| line.contains("Approve for me") && line.contains('›')), + "expected one Down from Ask for approval to select Approve for me: {popup}" ); chat.handle_key_event(KeyEvent::from(KeyCode::Enter)); diff --git a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs index ff1feb084c..4f425c25c9 100644 --- a/codex-rs/tui/src/chatwidget/tests/slash_commands.rs +++ b/codex-rs/tui/src/chatwidget/tests/slash_commands.rs @@ -1727,6 +1727,27 @@ async fn slash_clear_is_disabled_while_task_running() { assert!(rx.try_recv().is_err(), "expected no follow-up events"); } +#[tokio::test] +async fn slash_archive_is_disabled_while_task_running() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.bottom_pane.set_task_running(/*running*/ true); + + chat.dispatch_command(SlashCommand::Archive); + + let event = rx.try_recv().expect("expected disabled command error"); + match event { + AppEvent::InsertHistoryCell(cell) => { + let rendered = lines_to_single_string(&cell.display_lines(/*width*/ 80)); + assert!( + rendered.contains("'/archive' is disabled while a task is in progress."), + "expected /archive task-running error, got {rendered:?}" + ); + } + other => panic!("expected InsertHistoryCell error, got {other:?}"), + } + assert!(rx.try_recv().is_err(), "expected no follow-up events"); +} + #[tokio::test] async fn slash_memory_drop_reports_stubbed_feature() { let (mut chat, mut rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; @@ -1846,6 +1867,24 @@ async fn slash_resume_opens_picker() { assert_matches!(rx.try_recv(), Ok(AppEvent::OpenResumePicker)); } +#[tokio::test] +async fn slash_archive_confirmation_requests_current_thread_archive() { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + + chat.dispatch_command(SlashCommand::Archive); + + assert!(chat.bottom_pane.has_active_view()); + assert_matches!(rx.try_recv(), Err(TryRecvError::Empty)); + + let popup = render_bottom_popup(&chat, /*width*/ 80); + assert_chatwidget_snapshot!("slash_archive_confirmation_popup", popup); + + chat.handle_key_event(KeyEvent::from(KeyCode::Down)); + chat.handle_key_event(KeyEvent::from(KeyCode::Enter)); + + assert_matches!(rx.try_recv(), Ok(AppEvent::ArchiveCurrentThread)); +} + #[tokio::test] async fn slash_resume_with_arg_requests_named_session() { let (mut chat, mut rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; diff --git a/codex-rs/tui/src/slash_command.rs b/codex-rs/tui/src/slash_command.rs index 4a43409296..3fccdf5967 100644 --- a/codex-rs/tui/src/slash_command.rs +++ b/codex-rs/tui/src/slash_command.rs @@ -30,6 +30,7 @@ pub enum SlashCommand { Review, Rename, New, + Archive, Resume, Fork, Init, @@ -86,6 +87,7 @@ impl SlashCommand { SlashCommand::Review => "review my current changes and find issues", SlashCommand::Rename => "rename the current thread", SlashCommand::Resume => "resume a saved chat", + SlashCommand::Archive => "archive this session and exit", SlashCommand::Clear => "clear the terminal and start a new chat", SlashCommand::Fork => "fork the current chat", SlashCommand::Quit | SlashCommand::Exit => "exit Codex", @@ -180,6 +182,7 @@ impl SlashCommand { pub fn available_during_task(self) -> bool { match self { SlashCommand::New + | SlashCommand::Archive | SlashCommand::Resume | SlashCommand::Fork | SlashCommand::Init diff --git a/codex-rs/tui/src/status/card.rs b/codex-rs/tui/src/status/card.rs index 8da9166a3b..0117508ae5 100644 --- a/codex-rs/tui/src/status/card.rs +++ b/codex-rs/tui/src/status/card.rs @@ -682,13 +682,14 @@ fn status_approval_label( approvals_reviewer: ApprovalsReviewer, approval: &str, ) -> String { - if approval_policy == AskForApproval::OnRequest - && approvals_reviewer == ApprovalsReviewer::AutoReview - { - "auto-review".to_string() - } else { - approval.to_string() + if approval_policy == AskForApproval::OnRequest { + return match approvals_reviewer { + ApprovalsReviewer::AutoReview => "Approve for me".to_string(), + ApprovalsReviewer::User => "Ask for approval".to_string(), + }; } + + approval.to_string() } impl HistoryCell for StatusHistoryCell { diff --git a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_cached_limits_hide_credits_without_flag.snap b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_cached_limits_hide_credits_without_flag.snap index 6293c0d742..c7d89b17aa 100644 --- a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_cached_limits_hide_credits_without_flag.snap +++ b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_cached_limits_hide_credits_without_flag.snap @@ -4,20 +4,20 @@ expression: sanitized --- /status -╭───────────────────────────────────────────────────────────────────────╮ -│ >_ OpenAI Codex (v0.0.0) │ -│ │ -│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ -│ information on rate limits and credits │ -│ │ -│ Model: gpt-5.1-codex (reasoning none, summaries auto) │ -│ Directory: [[workspace]] │ -│ Permissions: Custom (workspace with network access, on-request) │ -│ Agents.md: │ -│ │ -│ Token usage: 1.05K total (700 input + 350 output) │ -│ Context window: 100% left (1.45K used / 272K) │ -│ 5h limit: [████████░░░░░░░░░░░░] 40% left (resets 11:32) │ -│ Weekly limit: [█████████████░░░░░░░] 65% left (resets 11:52) │ -│ Warning: limits may be stale - start new turn to refresh. │ -╰───────────────────────────────────────────────────────────────────────╯ +╭─────────────────────────────────────────────────────────────────────────────╮ +│ >_ OpenAI Codex (v0.0.0) │ +│ │ +│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ +│ information on rate limits and credits │ +│ │ +│ Model: gpt-5.1-codex (reasoning none, summaries auto) │ +│ Directory: [[workspace]] │ +│ Permissions: Custom (workspace with network access, Ask for approval) │ +│ Agents.md: │ +│ │ +│ Token usage: 1.05K total (700 input + 350 output) │ +│ Context window: 100% left (1.45K used / 272K) │ +│ 5h limit: [████████░░░░░░░░░░░░] 40% left (resets 11:32) │ +│ Weekly limit: [█████████████░░░░░░░] 65% left (resets 11:52) │ +│ Warning: limits may be stale - start new turn to refresh. │ +╰─────────────────────────────────────────────────────────────────────────────╯ diff --git a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_includes_credits_and_limits.snap b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_includes_credits_and_limits.snap index 0c8dae72cf..48fa575d1f 100644 --- a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_includes_credits_and_limits.snap +++ b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_includes_credits_and_limits.snap @@ -4,20 +4,20 @@ expression: sanitized --- /status -╭───────────────────────────────────────────────────────────────────────╮ -│ >_ OpenAI Codex (v0.0.0) │ -│ │ -│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ -│ information on rate limits and credits │ -│ │ -│ Model: gpt-5.1-codex (reasoning none, summaries auto) │ -│ Directory: [[workspace]] │ -│ Permissions: Custom (workspace with network access, on-request) │ -│ Agents.md: │ -│ │ -│ Token usage: 2K total (1.4K input + 600 output) │ -│ Context window: 100% left (2.2K used / 272K) │ -│ 5h limit: [███████████░░░░░░░░░] 55% left (resets 09:25) │ -│ Weekly limit: [██████████████░░░░░░] 70% left (resets 09:55) │ -│ Credits: 38 credits │ -╰───────────────────────────────────────────────────────────────────────╯ +╭─────────────────────────────────────────────────────────────────────────────╮ +│ >_ OpenAI Codex (v0.0.0) │ +│ │ +│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ +│ information on rate limits and credits │ +│ │ +│ Model: gpt-5.1-codex (reasoning none, summaries auto) │ +│ Directory: [[workspace]] │ +│ Permissions: Custom (workspace with network access, Ask for approval) │ +│ Agents.md: │ +│ │ +│ Token usage: 2K total (1.4K input + 600 output) │ +│ Context window: 100% left (2.2K used / 272K) │ +│ 5h limit: [███████████░░░░░░░░░] 55% left (resets 09:25) │ +│ Weekly limit: [██████████████░░░░░░] 70% left (resets 09:55) │ +│ Credits: 38 credits │ +╰─────────────────────────────────────────────────────────────────────────────╯ diff --git a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_includes_forked_from.snap b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_includes_forked_from.snap index f884d006ab..152575179e 100644 --- a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_includes_forked_from.snap +++ b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_includes_forked_from.snap @@ -4,20 +4,20 @@ expression: sanitized --- /status -╭───────────────────────────────────────────────────────────────────────╮ -│ >_ OpenAI Codex (v0.0.0) │ -│ │ -│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ -│ information on rate limits and credits │ -│ │ -│ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ -│ Directory: [[workspace]] │ -│ Permissions: Custom (workspace with network access, on-request) │ -│ Agents.md: │ -│ Session: 0f0f3c13-6cf9-4aa4-8b80-7d49c2f1be2e │ -│ Forked from: e9f18a88-8081-4e51-9d4e-8af5cde2d8dd │ -│ │ -│ Token usage: 1.2K total (800 input + 400 output) │ -│ Context window: 100% left (1.2K used / 272K) │ -│ Limits: data not available yet │ -╰───────────────────────────────────────────────────────────────────────╯ +╭─────────────────────────────────────────────────────────────────────────────╮ +│ >_ OpenAI Codex (v0.0.0) │ +│ │ +│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ +│ information on rate limits and credits │ +│ │ +│ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ +│ Directory: [[workspace]] │ +│ Permissions: Custom (workspace with network access, Ask for approval) │ +│ Agents.md: │ +│ Session: 0f0f3c13-6cf9-4aa4-8b80-7d49c2f1be2e │ +│ Forked from: e9f18a88-8081-4e51-9d4e-8af5cde2d8dd │ +│ │ +│ Token usage: 1.2K total (800 input + 400 output) │ +│ Context window: 100% left (1.2K used / 272K) │ +│ Limits: data not available yet │ +╰─────────────────────────────────────────────────────────────────────────────╯ diff --git a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_includes_monthly_limit.snap b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_includes_monthly_limit.snap index fc691f6bf7..7a4f58a823 100644 --- a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_includes_monthly_limit.snap +++ b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_includes_monthly_limit.snap @@ -4,18 +4,18 @@ expression: sanitized --- /status -╭────────────────────────────────────────────────────────────────────────────╮ -│ >_ OpenAI Codex (v0.0.0) │ -│ │ -│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ -│ information on rate limits and credits │ -│ │ -│ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ -│ Directory: [[workspace]] │ -│ Permissions: Custom (workspace with network access, on-request) │ -│ Agents.md: │ -│ │ -│ Token usage: 1.2K total (800 input + 400 output) │ -│ Context window: 100% left (1.2K used / 272K) │ -│ Monthly limit: [██████████████████░░] 88% left (resets 07:08 on 7 May) │ -╰────────────────────────────────────────────────────────────────────────────╯ +╭─────────────────────────────────────────────────────────────────────────────╮ +│ >_ OpenAI Codex (v0.0.0) │ +│ │ +│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ +│ information on rate limits and credits │ +│ │ +│ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ +│ Directory: [[workspace]] │ +│ Permissions: Custom (workspace with network access, Ask for approval) │ +│ Agents.md: │ +│ │ +│ Token usage: 1.2K total (800 input + 400 output) │ +│ Context window: 100% left (1.2K used / 272K) │ +│ Monthly limit: [██████████████████░░] 88% left (resets 07:08 on 7 May) │ +╰─────────────────────────────────────────────────────────────────────────────╯ diff --git a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_includes_reasoning_details.snap b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_includes_reasoning_details.snap index 8719230212..96b2ed1648 100644 --- a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_includes_reasoning_details.snap +++ b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_includes_reasoning_details.snap @@ -12,7 +12,7 @@ expression: sanitized │ │ │ Model: gpt-5.1-codex-max (reasoning high, summaries detailed) │ │ Directory: [[workspace]] │ -│ Permissions: Workspace (on-request) │ +│ Permissions: Workspace (Ask for approval) │ │ Agents.md: │ │ │ │ Token usage: 1.9K total (1K input + 900 output) │ diff --git a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_active_user_defined_profile.snap b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_active_user_defined_profile.snap index a57c34bdf1..4604b26f7e 100644 --- a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_active_user_defined_profile.snap +++ b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_active_user_defined_profile.snap @@ -12,7 +12,7 @@ expression: sanitized │ │ │ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ │ Directory: [[workspace]] │ -│ Permissions: Profile locked (read-only, on-request) │ +│ Permissions: Profile locked (read-only, Ask for approval) │ │ Agents.md: │ │ │ │ Token usage: 0 total (0 input + 0 output) │ diff --git a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_auto_review_permissions.snap b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_auto_review_permissions.snap index f1ceb21ab0..c739d55075 100644 --- a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_auto_review_permissions.snap +++ b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_auto_review_permissions.snap @@ -12,7 +12,7 @@ expression: sanitized │ │ │ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ │ Directory: [[workspace]] │ -│ Permissions: Workspace (auto-review) │ +│ Permissions: Workspace (Approve for me) │ │ Agents.md: │ │ │ │ Token usage: 0 total (0 input + 0 output) │ diff --git a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_missing_limits_message.snap b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_missing_limits_message.snap index 79344df582..533be27976 100644 --- a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_missing_limits_message.snap +++ b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_missing_limits_message.snap @@ -4,18 +4,18 @@ expression: sanitized --- /status -╭───────────────────────────────────────────────────────────────────────╮ -│ >_ OpenAI Codex (v0.0.0) │ -│ │ -│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ -│ information on rate limits and credits │ -│ │ -│ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ -│ Directory: [[workspace]] │ -│ Permissions: Custom (workspace with network access, on-request) │ -│ Agents.md: │ -│ │ -│ Token usage: 750 total (500 input + 250 output) │ -│ Context window: 100% left (750 used / 272K) │ -│ Limits: data not available yet │ -╰───────────────────────────────────────────────────────────────────────╯ +╭─────────────────────────────────────────────────────────────────────────────╮ +│ >_ OpenAI Codex (v0.0.0) │ +│ │ +│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ +│ information on rate limits and credits │ +│ │ +│ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ +│ Directory: [[workspace]] │ +│ Permissions: Custom (workspace with network access, Ask for approval) │ +│ Agents.md: │ +│ │ +│ Token usage: 750 total (500 input + 250 output) │ +│ Context window: 100% left (750 used / 272K) │ +│ Limits: data not available yet │ +╰─────────────────────────────────────────────────────────────────────────────╯ diff --git a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_refreshing_limits_notice.snap b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_refreshing_limits_notice.snap index f633cd7f6c..631982bcf9 100644 --- a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_refreshing_limits_notice.snap +++ b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_refreshing_limits_notice.snap @@ -4,19 +4,19 @@ expression: sanitized --- /status -╭───────────────────────────────────────────────────────────────────────╮ -│ >_ OpenAI Codex (v0.0.0) │ -│ │ -│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ -│ information on rate limits and credits │ -│ │ -│ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ -│ Directory: [[workspace]] │ -│ Permissions: Custom (workspace with network access, on-request) │ -│ Agents.md: │ -│ │ -│ Token usage: 750 total (500 input + 250 output) │ -│ Context window: 100% left (750 used / 272K) │ -│ 5h limit: [███████████░░░░░░░░░] 55% left (resets 08:24) │ -│ Weekly limit: [██████████████░░░░░░] 70% left (resets 08:54) │ -╰───────────────────────────────────────────────────────────────────────╯ +╭─────────────────────────────────────────────────────────────────────────────╮ +│ >_ OpenAI Codex (v0.0.0) │ +│ │ +│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ +│ information on rate limits and credits │ +│ │ +│ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ +│ Directory: [[workspace]] │ +│ Permissions: Custom (workspace with network access, Ask for approval) │ +│ Agents.md: │ +│ │ +│ Token usage: 750 total (500 input + 250 output) │ +│ Context window: 100% left (750 used / 272K) │ +│ 5h limit: [███████████░░░░░░░░░] 55% left (resets 08:24) │ +│ Weekly limit: [██████████████░░░░░░] 70% left (resets 08:54) │ +╰─────────────────────────────────────────────────────────────────────────────╯ diff --git a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_stale_limits_message.snap b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_stale_limits_message.snap index dc52e22ec4..8facd56275 100644 --- a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_stale_limits_message.snap +++ b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_stale_limits_message.snap @@ -4,20 +4,20 @@ expression: sanitized --- /status -╭───────────────────────────────────────────────────────────────────────╮ -│ >_ OpenAI Codex (v0.0.0) │ -│ │ -│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ -│ information on rate limits and credits │ -│ │ -│ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ -│ Directory: [[workspace]] │ -│ Permissions: Custom (workspace with network access, on-request) │ -│ Agents.md: │ -│ │ -│ Token usage: 1.9K total (1K input + 900 output) │ -│ Context window: 100% left (2.25K used / 272K) │ -│ 5h limit: [██████░░░░░░░░░░░░░░] 28% left (resets 03:14) │ -│ Weekly limit: [████████████░░░░░░░░] 60% left (resets 03:34) │ -│ Warning: limits may be stale - start new turn to refresh. │ -╰───────────────────────────────────────────────────────────────────────╯ +╭─────────────────────────────────────────────────────────────────────────────╮ +│ >_ OpenAI Codex (v0.0.0) │ +│ │ +│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ +│ information on rate limits and credits │ +│ │ +│ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ +│ Directory: [[workspace]] │ +│ Permissions: Custom (workspace with network access, Ask for approval) │ +│ Agents.md: │ +│ │ +│ Token usage: 1.9K total (1K input + 900 output) │ +│ Context window: 100% left (2.25K used / 272K) │ +│ 5h limit: [██████░░░░░░░░░░░░░░] 28% left (resets 03:14) │ +│ Weekly limit: [████████████░░░░░░░░] 60% left (resets 03:34) │ +│ Warning: limits may be stale - start new turn to refresh. │ +╰─────────────────────────────────────────────────────────────────────────────╯ diff --git a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_unavailable_limits_message.snap b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_unavailable_limits_message.snap index 83d30b5adc..d46a483e4b 100644 --- a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_unavailable_limits_message.snap +++ b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_shows_unavailable_limits_message.snap @@ -4,18 +4,18 @@ expression: sanitized --- /status -╭───────────────────────────────────────────────────────────────────────╮ -│ >_ OpenAI Codex (v0.0.0) │ -│ │ -│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ -│ information on rate limits and credits │ -│ │ -│ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ -│ Directory: [[workspace]] │ -│ Permissions: Custom (workspace with network access, on-request) │ -│ Agents.md: │ -│ │ -│ Token usage: 750 total (500 input + 250 output) │ -│ Context window: 100% left (750 used / 272K) │ -│ Limits: not available for this account │ -╰───────────────────────────────────────────────────────────────────────╯ +╭─────────────────────────────────────────────────────────────────────────────╮ +│ >_ OpenAI Codex (v0.0.0) │ +│ │ +│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ +│ information on rate limits and credits │ +│ │ +│ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ +│ Directory: [[workspace]] │ +│ Permissions: Custom (workspace with network access, Ask for approval) │ +│ Agents.md: │ +│ │ +│ Token usage: 750 total (500 input + 250 output) │ +│ Context window: 100% left (750 used / 272K) │ +│ Limits: not available for this account │ +╰─────────────────────────────────────────────────────────────────────────────╯ diff --git a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_treats_refreshing_empty_limits_as_unavailable.snap b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_treats_refreshing_empty_limits_as_unavailable.snap index 83d30b5adc..d46a483e4b 100644 --- a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_treats_refreshing_empty_limits_as_unavailable.snap +++ b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_treats_refreshing_empty_limits_as_unavailable.snap @@ -4,18 +4,18 @@ expression: sanitized --- /status -╭───────────────────────────────────────────────────────────────────────╮ -│ >_ OpenAI Codex (v0.0.0) │ -│ │ -│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ -│ information on rate limits and credits │ -│ │ -│ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ -│ Directory: [[workspace]] │ -│ Permissions: Custom (workspace with network access, on-request) │ -│ Agents.md: │ -│ │ -│ Token usage: 750 total (500 input + 250 output) │ -│ Context window: 100% left (750 used / 272K) │ -│ Limits: not available for this account │ -╰───────────────────────────────────────────────────────────────────────╯ +╭─────────────────────────────────────────────────────────────────────────────╮ +│ >_ OpenAI Codex (v0.0.0) │ +│ │ +│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ +│ information on rate limits and credits │ +│ │ +│ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ +│ Directory: [[workspace]] │ +│ Permissions: Custom (workspace with network access, Ask for approval) │ +│ Agents.md: │ +│ │ +│ Token usage: 750 total (500 input + 250 output) │ +│ Context window: 100% left (750 used / 272K) │ +│ Limits: not available for this account │ +╰─────────────────────────────────────────────────────────────────────────────╯ diff --git a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_truncates_in_narrow_terminal.snap b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_truncates_in_narrow_terminal.snap index a048726bb3..6b28b53b62 100644 --- a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_truncates_in_narrow_terminal.snap +++ b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_truncates_in_narrow_terminal.snap @@ -12,7 +12,7 @@ expression: sanitized │ │ │ Model: gpt-5.1-codex-max (reasoning high, summaries de │ │ Directory: [[workspace]] │ -│ Permissions: Custom (workspace with network access, on-reque │ +│ Permissions: Custom (workspace with network access, Ask for │ │ Agents.md: │ │ │ │ Token usage: 1.9K total (1K input + 900 output) │ diff --git a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_uses_default_reasoning_when_config_empty.snap b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_uses_default_reasoning_when_config_empty.snap index 2f6925c830..9683beed97 100644 --- a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_uses_default_reasoning_when_config_empty.snap +++ b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_uses_default_reasoning_when_config_empty.snap @@ -4,21 +4,21 @@ expression: sanitized --- /status -╭──────────────────────────────────────────────────────────────────────────╮ -│ >_ OpenAI Codex (v0.0.0) │ -│ │ -│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ -│ information on rate limits and credits │ -│ │ -│ Remote: unix:///tmp/codex-home/app-server-control/app-server- │ -│ control.sock (v0.133.0) │ -│ │ -│ Model: gpt-5.1-codex-max (reasoning medium, summaries auto) │ -│ Directory: [[workspace]] │ -│ Permissions: Custom (workspace with network access, on-request) │ -│ Agents.md: │ -│ │ -│ Token usage: 750 total (500 input + 250 output) │ -│ Context window: 100% left (750 used / 272K) │ -│ Limits: data not available yet │ -╰──────────────────────────────────────────────────────────────────────────╯ +╭─────────────────────────────────────────────────────────────────────────────╮ +│ >_ OpenAI Codex (v0.0.0) │ +│ │ +│ Visit https://chatgpt.com/codex/settings/usage for up-to-date │ +│ information on rate limits and credits │ +│ │ +│ Remote: unix:///tmp/codex-home/app-server-control/app-server- │ +│ control.sock (v0.133.0) │ +│ │ +│ Model: gpt-5.1-codex-max (reasoning medium, summaries auto) │ +│ Directory: [[workspace]] │ +│ Permissions: Custom (workspace with network access, Ask for approval) │ +│ Agents.md: │ +│ │ +│ Token usage: 750 total (500 input + 250 output) │ +│ Context window: 100% left (750 used / 272K) │ +│ Limits: data not available yet │ +╰─────────────────────────────────────────────────────────────────────────────╯ diff --git a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_uses_generic_limit_labels_for_unsupported_windows.snap b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_uses_generic_limit_labels_for_unsupported_windows.snap index f983cf478d..6b6fad683b 100644 --- a/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_uses_generic_limit_labels_for_unsupported_windows.snap +++ b/codex-rs/tui/src/status/snapshots/codex_tui__status__tests__status_snapshot_uses_generic_limit_labels_for_unsupported_windows.snap @@ -12,7 +12,7 @@ expression: sanitized │ │ │ Model: gpt-5.1-codex-max (reasoning none, summaries auto) │ │ Directory: [[workspace]] │ -│ Permissions: Custom (workspace with network access, on-request) │ +│ Permissions: Custom (workspace with network access, Ask for app │ │ Agents.md: │ │ │ │ Token usage: 1.2K total (800 input + 400 output) │ diff --git a/codex-rs/tui/src/status/tests.rs b/codex-rs/tui/src/status/tests.rs index 9609703b0d..8b4e41ce76 100644 --- a/codex-rs/tui/src/status/tests.rs +++ b/codex-rs/tui/src/status/tests.rs @@ -291,7 +291,7 @@ async fn status_permissions_non_default_workspace_write_uses_workspace_label() { assert_eq!( permissions_text_for(&config).as_deref(), - Some("Custom (workspace with network access, on-request)") + Some("Custom (workspace with network access, Ask for approval)") ); } @@ -314,7 +314,7 @@ async fn status_permissions_named_read_only_profile_shows_builtin_label() { assert_eq!( permissions_text_for(&config).as_deref(), - Some("Read Only (on-request)") + Some("Read Only (Ask for approval)") ); } @@ -344,7 +344,7 @@ async fn status_permissions_read_only_profile_shows_additional_writable_roots() assert_eq!( permissions_text_for(&config).as_deref(), - Some("Read Only (on-request)") + Some("Read Only (Ask for approval)") ); } @@ -367,7 +367,7 @@ async fn status_permissions_named_workspace_profile_shows_builtin_label() { assert_eq!( permissions_text_for(&config).as_deref(), - Some("Workspace (on-request)") + Some("Workspace (Ask for approval)") ); } @@ -391,7 +391,7 @@ async fn status_permissions_workspace_auto_review_shows_reviewer_label() { assert_eq!( permissions_text_for(&config).as_deref(), - Some("Workspace (auto-review)") + Some("Workspace (Approve for me)") ); } @@ -420,7 +420,7 @@ async fn status_permissions_named_profile_shows_additional_writable_roots() { assert_eq!( permissions_text_for(&config).as_deref(), - Some("Workspace (on-request)") + Some("Workspace (Ask for approval)") ); } @@ -449,7 +449,10 @@ async fn status_permissions_workspace_roots_show_additional_directories() { assert_eq!( permissions_text_for(&config), - Some(format!("Workspace [{}] (on-request)", extra_root.display())) + Some(format!( + "Workspace [{}] (Ask for approval)", + extra_root.display() + )) ); } @@ -483,7 +486,7 @@ async fn status_permissions_workspace_roots_include_profile_defined_directories( assert_eq!( permissions_text_for(&config), Some(format!( - "Workspace [{}] (on-request)", + "Workspace [{}] (Ask for approval)", profile_root.display() )) ); @@ -513,7 +516,7 @@ async fn status_permissions_broadened_workspace_profile_shows_builtin_label() { assert_eq!( permissions_text_for(&config).as_deref(), - Some("Workspace with network access (on-request)") + Some("Workspace with network access (Ask for approval)") ); } @@ -531,7 +534,7 @@ async fn status_permissions_user_defined_profile_shows_name() { assert_eq!( permissions_text_for(&config).as_deref(), - Some("Profile locked (read-only, on-request)") + Some("Profile locked (read-only, Ask for approval)") ); } @@ -756,7 +759,7 @@ async fn status_permissions_full_disk_managed_with_network_is_danger_full_access assert_eq!( permissions_text_for(&config).as_deref(), - Some("Custom (danger-full-access, on-request)") + Some("Custom (danger-full-access, Ask for approval)") ); } @@ -779,7 +782,7 @@ async fn status_permissions_full_disk_managed_without_network_is_external_sandbo assert_eq!( permissions_text_for(&config).as_deref(), - Some("Custom (external-sandbox, on-request)") + Some("Custom (external-sandbox, Ask for approval)") ); } diff --git a/codex-rs/windows-sandbox-rs/src/bin/setup_main/win.rs b/codex-rs/windows-sandbox-rs/src/bin/setup_main/win.rs index 9b178fdca9..41fa2f720e 100644 --- a/codex-rs/windows-sandbox-rs/src/bin/setup_main/win.rs +++ b/codex-rs/windows-sandbox-rs/src/bin/setup_main/win.rs @@ -105,6 +105,7 @@ struct Payload { enum SetupMode { #[default] Full, + ProvisionOnly, ReadAclsOnly, } @@ -476,6 +477,7 @@ fn real_main() -> Result<()> { fn run_setup(payload: &Payload, log: &mut dyn Write, sbx_dir: &Path) -> Result<()> { match payload.mode { SetupMode::ReadAclsOnly => run_read_acl_only(payload, log), + SetupMode::ProvisionOnly => run_provision_only(payload, log, sbx_dir), SetupMode::Full => run_setup_full(payload, log, sbx_dir), } } @@ -543,31 +545,182 @@ fn run_read_acl_only(payload: &Payload, log: &mut dyn Write) -> Result<()> { Ok(()) } +fn provision_and_hide_sandbox_users( + payload: &Payload, + log: &mut dyn Write, + sbx_dir: &Path, +) -> Result<()> { + let provision_result = provision_sandbox_users( + &payload.codex_home, + &payload.offline_username, + &payload.online_username, + &payload.proxy_ports, + payload.allow_local_binding, + log, + ); + if let Err(err) = provision_result { + if extract_setup_failure(&err).is_some() { + return Err(err); + } + return Err(anyhow::Error::new(SetupFailure::new( + SetupErrorCode::HelperUserProvisionFailed, + format!("provision sandbox users failed: {err}"), + ))); + } + let users = vec![ + payload.offline_username.clone(), + payload.online_username.clone(), + ]; + hide_newly_created_users(&users, sbx_dir); + Ok(()) +} + +fn configure_offline_sandbox_network( + payload: &Payload, + offline_sid_str: &str, + log: &mut dyn Write, +) -> Result<()> { + let proxy_allowlist_result = firewall::ensure_offline_proxy_allowlist( + offline_sid_str, + &payload.proxy_ports, + payload.allow_local_binding, + log, + ); + if let Err(err) = proxy_allowlist_result { + if extract_setup_failure(&err).is_some() { + return Err(err); + } + return Err(anyhow::Error::new(SetupFailure::new( + SetupErrorCode::HelperFirewallRuleCreateOrAddFailed, + format!("ensure offline proxy allowlist failed: {err}"), + ))); + } + let firewall_result = firewall::ensure_offline_outbound_block(offline_sid_str, log); + if let Err(err) = firewall_result { + if extract_setup_failure(&err).is_some() { + return Err(err); + } + return Err(anyhow::Error::new(SetupFailure::new( + SetupErrorCode::HelperFirewallRuleCreateOrAddFailed, + format!("ensure offline outbound block failed: {err}"), + ))); + } + install_wfp_filters( + &payload.codex_home, + &payload.offline_username, + payload.otel.as_ref(), + |message| { + let _ = log_line(log, message); + }, + ); + Ok(()) +} + +fn lock_persistent_sandbox_dirs( + payload: &Payload, + sandbox_group_sid: &[u8], + log: &mut dyn Write, +) -> Result<()> { + lock_sandbox_dir( + &sandbox_dir(&payload.codex_home), + &payload.real_user, + sandbox_group_sid, + GRANT_ACCESS, + FILE_GENERIC_READ | FILE_GENERIC_WRITE | FILE_GENERIC_EXECUTE | DELETE, + FILE_GENERIC_READ | FILE_GENERIC_WRITE | FILE_GENERIC_EXECUTE, + log, + ) + .map_err(|err| { + anyhow::Error::new(SetupFailure::new( + SetupErrorCode::HelperSandboxLockFailed, + format!( + "lock sandbox dir {} failed: {err}", + sandbox_dir(&payload.codex_home).display() + ), + )) + })?; + lock_sandbox_dir( + &sandbox_secrets_dir(&payload.codex_home), + &payload.real_user, + sandbox_group_sid, + DENY_ACCESS, + FILE_GENERIC_READ | FILE_GENERIC_WRITE | FILE_GENERIC_EXECUTE | DELETE, + FILE_GENERIC_READ | FILE_GENERIC_WRITE | FILE_GENERIC_EXECUTE, + log, + ) + .map_err(|err| { + anyhow::Error::new(SetupFailure::new( + SetupErrorCode::HelperSandboxLockFailed, + format!( + "lock sandbox secrets dir {} failed: {err}", + sandbox_secrets_dir(&payload.codex_home).display() + ), + )) + })?; + let legacy_users = sandbox_dir(&payload.codex_home).join("sandbox_users.json"); + if legacy_users.exists() { + let _ = std::fs::remove_file(&legacy_users); + } + Ok(()) +} + +fn lock_sandbox_bin_dir( + payload: &Payload, + sandbox_group_sid: &[u8], + log: &mut dyn Write, +) -> Result<()> { + lock_sandbox_dir( + &sandbox_bin_dir(&payload.codex_home), + &payload.real_user, + sandbox_group_sid, + GRANT_ACCESS, + FILE_GENERIC_READ | FILE_GENERIC_EXECUTE, + FILE_GENERIC_READ | FILE_GENERIC_WRITE | FILE_GENERIC_EXECUTE | DELETE, + log, + ) + .map_err(|err| { + anyhow::Error::new(SetupFailure::new( + SetupErrorCode::HelperSandboxLockFailed, + format!( + "lock sandbox bin dir {} failed: {err}", + sandbox_bin_dir(&payload.codex_home).display() + ), + )) + }) +} + +fn run_provision_only(payload: &Payload, log: &mut dyn Write, sbx_dir: &Path) -> Result<()> { + provision_and_hide_sandbox_users(payload, log, sbx_dir)?; + let offline_sid = resolve_sid(&payload.offline_username).map_err(|err| { + anyhow::Error::new(SetupFailure::new( + SetupErrorCode::HelperSidResolveFailed, + format!( + "resolve SID for offline user {} failed: {err}", + payload.offline_username + ), + )) + })?; + let offline_sid_str = string_from_sid_bytes(&offline_sid).map_err(anyhow::Error::msg)?; + + let sandbox_group_sid = resolve_sandbox_users_group_sid().map_err(|err| { + anyhow::Error::new(SetupFailure::new( + SetupErrorCode::HelperSidResolveFailed, + format!("resolve sandbox users group SID failed: {err}"), + )) + })?; + + configure_offline_sandbox_network(payload, &offline_sid_str, log)?; + + lock_sandbox_bin_dir(payload, &sandbox_group_sid, log)?; + lock_persistent_sandbox_dirs(payload, &sandbox_group_sid, log)?; + log_note("setup provisioning binary completed", Some(sbx_dir)); + Ok(()) +} + fn run_setup_full(payload: &Payload, log: &mut dyn Write, sbx_dir: &Path) -> Result<()> { let refresh_only = payload.refresh_only; if !refresh_only { - let provision_result = provision_sandbox_users( - &payload.codex_home, - &payload.offline_username, - &payload.online_username, - &payload.proxy_ports, - payload.allow_local_binding, - log, - ); - if let Err(err) = provision_result { - if extract_setup_failure(&err).is_some() { - return Err(err); - } - return Err(anyhow::Error::new(SetupFailure::new( - SetupErrorCode::HelperUserProvisionFailed, - format!("provision sandbox users failed: {err}"), - ))); - } - let users = vec![ - payload.offline_username.clone(), - payload.online_username.clone(), - ]; - hide_newly_created_users(&users, sbx_dir); + provision_and_hide_sandbox_users(payload, log, sbx_dir)?; } let offline_sid = resolve_sid(&payload.offline_username).map_err(|err| { anyhow::Error::new(SetupFailure::new( @@ -597,39 +750,7 @@ fn run_setup_full(payload: &Payload, log: &mut dyn Write, sbx_dir: &Path) -> Res let mut refresh_errors: Vec = Vec::new(); if !refresh_only { - let proxy_allowlist_result = firewall::ensure_offline_proxy_allowlist( - &offline_sid_str, - &payload.proxy_ports, - payload.allow_local_binding, - log, - ); - if let Err(err) = proxy_allowlist_result { - if extract_setup_failure(&err).is_some() { - return Err(err); - } - return Err(anyhow::Error::new(SetupFailure::new( - SetupErrorCode::HelperFirewallRuleCreateOrAddFailed, - format!("ensure offline proxy allowlist failed: {err}"), - ))); - } - let firewall_result = firewall::ensure_offline_outbound_block(&offline_sid_str, log); - if let Err(err) = firewall_result { - if extract_setup_failure(&err).is_some() { - return Err(err); - } - return Err(anyhow::Error::new(SetupFailure::new( - SetupErrorCode::HelperFirewallRuleCreateOrAddFailed, - format!("ensure offline outbound block failed: {err}"), - ))); - } - install_wfp_filters( - &payload.codex_home, - &payload.offline_username, - payload.otel.as_ref(), - |message| { - let _ = log_line(log, message); - }, - ); + configure_offline_sandbox_network(payload, &offline_sid_str, log)?; } // Deny-read ACEs must be present before the sandboxed command starts. Apply @@ -865,24 +986,7 @@ fn run_setup_full(payload: &Payload, log: &mut dyn Write, sbx_dir: &Path) -> Res } } - lock_sandbox_dir( - &sandbox_bin_dir(&payload.codex_home), - &payload.real_user, - &sandbox_group_sid, - GRANT_ACCESS, - FILE_GENERIC_READ | FILE_GENERIC_EXECUTE, - FILE_GENERIC_READ | FILE_GENERIC_WRITE | FILE_GENERIC_EXECUTE | DELETE, - log, - ) - .map_err(|err| { - anyhow::Error::new(SetupFailure::new( - SetupErrorCode::HelperSandboxLockFailed, - format!( - "lock sandbox bin dir {} failed: {err}", - sandbox_bin_dir(&payload.codex_home).display() - ), - )) - })?; + lock_sandbox_bin_dir(payload, &sandbox_group_sid, log)?; if refresh_only { log_line( @@ -895,46 +999,7 @@ fn run_setup_full(payload: &Payload, log: &mut dyn Write, sbx_dir: &Path) -> Res )?; } if !refresh_only { - lock_sandbox_dir( - &sandbox_dir(&payload.codex_home), - &payload.real_user, - &sandbox_group_sid, - GRANT_ACCESS, - FILE_GENERIC_READ | FILE_GENERIC_WRITE | FILE_GENERIC_EXECUTE | DELETE, - FILE_GENERIC_READ | FILE_GENERIC_WRITE | FILE_GENERIC_EXECUTE, - log, - ) - .map_err(|err| { - anyhow::Error::new(SetupFailure::new( - SetupErrorCode::HelperSandboxLockFailed, - format!( - "lock sandbox dir {} failed: {err}", - sandbox_dir(&payload.codex_home).display() - ), - )) - })?; - lock_sandbox_dir( - &sandbox_secrets_dir(&payload.codex_home), - &payload.real_user, - &sandbox_group_sid, - DENY_ACCESS, - FILE_GENERIC_READ | FILE_GENERIC_WRITE | FILE_GENERIC_EXECUTE | DELETE, - FILE_GENERIC_READ | FILE_GENERIC_WRITE | FILE_GENERIC_EXECUTE, - log, - ) - .map_err(|err| { - anyhow::Error::new(SetupFailure::new( - SetupErrorCode::HelperSandboxLockFailed, - format!( - "lock sandbox secrets dir {} failed: {err}", - sandbox_secrets_dir(&payload.codex_home).display() - ), - )) - })?; - let legacy_users = sandbox_dir(&payload.codex_home).join("sandbox_users.json"); - if legacy_users.exists() { - let _ = std::fs::remove_file(&legacy_users); - } + lock_persistent_sandbox_dirs(payload, &sandbox_group_sid, log)?; } unsafe { @@ -986,6 +1051,15 @@ mod tests { assert_eq!(payload.otel, None); } + #[test] + fn payload_accepts_provision_only_mode() { + let mut payload = payload_json(); + payload["mode"] = json!("provision-only"); + let payload: Payload = serde_json::from_value(payload).expect("payload"); + + assert_eq!(payload.mode, super::SetupMode::ProvisionOnly); + } + #[test] fn payload_accepts_otel_settings() { let mut payload = payload_json(); diff --git a/codex-rs/windows-sandbox-rs/src/lib.rs b/codex-rs/windows-sandbox-rs/src/lib.rs index d05bb9554c..3bd2639868 100644 --- a/codex-rs/windows-sandbox-rs/src/lib.rs +++ b/codex-rs/windows-sandbox-rs/src/lib.rs @@ -241,6 +241,8 @@ pub use setup::SandboxSetupRequest; #[cfg(target_os = "windows")] pub use setup::SetupRootOverrides; #[cfg(target_os = "windows")] +pub use setup::run_elevated_provisioning_setup; +#[cfg(target_os = "windows")] pub use setup::run_elevated_setup; #[cfg(target_os = "windows")] pub use setup::run_setup_refresh; diff --git a/codex-rs/windows-sandbox-rs/src/setup.rs b/codex-rs/windows-sandbox-rs/src/setup.rs index bbb3881968..a0f46fe3f8 100644 --- a/codex-rs/windows-sandbox-rs/src/setup.rs +++ b/codex-rs/windows-sandbox-rs/src/setup.rs @@ -203,6 +203,7 @@ fn run_setup_refresh_inner( allow_local_binding: offline_proxy_settings.allow_local_binding, otel: None, real_user: std::env::var("USERNAME").unwrap_or_else(|_| "Administrators".to_string()), + mode: SetupMode::Full, refresh_only: true, }; let json = serde_json::to_vec(&payload)?; @@ -496,10 +497,18 @@ struct ElevationPayload { allow_local_binding: bool, otel: Option, real_user: String, + mode: SetupMode, #[serde(default)] refresh_only: bool, } +#[derive(Clone, Copy, Serialize)] +#[serde(rename_all = "kebab-case")] +enum SetupMode { + Full, + ProvisionOnly, +} + #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct OfflineProxySettings { pub proxy_ports: Vec, @@ -808,6 +817,7 @@ pub fn run_elevated_setup( allow_local_binding: offline_proxy_settings.allow_local_binding, real_user: std::env::var("USERNAME").unwrap_or_else(|_| "Administrators".to_string()), otel: codex_otel::global_statsig_metrics_settings(), + mode: SetupMode::Full, refresh_only: false, }; let needs_elevation = !is_elevated().map_err(|err| { @@ -819,6 +829,45 @@ pub fn run_elevated_setup( run_setup_exe(&payload, needs_elevation, request.codex_home) } +pub fn run_elevated_provisioning_setup(codex_home: &Path, real_user: &str) -> Result<()> { + let sbx_dir = sandbox_dir(codex_home); + std::fs::create_dir_all(&sbx_dir).map_err(|err| { + failure( + SetupErrorCode::OrchestratorSandboxDirCreateFailed, + format!("failed to create sandbox dir {}: {err}", sbx_dir.display()), + ) + })?; + if !is_elevated().map_err(|err| { + failure( + SetupErrorCode::OrchestratorElevationCheckFailed, + format!("failed to determine elevation state: {err}"), + ) + })? { + return Err(failure( + SetupErrorCode::OrchestratorElevationRequired, + "sandbox provisioning setup must be run from an elevated process", + )); + } + let payload = ElevationPayload { + version: SETUP_VERSION, + offline_username: OFFLINE_USERNAME.to_string(), + online_username: ONLINE_USERNAME.to_string(), + codex_home: codex_home.to_path_buf(), + command_cwd: codex_home.to_path_buf(), + read_roots: Vec::new(), + write_roots: Vec::new(), + deny_read_paths: Vec::new(), + deny_write_paths: Vec::new(), + proxy_ports: Vec::new(), + allow_local_binding: false, + otel: codex_otel::global_statsig_metrics_settings(), + real_user: real_user.to_string(), + mode: SetupMode::ProvisionOnly, + refresh_only: false, + }; + run_setup_exe(&payload, /*needs_elevation*/ false, codex_home) +} + fn build_payload_roots( request: &SandboxSetupRequest<'_>, overrides: &SetupRootOverrides, diff --git a/codex-rs/windows-sandbox-rs/src/setup_error.rs b/codex-rs/windows-sandbox-rs/src/setup_error.rs index efaaa531e2..4d2a9d57df 100644 --- a/codex-rs/windows-sandbox-rs/src/setup_error.rs +++ b/codex-rs/windows-sandbox-rs/src/setup_error.rs @@ -19,6 +19,8 @@ pub enum SetupErrorCode { OrchestratorSandboxDirCreateFailed, /// Failed to determine whether the current process is elevated. OrchestratorElevationCheckFailed, + /// The setup command requires an already elevated process. + OrchestratorElevationRequired, /// Failed to serialize the elevation payload before launching the helper. OrchestratorPayloadSerializeFailed, /// Failed to launch the setup helper process (spawn or ShellExecuteExW). @@ -75,6 +77,7 @@ impl SetupErrorCode { match self { Self::OrchestratorSandboxDirCreateFailed => "orchestrator_sandbox_dir_create_failed", Self::OrchestratorElevationCheckFailed => "orchestrator_elevation_check_failed", + Self::OrchestratorElevationRequired => "orchestrator_elevation_required", Self::OrchestratorPayloadSerializeFailed => "orchestrator_payload_serialize_failed", Self::OrchestratorHelperLaunchFailed => "orchestrator_helper_launch_failed", Self::OrchestratorHelperLaunchCanceled => "orchestrator_helper_launch_canceled",