From 49fb25997f3c09c25684c2a729cb933939a7f830 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Thu, 23 Apr 2026 18:57:13 -0700 Subject: [PATCH] Add sticky environment API and thread state (#18897) ## Summary - add sticky environment selections to app-server v2 thread/start and turn/start request flow - carry thread-level selections through core session/thread state - add app-server coverage for sticky selections and turn overrides ## Stack 1. This PR: API and thread persistence 2. #18898: config.toml named environment loading 3. #18899: downstream tool/runtime consumers ## Validation - Not run locally; split only. --------- Co-authored-by: Codex --- .../schema/json/v2/ThreadStartParams.json | 15 ++ .../app-server-protocol/src/protocol/v2.rs | 15 +- codex-rs/app-server/README.md | 4 +- .../app-server/src/codex_message_processor.rs | 86 +++++-- .../app-server/tests/suite/v2/skills_list.rs | 1 + .../app-server/tests/suite/v2/thread_start.rs | 35 +++ .../app-server/tests/suite/v2/turn_start.rs | 238 ++++++++++++++++++ codex-rs/core/src/agent/control.rs | 4 + codex-rs/core/src/agent/control_tests.rs | 3 + codex-rs/core/src/codex_delegate.rs | 7 +- codex-rs/core/src/environment_selection.rs | 111 ++++++++ codex-rs/core/src/lib.rs | 2 + codex-rs/core/src/session/handlers.rs | 13 +- codex-rs/core/src/session/mod.rs | 12 +- codex-rs/core/src/session/session.rs | 13 +- codex-rs/core/src/session/tests.rs | 168 +++++++++---- .../core/src/session/tests/guardian_tests.rs | 1 + codex-rs/core/src/session/turn_context.rs | 91 ++++--- codex-rs/core/src/thread_manager.rs | 114 +++++++-- codex-rs/core/src/thread_manager_tests.rs | 140 +++++++++++ .../core/src/tools/handlers/agent_jobs.rs | 15 +- .../src/tools/handlers/multi_agents/spawn.rs | 39 +-- .../tools/handlers/multi_agents_v2/spawn.rs | 7 + codex-rs/exec-server/src/environment.rs | 13 + codex-rs/exec-server/src/lib.rs | 2 + codex-rs/protocol/src/protocol.rs | 4 +- 26 files changed, 988 insertions(+), 165 deletions(-) create mode 100644 codex-rs/core/src/environment_selection.rs diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json index 0c5f21764..ae599d3e6 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json @@ -388,6 +388,21 @@ "clear" ], "type": "string" + }, + "TurnEnvironmentParams": { + "properties": { + "cwd": { + "$ref": "#/definitions/AbsolutePathBuf" + }, + "environmentId": { + "type": "string" + } + }, + "required": [ + "cwd", + "environmentId" + ], + "type": "object" } }, "properties": { diff --git a/codex-rs/app-server-protocol/src/protocol/v2.rs b/codex-rs/app-server-protocol/src/protocol/v2.rs index 97165a509..a30306b34 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2.rs @@ -3313,6 +3313,15 @@ pub struct ThreadStartParams { pub ephemeral: Option, #[ts(optional = nullable)] pub session_start_source: Option, + /// Optional sticky environments for this thread. + /// + /// Omitted selects the default environment when environment access is + /// enabled. Empty disables environment access for turns that do not + /// provide a turn override. Non-empty selects the first environment as the + /// current turn environment. + #[experimental("thread/start.environments")] + #[ts(optional = nullable)] + pub environments: Option>, #[experimental("thread/start.dynamicTools")] #[ts(optional = nullable)] pub dynamic_tools: Option>, @@ -4993,7 +5002,11 @@ pub struct TurnStartParams { #[experimental("turn/start.responsesapiClientMetadata")] #[ts(optional = nullable)] pub responsesapi_client_metadata: Option>, - /// Optional turn-scoped environment selections. + /// Optional turn-scoped environments. + /// + /// Omitted uses the thread sticky environments. Empty disables + /// environment access for this turn. Non-empty selects the first + /// environment as the current turn environment for this turn. #[experimental("turn/start.environments")] #[ts(optional = nullable)] pub environments: Option>, diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index e46d785b3..a28021bb9 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -142,7 +142,7 @@ Example with notification opt-out: ## API Overview -- `thread/start` — create a new thread; emits `thread/started` (including the current `thread.status`) and auto-subscribes you to turn/item events for that thread. When the request includes a `cwd` and the resolved sandbox is `workspace-write` or full access, app-server also marks that project as trusted in the user `config.toml`. Pass `sessionStartSource: "clear"` when starting a replacement thread after clearing the current session so `SessionStart` hooks receive `source: "clear"` instead of the default `"startup"`. For permissions, prefer `permissionProfile`; the legacy `sandbox` shorthand is still accepted but cannot be combined with `permissionProfile`. +- `thread/start` — create a new thread; emits `thread/started` (including the current `thread.status`) and auto-subscribes you to turn/item events for that thread. When the request includes a `cwd` and the resolved sandbox is `workspace-write` or full access, app-server also marks that project as trusted in the user `config.toml`. Pass `sessionStartSource: "clear"` when starting a replacement thread after clearing the current session so `SessionStart` hooks receive `source: "clear"` instead of the default `"startup"`. For permissions, prefer `permissionProfile`; the legacy `sandbox` shorthand is still accepted but cannot be combined with `permissionProfile`. Experimental `environments` selects the sticky execution environments for turns on the thread; omit it to use the server default, pass `[]` to disable environments, or pass explicit environment ids with per-environment `cwd`. - `thread/resume` — reopen an existing thread by id so subsequent `turn/start` calls append to it. Accepts the same permission override rules as `thread/start`. - `thread/fork` — fork an existing thread into a new thread id by copying the stored history; if the source thread is currently mid-turn, the fork records the same interruption marker as `turn/interrupt` instead of inheriting an unmarked partial turn suffix. The returned `thread.forkedFromId` points at the source thread when known. Accepts `ephemeral: true` for an in-memory temporary fork, emits `thread/started` (including the current `thread.status`), and auto-subscribes you to turn/item events for the new thread. Pass `excludeTurns: true` when the client plans to page fork history via `thread/turns/list` instead of receiving the full turn array immediately. Accepts the same permission override rules as `thread/start`. - `thread/list` — page through stored rollouts; supports cursor-based pagination and optional `modelProviders`, `sourceKinds`, `archived`, `cwd`, and `searchTerm` filters. Each returned `thread` includes `status` (`ThreadStatus`), defaulting to `notLoaded` when the thread is not currently loaded. @@ -541,7 +541,7 @@ Turns attach user input (text or images) to a thread and trigger Codex generatio - `{"type":"image","url":"https://…png"}` - `{"type":"localImage","path":"/tmp/screenshot.png"}` -You can optionally specify config overrides on the new turn. If specified, these settings become the default for subsequent turns on the same thread. `outputSchema` applies only to the current turn. +You can optionally specify config overrides on the new turn. If specified, these settings become the default for subsequent turns on the same thread. `outputSchema` applies only to the current turn. Experimental `environments` is turn-scoped: omit it to inherit the thread's sticky environments, pass `[]` to run the turn with no environments, or pass explicit environment ids to override the sticky selection for this turn only. `approvalsReviewer` accepts: diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 9d671996e..347767e78 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -226,6 +226,7 @@ use codex_core::ForkSnapshot; use codex_core::NewThread; use codex_core::RolloutRecorder; use codex_core::SessionMeta; +use codex_core::StartThreadWithToolsOptions; use codex_core::SteerInputError; use codex_core::ThreadConfigSnapshot; use codex_core::ThreadManager; @@ -665,6 +666,13 @@ fn configured_thread_store(config: &Config) -> Arc { } } +fn environment_selection_error_message(err: CodexErr) -> String { + match err { + CodexErr::InvalidRequest(message) => message, + err => err.to_string(), + } +} + impl CodexMessageProcessor { async fn instruction_sources_from_config(config: &Config) -> Vec { codex_core::AgentsMdManager::new(config) @@ -2431,6 +2439,7 @@ impl CodexMessageProcessor { personality, ephemeral, session_start_source, + environments, persist_extended_history, } = params; if sandbox.is_some() && permission_profile.is_some() { @@ -2441,6 +2450,24 @@ impl CodexMessageProcessor { .await; return; } + let environments = environments.map(|environments| { + environments + .into_iter() + .map(|environment| TurnEnvironmentSelection { + environment_id: environment.environment_id, + cwd: environment.cwd, + }) + .collect::>() + }); + if let Some(environments) = environments.as_ref() + && let Err(err) = self + .thread_manager + .validate_environment_selections(environments) + { + self.send_invalid_request_error(request_id, environment_selection_error_message(err)) + .await; + return; + } let mut typesafe_overrides = self.build_thread_config_overrides( model, model_provider, @@ -2479,6 +2506,7 @@ impl CodexMessageProcessor { typesafe_overrides, dynamic_tools, session_start_source, + environments, persist_extended_history, service_name, experimental_raw_events, @@ -2553,6 +2581,7 @@ impl CodexMessageProcessor { typesafe_overrides: ConfigOverrides, dynamic_tools: Option>, session_start_source: Option, + environments: Option>, persist_extended_history: bool, service_name: Option, experimental_raw_events: bool, @@ -2652,6 +2681,11 @@ impl CodexMessageProcessor { } let instruction_sources = Self::instruction_sources_from_config(&config).await; + let environments = environments.unwrap_or_else(|| { + listener_task_context + .thread_manager + .default_environment_selections(&config.cwd) + }); let dynamic_tools = dynamic_tools.unwrap_or_default(); let core_dynamic_tools = if dynamic_tools.is_empty() { Vec::new() @@ -2683,19 +2717,20 @@ impl CodexMessageProcessor { match listener_task_context .thread_manager - .start_thread_with_tools_and_service_name( + .start_thread_with_tools_and_service_name(StartThreadWithToolsOptions { config, - match session_start_source + initial_history: match session_start_source .unwrap_or(codex_app_server_protocol::ThreadStartSource::Startup) { codex_app_server_protocol::ThreadStartSource::Startup => InitialHistory::New, codex_app_server_protocol::ThreadStartSource::Clear => InitialHistory::Cleared, }, - core_dynamic_tools, + dynamic_tools: core_dynamic_tools, persist_extended_history, - service_name, - request_trace, - ) + metrics_service_name: service_name, + parent_trace: request_trace, + environments, + }) .instrument(tracing::info_span!( "app_server.thread_start.create_thread", otel.name = "app_server.thread_start.create_thread", @@ -2827,6 +2862,17 @@ impl CodexMessageProcessor { )) .await; } + Err(CodexErr::InvalidRequest(message)) => { + let error = JSONRPCErrorError { + code: INVALID_REQUEST_ERROR_CODE, + message, + data: None, + }; + listener_task_context + .outgoing + .send_error(request_id, error) + .await; + } Err(err) => { let error = JSONRPCErrorError { code: INTERNAL_ERROR_CODE, @@ -6949,15 +6995,25 @@ impl CodexMessageProcessor { let collaboration_mode = params.collaboration_mode.map(|mode| { self.normalize_turn_start_collaboration_mode(mode, collaboration_modes_config) }); - let environments = params.environments.map(|environments| { - environments - .into_iter() - .map(|environment| TurnEnvironmentSelection { - environment_id: environment.environment_id, - cwd: environment.cwd, - }) - .collect() - }); + let environments: Option> = + params.environments.map(|environments| { + environments + .into_iter() + .map(|environment| TurnEnvironmentSelection { + environment_id: environment.environment_id, + cwd: environment.cwd, + }) + .collect() + }); + if let Some(environments) = environments.as_ref() + && let Err(err) = self + .thread_manager + .validate_environment_selections(environments) + { + self.send_invalid_request_error(request_id, environment_selection_error_message(err)) + .await; + return; + } // Map v2 input items to core input items. let mapped_items: Vec = params diff --git a/codex-rs/app-server/tests/suite/v2/skills_list.rs b/codex-rs/app-server/tests/suite/v2/skills_list.rs index 195c4a86d..0d3bf4b49 100644 --- a/codex-rs/app-server/tests/suite/v2/skills_list.rs +++ b/codex-rs/app-server/tests/suite/v2/skills_list.rs @@ -324,6 +324,7 @@ async fn skills_changed_notification_is_emitted_after_skill_change() -> Result<( ephemeral: None, session_start_source: None, dynamic_tools: None, + environments: None, mock_experimental_field: None, experimental_raw_events: false, persist_extended_history: false, diff --git a/codex-rs/app-server/tests/suite/v2/thread_start.rs b/codex-rs/app-server/tests/suite/v2/thread_start.rs index 2aae14e20..3177003dd 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_start.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_start.rs @@ -19,6 +19,7 @@ use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::ThreadStartedNotification; use codex_app_server_protocol::ThreadStatus; use codex_app_server_protocol::ThreadStatusChangedNotification; +use codex_app_server_protocol::TurnEnvironmentParams; use codex_config::types::AuthCredentialsStoreMode; use codex_core::config::set_project_trust_level; use codex_core::config_loader::project_trust_key; @@ -48,6 +49,7 @@ use super::analytics::thread_initialized_event; use super::analytics::wait_for_analytics_payload; const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); +const INVALID_REQUEST_ERROR_CODE: i64 = -32600; #[tokio::test] async fn thread_start_creates_thread_and_emits_started() -> Result<()> { @@ -166,6 +168,39 @@ async fn thread_start_creates_thread_and_emits_started() -> Result<()> { Ok(()) } +#[tokio::test] +async fn thread_start_rejects_unknown_environment_as_invalid_request() -> Result<()> { + let server = create_mock_responses_server_repeating_assistant("Done").await; + + let codex_home = TempDir::new()?; + create_config_toml_without_approval_policy(codex_home.path(), &server.uri())?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp + .send_thread_start_request(ThreadStartParams { + environments: Some(vec![TurnEnvironmentParams { + environment_id: "missing".to_string(), + cwd: codex_home.path().to_path_buf().try_into()?, + }]), + ..Default::default() + }) + .await?; + + let error: JSONRPCError = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + + assert_eq!(error.id, RequestId::Integer(request_id)); + assert_eq!(error.error.code, INVALID_REQUEST_ERROR_CODE); + assert_eq!(error.error.message, "unknown turn environment id `missing`"); + + Ok(()) +} + #[tokio::test] async fn thread_start_response_includes_loaded_instruction_sources() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; diff --git a/codex-rs/app-server/tests/suite/v2/turn_start.rs b/codex-rs/app-server/tests/suite/v2/turn_start.rs index c6eb58a43..4a9b2c6a1 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start.rs @@ -5,6 +5,7 @@ use app_test_support::create_apply_patch_sse_response; use app_test_support::create_exec_command_sse_response; use app_test_support::create_fake_rollout; use app_test_support::create_final_assistant_message_sse_response; +use app_test_support::create_mock_responses_server_repeating_assistant; use app_test_support::create_mock_responses_server_sequence; use app_test_support::create_mock_responses_server_sequence_unchecked; use app_test_support::create_shell_command_sse_response; @@ -47,6 +48,7 @@ use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnCompletedNotification; +use codex_app_server_protocol::TurnEnvironmentParams; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::TurnStartedNotification; @@ -820,6 +822,69 @@ async fn turn_start_rejects_invalid_permission_profile_before_starting_turn() -> Ok(()) } +#[tokio::test] +async fn turn_start_rejects_unknown_environment_before_starting_turn() -> Result<()> { + let server = create_mock_responses_server_repeating_assistant("Done").await; + let codex_home = TempDir::new()?; + create_config_toml( + codex_home.path(), + &server.uri(), + "never", + &BTreeMap::default(), + )?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let thread_req = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + ..Default::default() + }) + .await?; + let thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + + let turn_req = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id, + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + environments: Some(vec![TurnEnvironmentParams { + environment_id: "missing".to_string(), + cwd: codex_home.path().to_path_buf().try_into()?, + }]), + ..Default::default() + }) + .await?; + let err: JSONRPCError = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(turn_req)), + ) + .await??; + + assert_eq!(err.id, RequestId::Integer(turn_req)); + assert_eq!(err.error.code, INVALID_REQUEST_ERROR_CODE); + assert_eq!(err.error.message, "unknown turn environment id `missing`"); + let turn_started = tokio::time::timeout( + std::time::Duration::from_millis(250), + mcp.read_stream_until_notification_message("turn/started"), + ) + .await; + assert!( + turn_started.is_err(), + "did not expect a turn/started notification after rejected environments" + ); + + Ok(()) +} + #[tokio::test] async fn turn_start_emits_notifications_and_accepts_model_override() -> Result<()> { // Provide a mock server and config so model wiring is valid. @@ -1926,6 +1991,179 @@ async fn turn_start_updates_sandbox_and_cwd_between_turns_v2() -> Result<()> { Ok(()) } +#[tokio::test] +async fn turn_start_resolves_sticky_thread_environments_and_turn_overrides() -> Result<()> { + let tmp = TempDir::new()?; + let codex_home = tmp.path().join("codex_home"); + std::fs::create_dir(&codex_home)?; + let workspace = tmp.path().join("workspace"); + std::fs::create_dir(&workspace)?; + + let server = create_mock_responses_server_repeating_assistant("done").await; + create_config_toml(&codex_home, &server.uri(), "never", &BTreeMap::default())?; + + let mut mcp = McpProcess::new_with_env( + &codex_home, + &[("CODEX_EXEC_SERVER_URL", Some("http://127.0.0.1:1"))], + ) + .await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + for case in [ + EnvironmentSelectionCase { + name: "sticky_unset_turn_unset", + sticky: None, + turn: None, + }, + EnvironmentSelectionCase { + name: "sticky_empty_turn_unset", + sticky: Some(&[]), + turn: None, + }, + EnvironmentSelectionCase { + name: "sticky_local_turn_unset", + sticky: Some(&["local"]), + turn: None, + }, + EnvironmentSelectionCase { + name: "sticky_remote_turn_unset", + sticky: Some(&["remote"]), + turn: None, + }, + EnvironmentSelectionCase { + name: "sticky_local_remote_turn_unset", + sticky: Some(&["local", "remote"]), + turn: None, + }, + EnvironmentSelectionCase { + name: "sticky_local_turn_empty", + sticky: Some(&["local"]), + turn: Some(&[]), + }, + EnvironmentSelectionCase { + name: "sticky_empty_turn_local", + sticky: Some(&[]), + turn: Some(&["local"]), + }, + EnvironmentSelectionCase { + name: "sticky_local_turn_remote", + sticky: Some(&["local"]), + turn: Some(&["remote"]), + }, + EnvironmentSelectionCase { + name: "sticky_remote_turn_local", + sticky: Some(&["remote"]), + turn: Some(&["local"]), + }, + EnvironmentSelectionCase { + name: "sticky_unset_turn_local_remote", + sticky: None, + turn: Some(&["local", "remote"]), + }, + ] { + run_environment_selection_case(&mut mcp, &workspace, case).await?; + } + + Ok(()) +} + +struct EnvironmentSelectionCase { + name: &'static str, + sticky: Option<&'static [&'static str]>, + turn: Option<&'static [&'static str]>, +} + +async fn run_environment_selection_case( + mcp: &mut McpProcess, + workspace: &Path, + case: EnvironmentSelectionCase, +) -> Result<()> { + let thread_req = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + cwd: Some(workspace.to_string_lossy().into_owned()), + environments: environment_params(case.sticky, workspace)?, + ..Default::default() + }) + .await?; + let thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + + let turn_req = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id.clone(), + input: vec![V2UserInput::Text { + text: format!("run {}", case.name), + text_elements: Vec::new(), + }], + environments: environment_params(case.turn, workspace)?, + cwd: Some(workspace.to_path_buf()), + model: Some("mock-model".to_string()), + ..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)?; + + let started_notification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/started"), + ) + .await??; + let started: TurnStartedNotification = serde_json::from_value( + started_notification + .params + .ok_or_else(|| anyhow::anyhow!("turn/started notification should include params"))?, + )?; + assert_eq!(started.turn.id, turn.id, "{}", case.name); + + let completed_notification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + let completed: TurnCompletedNotification = + serde_json::from_value(completed_notification.params.ok_or_else(|| { + anyhow::anyhow!("turn/completed notification should include params") + })?)?; + assert_eq!(completed.turn.id, turn.id, "{}", case.name); + assert_eq!( + completed.turn.status, + TurnStatus::Completed, + "{}", + case.name + ); + + mcp.clear_message_buffer(); + + Ok(()) +} + +fn environment_params( + ids: Option<&[&str]>, + cwd: &Path, +) -> Result>> { + ids.map(|ids| { + ids.iter() + .map(|id| { + Ok(TurnEnvironmentParams { + environment_id: (*id).to_string(), + cwd: cwd.to_path_buf().try_into()?, + }) + }) + .collect() + }) + .transpose() +} + #[tokio::test] async fn turn_start_file_change_approval_v2() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index c54d0663b..d4ec6858d 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -28,6 +28,7 @@ use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TokenUsage; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::user_input::UserInput; use codex_rollout::state_db; use codex_state::DirectionalThreadSpawnEdgeStatus; @@ -52,6 +53,7 @@ pub(crate) enum SpawnAgentForkMode { pub(crate) struct SpawnAgentOptions { pub(crate) fork_parent_spawn_call_id: Option, pub(crate) fork_mode: Option, + pub(crate) environments: Option>, } #[derive(Clone, Debug)] @@ -246,6 +248,7 @@ impl AgentControl { /*metrics_service_name*/ None, inherited_shell_snapshot, inherited_exec_policy, + options.environments.clone(), ) .await? } @@ -405,6 +408,7 @@ impl AgentControl { /*persist_extended_history*/ false, inherited_shell_snapshot, inherited_exec_policy, + options.environments.clone(), ) .await } diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index e83ee6d4b..6018c3747 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -657,6 +657,7 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { SpawnAgentOptions { fork_parent_spawn_call_id: Some(parent_spawn_call_id.clone()), fork_mode: Some(SpawnAgentForkMode::FullHistory), + ..Default::default() }, ) .await @@ -751,6 +752,7 @@ async fn spawn_agent_fork_flushes_parent_rollout_before_loading_history() { SpawnAgentOptions { fork_parent_spawn_call_id: Some(parent_spawn_call_id.clone()), fork_mode: Some(SpawnAgentForkMode::FullHistory), + ..Default::default() }, ) .await @@ -860,6 +862,7 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { SpawnAgentOptions { fork_parent_spawn_call_id: Some(parent_spawn_call_id.clone()), fork_mode: Some(SpawnAgentForkMode::LastNTurns(2)), + ..Default::default() }, ) .await diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index e3ac8b1cf..eb3876f60 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -47,6 +47,7 @@ use crate::session::SUBMISSION_CHANNEL_CAPACITY; use crate::session::emit_subagent_session_started; use crate::session::session::Session; use crate::session::turn_context::TurnContext; +use crate::session::turn_context::TurnEnvironment; use codex_login::AuthManager; use codex_models_manager::manager::ModelsManager; use codex_protocol::error::CodexErr; @@ -73,7 +74,6 @@ pub(crate) async fn run_codex_thread_interactive( ) -> Result { let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let (tx_ops, rx_ops) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); - let CodexSpawnOk { codex, .. } = Box::pin(Codex::spawn(CodexSpawnArgs { config, auth_manager, @@ -94,6 +94,11 @@ pub(crate) async fn run_codex_thread_interactive( inherited_exec_policy: Some(Arc::clone(&parent_session.services.exec_policy)), parent_rollout_thread_trace: codex_rollout_trace::ThreadTraceContext::disabled(), parent_trace: None, + environments: parent_ctx + .environments + .iter() + .map(TurnEnvironment::selection) + .collect(), analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), thread_store: Arc::clone(&parent_session.services.thread_store), })) diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs new file mode 100644 index 000000000..7f93ba384 --- /dev/null +++ b/codex-rs/core/src/environment_selection.rs @@ -0,0 +1,111 @@ +use std::sync::Arc; + +use codex_exec_server::Environment; +use codex_exec_server::EnvironmentManager; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; +use codex_protocol::protocol::TurnEnvironmentSelection; +use codex_utils_absolute_path::AbsolutePathBuf; + +pub(crate) fn default_thread_environment_selections( + environment_manager: &EnvironmentManager, + cwd: &AbsolutePathBuf, +) -> Vec { + environment_manager + .default_environment_id() + .map(|environment_id| TurnEnvironmentSelection { + environment_id: environment_id.to_string(), + cwd: cwd.clone(), + }) + .into_iter() + .collect() +} + +pub(crate) fn validate_environment_selections( + environment_manager: &EnvironmentManager, + environments: &[TurnEnvironmentSelection], +) -> CodexResult<()> { + for selected_environment in environments { + if environment_manager + .get_environment(&selected_environment.environment_id) + .is_none() + { + return Err(CodexErr::InvalidRequest(format!( + "unknown turn environment id `{}`", + selected_environment.environment_id + ))); + } + } + + Ok(()) +} + +pub(crate) fn selected_primary_environment( + environment_manager: &EnvironmentManager, + environments: &[TurnEnvironmentSelection], +) -> CodexResult>> { + environments + .first() + .map(|selected_environment| { + environment_manager + .get_environment(&selected_environment.environment_id) + .ok_or_else(|| { + CodexErr::InvalidRequest(format!( + "unknown turn environment id `{}`", + selected_environment.environment_id + )) + }) + }) + .transpose() +} + +#[cfg(test)] +mod tests { + use codex_exec_server::EnvironmentManagerArgs; + use codex_exec_server::ExecServerRuntimePaths; + use codex_exec_server::REMOTE_ENVIRONMENT_ID; + use codex_protocol::protocol::TurnEnvironmentSelection; + use codex_utils_absolute_path::AbsolutePathBuf; + use pretty_assertions::assert_eq; + + use super::*; + + fn test_runtime_paths() -> ExecServerRuntimePaths { + ExecServerRuntimePaths::new( + std::env::current_exe().expect("current exe"), + /*codex_linux_sandbox_exe*/ None, + ) + .expect("runtime paths") + } + + #[tokio::test] + async fn default_thread_environment_selections_use_manager_default_id() { + let cwd = AbsolutePathBuf::current_dir().expect("cwd"); + let manager = EnvironmentManager::new(EnvironmentManagerArgs { + exec_server_url: Some("ws://127.0.0.1:8765".to_string()), + local_runtime_paths: test_runtime_paths(), + }); + + assert_eq!( + default_thread_environment_selections(&manager, &cwd), + vec![TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd, + }] + ); + } + + #[tokio::test] + async fn default_thread_environment_selections_empty_when_default_disabled() { + let cwd = AbsolutePathBuf::current_dir().expect("cwd"); + let manager = EnvironmentManager::new(EnvironmentManagerArgs { + exec_server_url: Some("none".to_string()), + local_runtime_paths: test_runtime_paths(), + }); + + assert_eq!( + default_thread_environment_selections(&manager, &cwd), + Vec::::new() + ); + } +} diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 548287958..54fadc6fd 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -29,6 +29,7 @@ pub mod config_loader; pub mod connectors; pub mod context; mod context_manager; +mod environment_selection; pub mod exec; pub mod exec_env; mod exec_policy; @@ -118,6 +119,7 @@ pub(crate) mod web_search; pub(crate) mod windows_sandbox_read_grants; pub use thread_manager::ForkSnapshot; pub use thread_manager::NewThread; +pub use thread_manager::StartThreadWithToolsOptions; pub use thread_manager::ThreadManager; pub use thread_manager::build_models_manager; pub use web_search::web_search_action_detail; diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 206c48fba..16f4f9b6b 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -127,7 +127,7 @@ pub(super) async fn user_input_or_turn_inner( op: Op, mirror_user_text_to_realtime: Option<()>, ) { - let (items, updates, responsesapi_client_metadata, environments) = match op { + let (items, updates, responsesapi_client_metadata) = match op { Op::UserTurn { cwd, approval_policy, @@ -167,12 +167,12 @@ pub(super) async fn user_input_or_turn_inner( reasoning_summary: summary, service_tier, final_output_json_schema: Some(final_output_json_schema), + environments, personality, app_server_client_name: None, app_server_client_version: None, }, None, - environments, ) } Op::UserInputWithTurnContext { @@ -217,12 +217,12 @@ pub(super) async fn user_input_or_turn_inner( reasoning_summary: summary, service_tier, final_output_json_schema: Some(final_output_json_schema), + environments, personality, app_server_client_name: None, app_server_client_version: None, }, responsesapi_client_metadata, - environments, ) } Op::UserInput { @@ -234,18 +234,15 @@ pub(super) async fn user_input_or_turn_inner( items, SessionSettingsUpdate { final_output_json_schema: Some(final_output_json_schema), + environments, ..Default::default() }, responsesapi_client_metadata, - environments, ), _ => unreachable!(), }; - let Ok(current_context) = sess - .new_turn_with_sub_id(sub_id.clone(), updates, environments) - .await - else { + let Ok(current_context) = sess.new_turn_with_sub_id(sub_id.clone(), updates).await else { // new_turn_with_sub_id already emits the error event. return; }; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 8365fba35..e5c675f16 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -29,6 +29,8 @@ use crate::context::NetworkRuleSaved; use crate::context::PermissionsInstructions; use crate::context::PersonalitySpecInstructions; use crate::default_skill_metadata_budget; +use crate::environment_selection::selected_primary_environment; +use crate::environment_selection::validate_environment_selections; use crate::exec_policy::ExecPolicyManager; use crate::installation_id::resolve_installation_id; use crate::parse_turn_item; @@ -110,6 +112,7 @@ use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnContextItem; use codex_protocol::protocol::TurnContextNetworkItem; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::W3cTraceContext; use codex_protocol::request_permissions::PermissionGrantScope; use codex_protocol::request_permissions::RequestPermissionProfile; @@ -409,6 +412,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) parent_rollout_thread_trace: ThreadTraceContext, pub(crate) user_shell_override: Option, pub(crate) parent_trace: Option, + pub(crate) environments: Vec, pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, } @@ -465,13 +469,15 @@ impl Codex { inherited_exec_policy, parent_rollout_thread_trace, parent_trace: _, + environments, analytics_events_client, thread_store, } = args; let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let (tx_event, rx_event) = async_channel::unbounded(); - - let environment = environment_manager.default_environment(); + validate_environment_selections(environment_manager.as_ref(), &environments)?; + let environment = + selected_primary_environment(environment_manager.as_ref(), &environments)?; let fs = environment .as_ref() .map(|environment| environment.get_filesystem()); @@ -598,7 +604,6 @@ impl Codex { } else { dynamic_tools }; - // TODO (aibrahim): Consolidate config.model and config.model_reasoning_effort into config.collaboration_mode // to avoid extracting these fields separately and constructing CollaborationMode here. let collaboration_mode = CollaborationMode { @@ -637,6 +642,7 @@ impl Codex { cwd: config.cwd.clone(), codex_home: config.codex_home.clone(), thread_name: None, + environments, original_config_do_not_use: Arc::clone(&config), metrics_service_name, app_server_client_name: None, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 512fd1af5..2918e21a0 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -71,6 +71,8 @@ pub(crate) struct SessionConfiguration { pub(super) codex_home: AbsolutePathBuf, /// Optional user-facing name for the thread, updated during the session. pub(super) thread_name: Option, + /// Sticky environments for turns that do not provide a turn-local override. + pub(super) environments: Vec, // TODO(pakrym): Remove config from here pub(super) original_config_do_not_use: Arc, @@ -159,7 +161,12 @@ impl SessionConfiguration { .unwrap_or_else(|| self.cwd.clone()); let cwd_changed = absolute_cwd.as_path() != self.cwd.as_path(); - next_configuration.cwd = absolute_cwd; + next_configuration.cwd = absolute_cwd.clone(); + if cwd_changed + && let Some(primary_environment) = next_configuration.environments.first_mut() + { + primary_environment.cwd = absolute_cwd; + } if let Some(permission_profile) = updates.permission_profile.clone() { let sandbox_policy = permission_profile @@ -238,6 +245,10 @@ pub(crate) struct SessionSettingsUpdate { pub(crate) reasoning_summary: Option, pub(crate) service_tier: Option>, pub(crate) final_output_json_schema: Option>, + /// Turn-local environment override. `None` inherits the sticky thread + /// environments stored on `SessionConfiguration`; `Some([])` explicitly + /// disables environments for this turn. + pub(crate) environments: Option>, pub(crate) personality: Option, pub(crate) app_server_client_name: Option, pub(crate) app_server_client_version: Option, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 0f7930907..1b18c8378 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -1,3 +1,4 @@ +use super::turn_context::TurnEnvironment; use super::*; use crate::config::ConfigBuilder; use crate::config::test_config; @@ -781,7 +782,6 @@ async fn new_turn_refreshes_managed_network_proxy_for_sandbox_change() -> anyhow sandbox_policy: Some(SandboxPolicy::DangerFullAccess), ..Default::default() }, - /*environment_selections*/ None, ) .await?; @@ -2239,6 +2239,7 @@ async fn set_rate_limits_retains_previous_credits() { cwd: config.cwd.clone(), codex_home: config.codex_home.clone(), thread_name: None, + environments: Vec::new(), original_config_do_not_use: Arc::clone(&config), metrics_service_name: None, app_server_client_name: None, @@ -2344,6 +2345,7 @@ async fn set_rate_limits_updates_plan_type_when_present() { cwd: config.cwd.clone(), codex_home: config.codex_home.clone(), thread_name: None, + environments: Vec::new(), original_config_do_not_use: Arc::clone(&config), metrics_service_name: None, app_server_client_name: None, @@ -2794,6 +2796,7 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati cwd: config.cwd.clone(), codex_home: config.codex_home.clone(), thread_name: None, + environments: Vec::new(), original_config_do_not_use: Arc::clone(&config), metrics_service_name: None, app_server_client_name: None, @@ -2806,6 +2809,17 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati } } +fn turn_environments_for_tests( + environment: &Arc, + cwd: &codex_utils_absolute_path::AbsolutePathBuf, +) -> Vec { + vec![TurnEnvironment { + environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(), + environment: Arc::clone(environment), + cwd: cwd.clone(), + }] +} + #[tokio::test] async fn session_configuration_apply_preserves_split_file_system_policy_on_cwd_only_update() { let mut session_configuration = make_session_configuration_for_tests().await; @@ -3111,6 +3125,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_zsh_path() { cwd: config.cwd.clone(), codex_home: config.codex_home.clone(), thread_name: None, + environments: Vec::new(), original_config_do_not_use: Arc::clone(&config), metrics_service_name: None, app_server_client_name: None, @@ -3193,6 +3208,10 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { developer_instructions: None, }, }; + let default_environments = vec![TurnEnvironmentSelection { + environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(), + cwd: config.cwd.clone(), + }]; let session_configuration = SessionConfiguration { provider: config.model_provider.clone(), collaboration_mode, @@ -3215,6 +3234,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { cwd: config.cwd.clone(), codex_home: config.codex_home.clone(), thread_name: None, + environments: default_environments, original_config_do_not_use: Arc::clone(&config), metrics_service_name: None, app_server_client_name: None, @@ -3331,6 +3351,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { .skills_for_config(&skills_input, Some(Arc::clone(&skill_fs))) .await, ); + let turn_environments = turn_environments_for_tests(&environment, &session_configuration.cwd); let turn_context = Session::make_turn_context( conversation_id, Some(Arc::clone(&auth_manager)), @@ -3345,7 +3366,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { &models_manager, /*network*/ None, Some(environment), - /*environments*/ None, + turn_environments, session_configuration.cwd.clone(), "turn_id".to_string(), Arc::clone(&js_repl), @@ -3410,6 +3431,10 @@ async fn make_session_with_config_and_rx( developer_instructions: None, }, }; + let default_environments = vec![TurnEnvironmentSelection { + environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(), + cwd: config.cwd.clone(), + }]; let session_configuration = SessionConfiguration { provider: config.model_provider.clone(), collaboration_mode, @@ -3432,6 +3457,7 @@ async fn make_session_with_config_and_rx( cwd: config.cwd.clone(), codex_home: config.codex_home.clone(), thread_name: None, + environments: default_environments, original_config_do_not_use: Arc::clone(&config), metrics_service_name: None, app_server_client_name: None, @@ -4059,7 +4085,7 @@ async fn user_turn_updates_approvals_reviewer() { } #[tokio::test] -async fn turn_environment_selection_sets_primary_environment() { +async fn turn_environments_set_primary_environment() { let (session, _turn_context, _rx) = make_session_and_context_with_rx().await; let selected_cwd = AbsolutePathBuf::try_from(session.get_config().await.cwd.as_path().join("selected")) @@ -4068,21 +4094,19 @@ async fn turn_environment_selection_sets_primary_environment() { let turn_context = session .new_turn_with_sub_id( "sub-1".to_string(), - SessionSettingsUpdate::default(), - Some(vec![codex_protocol::protocol::TurnEnvironmentSelection { - environment_id: "local".to_string(), - cwd: selected_cwd.clone(), - }]), + SessionSettingsUpdate { + environments: Some(vec![TurnEnvironmentSelection { + environment_id: "local".to_string(), + cwd: selected_cwd.clone(), + }]), + ..Default::default() + }, ) .await .expect("turn should start"); - let turn_environments = turn_context - .environments - .as_ref() - .expect("turn environments should be recorded"); + let turn_environments = &turn_context.environments; assert_eq!(turn_environments.len(), 1); - assert_eq!(turn_environments[0].environment_id, "local"); assert!(std::sync::Arc::ptr_eq( turn_context .environment @@ -4095,7 +4119,55 @@ async fn turn_environment_selection_sets_primary_environment() { } #[tokio::test] -async fn multiple_turn_environment_selections_use_first_as_primary_environment() { +async fn default_turn_uses_stored_thread_environments() { + let (session, _turn_context, _rx) = make_session_and_context_with_rx().await; + let selected_cwd = + AbsolutePathBuf::try_from(session.get_config().await.cwd.as_path().join("selected")) + .expect("absolute path"); + + { + let mut state = session.state.lock().await; + state.session_configuration.environments = vec![TurnEnvironmentSelection { + environment_id: "local".to_string(), + cwd: selected_cwd.clone(), + }]; + } + + let turn_context = session.new_default_turn().await; + + let turn_environments = &turn_context.environments; + assert_eq!(turn_environments.len(), 1); + assert!(std::sync::Arc::ptr_eq( + turn_context + .environment + .as_ref() + .expect("primary environment should be set"), + &turn_environments[0].environment + )); + assert_eq!(turn_context.cwd, selected_cwd); + assert_eq!(turn_context.config.cwd, selected_cwd); +} + +#[tokio::test] +async fn default_turn_honors_empty_stored_thread_environments() { + let (session, _turn_context, _rx) = make_session_and_context_with_rx().await; + let session_cwd = session.get_config().await.cwd.clone(); + + { + let mut state = session.state.lock().await; + state.session_configuration.environments = Vec::new(); + } + + let turn_context = session.new_default_turn().await; + + assert!(turn_context.environment.is_none()); + assert_eq!(turn_context.cwd, session_cwd); + assert_eq!(turn_context.config.cwd, session_cwd); + assert_eq!(turn_context.environments.len(), 0); +} + +#[tokio::test] +async fn multiple_turn_environments_use_first_as_primary_environment() { let (session, _turn_context, _rx) = make_session_and_context_with_rx().await; let session_cwd = session.get_config().await.cwd.clone(); let first_cwd = @@ -4106,25 +4178,24 @@ async fn multiple_turn_environment_selections_use_first_as_primary_environment() let turn_context = session .new_turn_with_sub_id( "sub-1".to_string(), - SessionSettingsUpdate::default(), - Some(vec![ - codex_protocol::protocol::TurnEnvironmentSelection { - environment_id: "local".to_string(), - cwd: first_cwd.clone(), - }, - codex_protocol::protocol::TurnEnvironmentSelection { - environment_id: "local".to_string(), - cwd: second_cwd.clone(), - }, - ]), + SessionSettingsUpdate { + environments: Some(vec![ + TurnEnvironmentSelection { + environment_id: "local".to_string(), + cwd: first_cwd.clone(), + }, + TurnEnvironmentSelection { + environment_id: "local".to_string(), + cwd: second_cwd.clone(), + }, + ]), + ..Default::default() + }, ) .await .expect("turn should start"); - let turn_environments = turn_context - .environments - .as_ref() - .expect("turn environments should be recorded"); + let turn_environments = &turn_context.environments; assert_eq!(turn_environments.len(), 2); assert_eq!(turn_environments[0].cwd, first_cwd); assert_eq!(turn_environments[1].cwd, second_cwd); @@ -4140,14 +4211,16 @@ async fn multiple_turn_environment_selections_use_first_as_primary_environment() } #[tokio::test] -async fn empty_turn_environment_selection_clears_primary_environment() { +async fn empty_turn_environments_clear_primary_environment() { let (session, _turn_context, _rx) = make_session_and_context_with_rx().await; let turn_context = session .new_turn_with_sub_id( "sub-1".to_string(), - SessionSettingsUpdate::default(), - Some(vec![]), + SessionSettingsUpdate { + environments: Some(vec![]), + ..Default::default() + }, ) .await .expect("turn should start"); @@ -4155,28 +4228,23 @@ async fn empty_turn_environment_selection_clears_primary_environment() { assert!(turn_context.environment.is_none()); assert_eq!(turn_context.cwd, session.get_config().await.cwd); assert_eq!(turn_context.config.cwd, session.get_config().await.cwd); - assert_eq!( - turn_context - .environments - .as_ref() - .expect("turn environments should be recorded") - .len(), - 0 - ); + assert_eq!(turn_context.environments.len(), 0); } #[tokio::test] -async fn unknown_turn_environment_selection_returns_error() { +async fn unknown_turn_environment_returns_error() { let (session, _turn_context, _rx) = make_session_and_context_with_rx().await; let err = session .new_turn_with_sub_id( "sub-1".to_string(), - SessionSettingsUpdate::default(), - Some(vec![codex_protocol::protocol::TurnEnvironmentSelection { - environment_id: "missing".to_string(), - cwd: session.get_config().await.cwd.clone(), - }]), + SessionSettingsUpdate { + environments: Some(vec![TurnEnvironmentSelection { + environment_id: "missing".to_string(), + cwd: session.get_config().await.cwd.clone(), + }]), + ..Default::default() + }, ) .await .expect_err("unknown environment should fail"); @@ -4509,6 +4577,10 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( developer_instructions: None, }, }; + let default_environments = vec![TurnEnvironmentSelection { + environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(), + cwd: config.cwd.clone(), + }]; let session_configuration = SessionConfiguration { provider: config.model_provider.clone(), collaboration_mode, @@ -4531,6 +4603,7 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( cwd: config.cwd.clone(), codex_home: config.codex_home.clone(), thread_name: None, + environments: default_environments, original_config_do_not_use: Arc::clone(&config), metrics_service_name: None, app_server_client_name: None, @@ -4647,6 +4720,7 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( .skills_for_config(&skills_input, Some(Arc::clone(&skill_fs))) .await, ); + let turn_environments = turn_environments_for_tests(&environment, &session_configuration.cwd); let turn_context = Arc::new(Session::make_turn_context( conversation_id, Some(Arc::clone(&auth_manager)), @@ -4661,7 +4735,7 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( &models_manager, /*network*/ None, Some(environment), - /*environments*/ None, + turn_environments, session_configuration.cwd.clone(), "turn_id".to_string(), Arc::clone(&js_repl), diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 84865190d..7844070ac 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -775,6 +775,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { parent_rollout_thread_trace: codex_rollout_trace::ThreadTraceContext::disabled(), user_shell_override: None, parent_trace: None, + environments: Vec::new(), analytics_events_client: None, thread_store, }) diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index e9ecb66e7..14e3ce4b7 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -27,12 +27,20 @@ impl TurnSkillsContext { #[derive(Clone, Debug)] pub(crate) struct TurnEnvironment { - #[allow(dead_code)] pub(crate) environment_id: String, pub(crate) environment: Arc, pub(crate) cwd: AbsolutePathBuf, } +impl TurnEnvironment { + pub(crate) fn selection(&self) -> TurnEnvironmentSelection { + TurnEnvironmentSelection { + environment_id: self.environment_id.clone(), + cwd: self.cwd.clone(), + } + } +} + /// The context needed for a single turn of the thread. #[derive(Debug)] pub(crate) struct TurnContext { @@ -48,7 +56,7 @@ pub(crate) struct TurnContext { pub(crate) reasoning_summary: ReasoningSummaryConfig, pub(crate) session_source: SessionSource, pub(crate) environment: Option>, - pub(crate) environments: Option>, + pub(crate) environments: Vec, /// The session's absolute working directory. All relative paths provided /// by the model as well as sandbox policies are resolved against this path /// instead of `std::env::current_dir()`. @@ -376,7 +384,7 @@ impl Session { models_manager: &ModelsManager, network: Option, environment: Option>, - environments: Option>, + environments: Vec, cwd: AbsolutePathBuf, sub_id: String, js_repl: Arc, @@ -483,26 +491,17 @@ impl Session { &self, sub_id: String, updates: SessionSettingsUpdate, - environment_selections: Option>, ) -> CodexResult> { - let turn_environments = match self.resolve_turn_environments(environment_selections) { - Ok(turn_environments) => turn_environments, - Err(err) => { - self.send_event_raw(Event { - id: sub_id.clone(), - msg: EventMsg::Error(ErrorEvent { - message: err.to_string(), - codex_error_info: Some(CodexErrorInfo::BadRequest), - }), - }) - .await; - return Err(err); - } - }; - let update_result = { + let update_result: CodexResult<_> = { let mut state = self.state.lock().await; match state.session_configuration.clone().apply(&updates) { Ok(next) => { + let effective_environments = updates + .environments + .clone() + .unwrap_or_else(|| next.environments.clone()); + let turn_environments = + self.resolve_turn_environments(&effective_environments)?; let previous_cwd = state.session_configuration.cwd.clone(); let sandbox_policy_changed = state.session_configuration.sandbox_policy != next.sandbox_policy; @@ -511,18 +510,20 @@ impl Session { state.session_configuration = next.clone(); Ok(( next, + turn_environments, sandbox_policy_changed, previous_cwd, codex_home, session_source, )) } - Err(err) => Err(err), + Err(err) => Err(CodexErr::InvalidRequest(err.to_string())), } }; let ( session_configuration, + turn_environments, sandbox_policy_changed, previous_cwd, codex_home, @@ -567,33 +568,29 @@ impl Session { fn resolve_turn_environments( &self, - environment_selections: Option>, - ) -> CodexResult>> { - let Some(environment_selections) = environment_selections else { - return Ok(None); - }; - - let mut turn_environments = Vec::with_capacity(environment_selections.len()); - for environment_selection in environment_selections { + environments: &[TurnEnvironmentSelection], + ) -> CodexResult> { + let mut turn_environments = Vec::with_capacity(environments.len()); + for selected_environment in environments { + let environment_id = selected_environment.environment_id.clone(); let environment = self .services .environment_manager - .get_environment(&environment_selection.environment_id) + .get_environment(&environment_id) .ok_or_else(|| { CodexErr::InvalidRequest(format!( - "unknown turn environment id `{}`", - environment_selection.environment_id + "unknown turn environment id `{environment_id}`" )) })?; - let cwd = environment_selection.cwd; + let cwd = selected_environment.cwd.clone(); turn_environments.push(TurnEnvironment { - environment_id: environment_selection.environment_id, + environment_id, environment, cwd, }); } - Ok(Some(turn_environments)) + Ok(turn_environments) } async fn new_turn_from_configuration( @@ -601,18 +598,11 @@ impl Session { sub_id: String, session_configuration: SessionConfiguration, final_output_json_schema: Option>, - turn_environments: Option>, + turn_environments: Vec, ) -> Arc { - // `None` means use the thread's default environment. `Some([])` is an - // explicit no-environment turn, so do not fall back in that case. - let primary_turn_environment = turn_environments - .as_ref() - .and_then(|turn_environments| turn_environments.first()); - let environment = match primary_turn_environment { - Some(turn_environment) => Some(Arc::clone(&turn_environment.environment)), - None if turn_environments.is_some() => None, - None => self.services.environment_manager.default_environment(), - }; + let primary_turn_environment = turn_environments.first(); + let environment = primary_turn_environment + .map(|turn_environment| Arc::clone(&turn_environment.environment)); let cwd = primary_turn_environment .map(|turn_environment| turn_environment.cwd.clone()) .unwrap_or_else(|| session_configuration.cwd.clone()); @@ -710,11 +700,20 @@ impl Session { let state = self.state.lock().await; state.session_configuration.clone() }; + let turn_environments = + match self.resolve_turn_environments(&session_configuration.environments) { + Ok(turn_environments) => turn_environments, + Err(err) => { + warn!("failed to resolve stored session environments: {err}"); + Vec::new() + } + }; + self.new_turn_from_configuration( sub_id, session_configuration, /*final_output_json_schema*/ None, - /*turn_environments*/ None, + turn_environments, ) .await } diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 1d4bdfe6f..3988360b9 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -2,6 +2,9 @@ use crate::SkillsManager; use crate::agent::AgentControl; use crate::codex_thread::CodexThread; use crate::config::Config; +use crate::environment_selection::default_thread_environment_selections; +use crate::environment_selection::selected_primary_environment; +use crate::environment_selection::validate_environment_selections; use crate::file_watcher::FileWatcher; use crate::mcp::McpManager; use crate::plugins::PluginsManager; @@ -44,6 +47,7 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::W3cTraceContext; use codex_rollout::RolloutConfig; use codex_state::DirectionalThreadSpawnEdgeStatus; @@ -203,6 +207,16 @@ pub struct ThreadManager { _test_codex_home_guard: Option, } +pub struct StartThreadWithToolsOptions { + pub config: Config, + pub initial_history: InitialHistory, + pub dynamic_tools: Vec, + pub persist_extended_history: bool, + pub metrics_service_name: Option, + pub parent_trace: Option, + pub environments: Vec, +} + /// Shared, `Arc`-owned state for [`ThreadManager`]. This `Arc` is required to have a single /// `Arc` reference that can be downgraded to by `AgentControl` while preventing every single /// function to require an `Arc<&Self>`. @@ -394,6 +408,20 @@ impl ThreadManager { self.state.environment_manager.clone() } + pub fn default_environment_selections( + &self, + cwd: &AbsolutePathBuf, + ) -> Vec { + default_thread_environment_selections(self.state.environment_manager.as_ref(), cwd) + } + + pub fn validate_environment_selections( + &self, + environments: &[TurnEnvironmentSelection], + ) -> CodexResult<()> { + validate_environment_selections(self.state.environment_manager.as_ref(), environments) + } + pub fn get_models_manager(&self) -> Arc { self.state.models_manager.clone() } @@ -506,37 +534,40 @@ impl ThreadManager { dynamic_tools: Vec, persist_extended_history: bool, ) -> CodexResult { - Box::pin(self.start_thread_with_tools_and_service_name( - config, - InitialHistory::New, - dynamic_tools, - persist_extended_history, - /*metrics_service_name*/ None, - /*parent_trace*/ None, - )) + let environments = default_thread_environment_selections( + self.state.environment_manager.as_ref(), + &config.cwd, + ); + Box::pin( + self.start_thread_with_tools_and_service_name(StartThreadWithToolsOptions { + config, + initial_history: InitialHistory::New, + dynamic_tools, + persist_extended_history, + metrics_service_name: None, + parent_trace: None, + environments, + }), + ) .await } pub async fn start_thread_with_tools_and_service_name( &self, - config: Config, - initial_history: InitialHistory, - dynamic_tools: Vec, - persist_extended_history: bool, - metrics_service_name: Option, - parent_trace: Option, + options: StartThreadWithToolsOptions, ) -> CodexResult { - let thread_store = configured_thread_store(&config); + let thread_store = configured_thread_store(&options.config); Box::pin(self.state.spawn_thread( - config, + options.config, thread_store, - initial_history, + options.initial_history, Arc::clone(&self.state.auth_manager), self.agent_control(), - dynamic_tools, - persist_extended_history, - metrics_service_name, - parent_trace, + options.dynamic_tools, + options.persist_extended_history, + options.metrics_service_name, + options.parent_trace, + options.environments, /*user_shell_override*/ None, )) .await @@ -569,6 +600,10 @@ impl ThreadManager { parent_trace: Option, ) -> CodexResult { let thread_store = configured_thread_store(&config); + let environments = default_thread_environment_selections( + self.state.environment_manager.as_ref(), + &config.cwd, + ); Box::pin(self.state.spawn_thread( config, thread_store, @@ -579,6 +614,7 @@ impl ThreadManager { persist_extended_history, /*metrics_service_name*/ None, parent_trace, + environments, /*user_shell_override*/ None, )) .await @@ -590,6 +626,10 @@ impl ThreadManager { user_shell_override: crate::shell::Shell, ) -> CodexResult { let thread_store = configured_thread_store(&config); + let environments = default_thread_environment_selections( + self.state.environment_manager.as_ref(), + &config.cwd, + ); Box::pin(self.state.spawn_thread( config, thread_store, @@ -600,6 +640,7 @@ impl ThreadManager { /*persist_extended_history*/ false, /*metrics_service_name*/ None, /*parent_trace*/ None, + environments, /*user_shell_override*/ Some(user_shell_override), )) .await @@ -614,6 +655,10 @@ impl ThreadManager { ) -> CodexResult { let initial_history = RolloutRecorder::get_rollout_history(&rollout_path).await?; let thread_store = configured_thread_store(&config); + let environments = default_thread_environment_selections( + self.state.environment_manager.as_ref(), + &config.cwd, + ); Box::pin(self.state.spawn_thread( config, thread_store, @@ -624,6 +669,7 @@ impl ThreadManager { /*persist_extended_history*/ false, /*metrics_service_name*/ None, /*parent_trace*/ None, + environments, /*user_shell_override*/ Some(user_shell_override), )) .await @@ -724,6 +770,10 @@ impl ThreadManager { } }; let thread_store = configured_thread_store(&config); + let environments = default_thread_environment_selections( + self.state.environment_manager.as_ref(), + &config.cwd, + ); Box::pin(self.state.spawn_thread( config, thread_store, @@ -734,6 +784,7 @@ impl ThreadManager { persist_extended_history, /*metrics_service_name*/ None, parent_trace, + environments, /*user_shell_override*/ None, )) .await @@ -808,6 +859,7 @@ impl ThreadManagerState { /*metrics_service_name*/ None, /*inherited_shell_snapshot*/ None, /*inherited_exec_policy*/ None, + /*environments*/ None, )) .await } @@ -822,8 +874,12 @@ impl ThreadManagerState { metrics_service_name: Option, inherited_shell_snapshot: Option>, inherited_exec_policy: Option>, + environments: Option>, ) -> CodexResult { let thread_store = configured_thread_store(&config); + let environments = environments.unwrap_or_else(|| { + default_thread_environment_selections(self.environment_manager.as_ref(), &config.cwd) + }); Box::pin(self.spawn_thread_with_source( config, thread_store, @@ -837,6 +893,7 @@ impl ThreadManagerState { inherited_shell_snapshot, inherited_exec_policy, /*parent_trace*/ None, + environments, /*user_shell_override*/ None, )) .await @@ -853,6 +910,8 @@ impl ThreadManagerState { ) -> CodexResult { let initial_history = RolloutRecorder::get_rollout_history(&rollout_path).await?; let thread_store = configured_thread_store(&config); + let environments = + default_thread_environment_selections(self.environment_manager.as_ref(), &config.cwd); Box::pin(self.spawn_thread_with_source( config, thread_store, @@ -866,6 +925,7 @@ impl ThreadManagerState { inherited_shell_snapshot, inherited_exec_policy, /*parent_trace*/ None, + environments, /*user_shell_override*/ None, )) .await @@ -881,8 +941,12 @@ impl ThreadManagerState { persist_extended_history: bool, inherited_shell_snapshot: Option>, inherited_exec_policy: Option>, + environments: Option>, ) -> CodexResult { let thread_store = configured_thread_store(&config); + let environments = environments.unwrap_or_else(|| { + default_thread_environment_selections(self.environment_manager.as_ref(), &config.cwd) + }); Box::pin(self.spawn_thread_with_source( config, thread_store, @@ -896,6 +960,7 @@ impl ThreadManagerState { inherited_shell_snapshot, inherited_exec_policy, /*parent_trace*/ None, + environments, /*user_shell_override*/ None, )) .await @@ -914,6 +979,7 @@ impl ThreadManagerState { persist_extended_history: bool, metrics_service_name: Option, parent_trace: Option, + environments: Vec, user_shell_override: Option, ) -> CodexResult { Box::pin(self.spawn_thread_with_source( @@ -929,6 +995,7 @@ impl ThreadManagerState { /*inherited_shell_snapshot*/ None, /*inherited_exec_policy*/ None, parent_trace, + environments, user_shell_override, )) .await @@ -949,9 +1016,11 @@ impl ThreadManagerState { inherited_shell_snapshot: Option>, inherited_exec_policy: Option>, parent_trace: Option, + environments: Vec, user_shell_override: Option, ) -> CodexResult { - let environment = self.environment_manager.default_environment(); + let environment = + selected_primary_environment(self.environment_manager.as_ref(), &environments)?; let watch_registration = match environment.as_ref() { Some(environment) if !environment.is_remote() => { self.skills_watcher @@ -990,6 +1059,7 @@ impl ThreadManagerState { parent_rollout_thread_trace, user_shell_override, parent_trace, + environments, analytics_events_client: self.analytics_events_client.clone(), thread_store, }) diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 4dcc29f56..dc2cb004f 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -1,6 +1,7 @@ use super::*; use crate::config::test_config; use crate::rollout::RolloutRecorder; +use crate::session::session::SessionSettingsUpdate; use crate::session::tests::make_session_and_context; use crate::tasks::interrupted_turn_history_marker; use codex_models_manager::collaboration_mode_presets::CollaborationModesConfig; @@ -43,6 +44,20 @@ fn assistant_msg(text: &str) -> ResponseItem { } } +fn disabled_environment_manager_for_tests() -> Arc { + let runtime_paths = codex_exec_server::ExecServerRuntimePaths::new( + std::env::current_exe().expect("current exe path"), + /*codex_linux_sandbox_exe*/ None, + ) + .expect("runtime paths"); + Arc::new(codex_exec_server::EnvironmentManager::new( + codex_exec_server::EnvironmentManagerArgs { + exec_server_url: Some("none".to_string()), + local_runtime_paths: runtime_paths, + }, + )) +} + #[test] fn truncates_before_requested_user_message() { let items = [ @@ -271,6 +286,131 @@ async fn shutdown_all_threads_bounded_submits_shutdown_to_every_thread() { assert!(manager.list_thread_ids().await.is_empty()); } +#[tokio::test] +async fn start_thread_accepts_explicit_environment_when_default_environment_is_disabled() { + let temp_dir = tempdir().expect("tempdir"); + let mut config = test_config().await; + config.codex_home = temp_dir.path().join("codex-home").abs(); + config.cwd = config.codex_home.abs(); + std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + + let manager = ThreadManager::with_models_provider_and_home_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.to_path_buf(), + disabled_environment_manager_for_tests(), + ); + + let thread = manager + .start_thread_with_tools_and_service_name(StartThreadWithToolsOptions { + config: config.clone(), + initial_history: InitialHistory::New, + dynamic_tools: Vec::new(), + persist_extended_history: false, + metrics_service_name: None, + parent_trace: None, + environments: vec![TurnEnvironmentSelection { + environment_id: "local".to_string(), + cwd: config.cwd.clone(), + }], + }) + .await + .expect("explicit sticky environment should resolve by id"); + + assert_eq!(manager.list_thread_ids().await, vec![thread.thread_id]); +} + +#[tokio::test] +async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { + let temp_dir = tempdir().expect("tempdir"); + let mut config = test_config().await; + config.codex_home = temp_dir.path().join("codex-home").abs(); + config.cwd = config.codex_home.abs(); + std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + + let auth_manager = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let manager = ThreadManager::new( + &config, + auth_manager.clone(), + SessionSource::Exec, + CollaborationModesConfig::default(), + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + /*analytics_events_client*/ None, + ); + let selected_cwd = + AbsolutePathBuf::try_from(config.cwd.as_path().join("selected")).expect("absolute path"); + let environments = vec![TurnEnvironmentSelection { + environment_id: "local".to_string(), + cwd: selected_cwd.clone(), + }]; + let default_cwd = config.cwd.clone(); + + let source = manager + .start_thread_with_tools_and_service_name(StartThreadWithToolsOptions { + config: config.clone(), + initial_history: InitialHistory::New, + dynamic_tools: Vec::new(), + persist_extended_history: false, + metrics_service_name: None, + parent_trace: None, + environments: environments.clone(), + }) + .await + .expect("start source thread"); + source.thread.ensure_rollout_materialized().await; + source + .thread + .flush_rollout() + .await + .expect("flush source rollout"); + let rollout_path = source + .thread + .rollout_path() + .expect("source rollout path should exist"); + + let resumed = manager + .resume_thread_from_rollout( + config.clone(), + rollout_path.clone(), + auth_manager, + /*parent_trace*/ None, + ) + .await + .expect("resume source thread"); + let resumed_turn = resumed + .thread + .codex + .session + .new_turn_with_sub_id("resume-turn".to_string(), SessionSettingsUpdate::default()) + .await + .expect("build resumed turn context"); + assert_eq!(resumed_turn.environments.len(), 1); + assert_eq!(resumed_turn.environments[0].cwd, default_cwd); + assert_ne!(resumed_turn.environments[0].cwd, selected_cwd); + + let forked = manager + .fork_thread( + ForkSnapshot::Interrupted, + config, + rollout_path, + /*persist_extended_history*/ false, + /*parent_trace*/ None, + ) + .await + .expect("fork source thread"); + let forked_turn = forked + .thread + .codex + .session + .new_turn_with_sub_id("fork-turn".to_string(), SessionSettingsUpdate::default()) + .await + .expect("build forked turn context"); + assert_eq!(forked_turn.environments.len(), 1); + assert_eq!(forked_turn.environments[0].cwd, default_cwd); + assert_ne!(forked_turn.environments[0].cwd, selected_cwd); +} + #[tokio::test] async fn new_uses_configured_openai_provider_for_model_refresh() { let server = MockServer::start().await; diff --git a/codex-rs/core/src/tools/handlers/agent_jobs.rs b/codex-rs/core/src/tools/handlers/agent_jobs.rs index e3b00960b..adf777fff 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs.rs @@ -1,3 +1,4 @@ +use crate::agent::control::SpawnAgentOptions; use crate::agent::exceeds_thread_spawn_depth_limit; use crate::agent::next_thread_spawn_depth; use crate::agent::status::is_final; @@ -5,6 +6,7 @@ use crate::config::Config; use crate::function_tool::FunctionCallError; use crate::session::session::Session; use crate::session::turn_context::TurnContext; +use crate::session::turn_context::TurnEnvironment; use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; @@ -629,16 +631,25 @@ async fn run_agent_job_loop( let thread_id = match session .services .agent_control - .spawn_agent( + .spawn_agent_with_metadata( options.spawn_config.clone(), items.into(), Some(SessionSource::SubAgent(SubAgentSource::Other(format!( "agent_job:{job_id}" )))), + SpawnAgentOptions { + environments: Some( + turn.environments + .iter() + .map(TurnEnvironment::selection) + .collect(), + ), + ..Default::default() + }, ) .await { - Ok(thread_id) => thread_id, + Ok(spawned_agent) => spawned_agent.thread_id, Err(CodexErr::AgentLimitReached { .. }) => { db.mark_agent_job_item_pending( job_id.as_str(), diff --git a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs index f1be8951e..777cb9be1 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs @@ -6,6 +6,7 @@ use crate::agent::exceeds_thread_spawn_depth_limit; use crate::agent::next_thread_spawn_depth; use crate::agent::role::DEFAULT_ROLE_NAME; use crate::agent::role::apply_role_to_config; +use crate::session::turn_context::TurnEnvironment; pub(crate) struct Handler; @@ -82,21 +83,29 @@ impl ToolHandler for Handler { apply_spawn_agent_runtime_overrides(&mut config, turn.as_ref())?; apply_spawn_agent_overrides(&mut config, child_depth); - let result = Box::pin(session.services.agent_control.spawn_agent_with_metadata( - config, - input_items, - Some(thread_spawn_source( - session.conversation_id, - &turn.session_source, - child_depth, - role_name, - /*task_name*/ None, - )?), - SpawnAgentOptions { - fork_parent_spawn_call_id: args.fork_context.then(|| call_id.clone()), - fork_mode: args.fork_context.then_some(SpawnAgentForkMode::FullHistory), - }, - )) + let result = Box::pin( + session.services.agent_control.spawn_agent_with_metadata( + config, + input_items, + Some(thread_spawn_source( + session.conversation_id, + &turn.session_source, + child_depth, + role_name, + /*task_name*/ None, + )?), + SpawnAgentOptions { + fork_parent_spawn_call_id: args.fork_context.then(|| call_id.clone()), + fork_mode: args.fork_context.then_some(SpawnAgentForkMode::FullHistory), + environments: Some( + turn.environments + .iter() + .map(TurnEnvironment::selection) + .collect(), + ), + }, + ), + ) .await .map_err(collab_spawn_error); let (new_thread_id, new_agent_metadata, status) = match &result { diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs index 00986311a..21b4638c0 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs @@ -5,6 +5,7 @@ use crate::agent::control::render_input_preview; use crate::agent::next_thread_spawn_depth; use crate::agent::role::DEFAULT_ROLE_NAME; use crate::agent::role::apply_role_to_config; +use crate::session::turn_context::TurnEnvironment; use codex_protocol::AgentPath; use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::Op; @@ -123,6 +124,12 @@ impl ToolHandler for Handler { SpawnAgentOptions { fork_parent_spawn_call_id: fork_mode.as_ref().map(|_| call_id.clone()), fork_mode, + environments: Some( + turn.environments + .iter() + .map(TurnEnvironment::selection) + .collect(), + ), }, ) .await diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 9e4c69c41..377a7a38d 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -115,6 +115,11 @@ impl EnvironmentManager { .and_then(|environment_id| self.get_environment(environment_id)) } + /// Returns the id of the default environment. + pub fn default_environment_id(&self) -> Option<&str> { + self.default_environment.as_deref() + } + /// Returns the local environment instance used for internal runtime work. pub fn local_environment(&self) -> Arc { match self.get_environment(LOCAL_ENVIRONMENT_ID) { @@ -304,6 +309,7 @@ mod tests { }); let environment = manager.default_environment().expect("default environment"); + assert_eq!(manager.default_environment_id(), Some(LOCAL_ENVIRONMENT_ID)); assert!(!environment.is_remote()); assert!( !manager @@ -322,6 +328,7 @@ mod tests { }); assert!(manager.default_environment().is_none()); + assert_eq!(manager.default_environment_id(), None); assert!( !manager .get_environment(LOCAL_ENVIRONMENT_ID) @@ -339,6 +346,10 @@ mod tests { }); let environment = manager.default_environment().expect("default environment"); + assert_eq!( + manager.default_environment_id(), + Some(REMOTE_ENVIRONMENT_ID) + ); assert!(environment.is_remote()); assert_eq!(environment.exec_server_url(), Some("ws://127.0.0.1:8765")); assert!(Arc::ptr_eq( @@ -399,6 +410,7 @@ mod tests { }); assert!(manager.default_environment().is_none()); + assert_eq!(manager.default_environment_id(), None); } #[tokio::test] @@ -409,6 +421,7 @@ mod tests { }); assert!(manager.default_environment().is_none()); + assert_eq!(manager.default_environment_id(), None); assert!( !manager .get_environment(LOCAL_ENVIRONMENT_ID) diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 3da7aa73f..a585a4b29 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -29,6 +29,8 @@ pub use environment::CODEX_EXEC_SERVER_URL_ENV_VAR; pub use environment::Environment; pub use environment::EnvironmentManager; pub use environment::EnvironmentManagerArgs; +pub use environment::LOCAL_ENVIRONMENT_ID; +pub use environment::REMOTE_ENVIRONMENT_ID; pub use file_system::CopyOptions; pub use file_system::CreateDirectoryOptions; pub use file_system::ExecutorFileSystem; diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 1f17b53c1..6e0a62fc4 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -432,7 +432,7 @@ pub enum Op { UserInput { /// User input items, see `InputItem` items: Vec, - /// Optional turn-scoped environment selections. + /// Optional turn-scoped environments. #[serde(default, skip_serializing_if = "Option::is_none")] environments: Option>, /// Optional JSON Schema used to constrain the final assistant message for this turn. @@ -579,7 +579,7 @@ pub enum Op { #[serde(skip_serializing_if = "Option::is_none")] personality: Option, - /// Optional turn-scoped environment selections. + /// Optional turn-scoped environments. #[serde(default, skip_serializing_if = "Option::is_none")] environments: Option>, },