From 31a9c54dbaf4fc24041c60a8c858355c3301343d Mon Sep 17 00:00:00 2001 From: Ben Schoepke Date: Tue, 3 Mar 2026 10:44:24 -0800 Subject: [PATCH] initial vibe code --- codex-rs/core/config.schema.json | 22 + codex-rs/core/src/agent/role.rs | 1 + codex-rs/core/src/codex.rs | 27 + codex-rs/core/src/config/mod.rs | 16 + codex-rs/core/src/config/profile.rs | 2 + codex-rs/core/src/features.rs | 14 + codex-rs/core/src/project_doc.rs | 39 + codex-rs/core/src/tools/context.rs | 1 + codex-rs/core/src/tools/handlers/mod.rs | 3 + codex-rs/core/src/tools/handlers/ps_repl.rs | 363 ++++ codex-rs/core/src/tools/mod.rs | 1 + codex-rs/core/src/tools/ps_repl/kernel.ps1 | 187 ++ codex-rs/core/src/tools/ps_repl/mod.rs | 1686 +++++++++++++++++++ codex-rs/core/src/tools/spec.rs | 111 ++ codex-rs/core/tests/suite/mod.rs | 1 + codex-rs/core/tests/suite/ps_repl.rs | 484 ++++++ codex-rs/exec/src/lib.rs | 1 + docs/ps_repl.md | 68 + 18 files changed, 3027 insertions(+) create mode 100644 codex-rs/core/src/tools/handlers/ps_repl.rs create mode 100644 codex-rs/core/src/tools/ps_repl/kernel.ps1 create mode 100644 codex-rs/core/src/tools/ps_repl/mod.rs create mode 100644 codex-rs/core/tests/suite/ps_repl.rs create mode 100644 docs/ps_repl.md diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 689538102c..b5efa7326d 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -379,6 +379,9 @@ "prevent_idle_sleep": { "type": "boolean" }, + "ps_repl": { + "type": "boolean" + }, "realtime_conversation": { "type": "boolean" }, @@ -502,6 +505,14 @@ "plan_mode_reasoning_effort": { "$ref": "#/definitions/ReasoningEffort" }, + "ps_repl_path": { + "allOf": [ + { + "$ref": "#/definitions/AbsolutePathBuf" + } + ], + "description": "Optional absolute path to the PowerShell 7 runtime used by `ps_repl`." + }, "sandbox_mode": { "$ref": "#/definitions/SandboxMode" }, @@ -1750,6 +1761,9 @@ "prevent_idle_sleep": { "type": "boolean" }, + "ps_repl": { + "type": "boolean" + }, "realtime_conversation": { "type": "boolean" }, @@ -2088,6 +2102,14 @@ }, "type": "object" }, + "ps_repl_path": { + "allOf": [ + { + "$ref": "#/definitions/AbsolutePathBuf" + } + ], + "description": "Optional absolute path to the PowerShell 7 runtime used by `ps_repl`." + }, "review_model": { "description": "Review model override used by the `/review` feature.", "type": "string" diff --git a/codex-rs/core/src/agent/role.rs b/codex-rs/core/src/agent/role.rs index 0ff5a7d560..9a07f0a690 100644 --- a/codex-rs/core/src/agent/role.rs +++ b/codex-rs/core/src/agent/role.rs @@ -88,6 +88,7 @@ pub(crate) async fn apply_role_to_config( cwd: Some(config.cwd.clone()), codex_linux_sandbox_exe: config.codex_linux_sandbox_exe.clone(), main_execve_wrapper_exe: config.main_execve_wrapper_exe.clone(), + ps_repl_path: config.ps_repl_path.clone(), js_repl_node_path: config.js_repl_node_path.clone(), ..Default::default() }, diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 05981029c5..ddff5b1644 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -277,6 +277,8 @@ use crate::tools::network_approval::NetworkApprovalService; use crate::tools::network_approval::build_blocked_request_observer; use crate::tools::network_approval::build_network_policy_decider; use crate::tools::parallel::ToolCallRuntime; +use crate::tools::ps_repl::PsReplHandle; +use crate::tools::ps_repl::resolve_compatible_pwsh; use crate::tools::sandboxing::ApprovalStore; use crate::tools::spec::ToolsConfig; use crate::tools::spec::ToolsConfigParams; @@ -378,6 +380,16 @@ impl Codex { config.features.disable(Feature::JsReplToolsOnly); config.startup_warnings.push(message); } + if config.features.enabled(Feature::PsRepl) + && let Err(err) = resolve_compatible_pwsh(config.ps_repl_path.as_deref()).await + { + let message = format!( + "Disabled `ps_repl` for this session because the configured PowerShell runtime is unavailable or incompatible. {err}" + ); + warn!("{message}"); + config.features.disable(Feature::PsRepl); + config.startup_warnings.push(message); + } let allowed_skills_for_implicit_invocation = loaded_skills.allowed_skills_for_implicit_invocation(); @@ -608,6 +620,7 @@ pub(crate) struct Session { pub(crate) conversation: Arc, pub(crate) active_turn: Mutex>, pub(crate) services: SessionServices, + ps_repl: Arc, js_repl: Arc, next_internal_sub_id: AtomicU64, } @@ -663,6 +676,7 @@ pub(crate) struct TurnContext { pub(crate) codex_linux_sandbox_exe: Option, pub(crate) tool_call_gate: Arc, pub(crate) truncation_policy: TruncationPolicy, + pub(crate) ps_repl: Arc, pub(crate) js_repl: Arc, pub(crate) dynamic_tools: Vec, pub(crate) turn_metadata_state: Arc, @@ -751,6 +765,7 @@ impl TurnContext { codex_linux_sandbox_exe: self.codex_linux_sandbox_exe.clone(), tool_call_gate: Arc::new(ReadinessFlag::new()), truncation_policy, + ps_repl: Arc::clone(&self.ps_repl), js_repl: Arc::clone(&self.js_repl), dynamic_tools: self.dynamic_tools.clone(), turn_metadata_state: self.turn_metadata_state.clone(), @@ -1056,6 +1071,7 @@ impl Session { model_info: ModelInfo, network: Option, sub_id: String, + ps_repl: Arc, js_repl: Arc, skills_outcome: Arc, ) -> TurnContext { @@ -1125,6 +1141,7 @@ impl Session { codex_linux_sandbox_exe: per_turn_config.codex_linux_sandbox_exe.clone(), tool_call_gate: Arc::new(ReadinessFlag::new()), truncation_policy: model_info.truncation_policy.into(), + ps_repl, js_repl, dynamic_tools: session_configuration.dynamic_tools.clone(), turn_metadata_state, @@ -1519,6 +1536,7 @@ impl Session { config.js_repl_node_path.clone(), config.js_repl_node_module_dirs.clone(), )); + let ps_repl = Arc::new(PsReplHandle::with_pwsh_path(config.ps_repl_path.clone())); let sess = Arc::new(Session { conversation_id, @@ -1530,6 +1548,7 @@ impl Session { conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), services, + ps_repl, js_repl, next_internal_sub_id: AtomicU64::new(0), }); @@ -2167,6 +2186,7 @@ impl Session { .as_ref() .map(StartedNetworkProxy::proxy), sub_id, + Arc::clone(&self.ps_repl), Arc::clone(&self.js_repl), skills_outcome, ); @@ -4695,6 +4715,7 @@ async fn spawn_review_thread( final_output_json_schema: None, codex_linux_sandbox_exe: parent_turn_context.codex_linux_sandbox_exe.clone(), tool_call_gate: Arc::new(ReadinessFlag::new()), + ps_repl: Arc::clone(&sess.ps_repl), js_repl: Arc::clone(&sess.js_repl), dynamic_tools: parent_turn_context.dynamic_tools.clone(), truncation_policy: model_info.truncation_policy.into(), @@ -8334,6 +8355,7 @@ mod tests { config.js_repl_node_path.clone(), config.js_repl_node_module_dirs.clone(), )); + let ps_repl = Arc::new(PsReplHandle::with_pwsh_path(config.ps_repl_path.clone())); let skills_outcome = Arc::new(services.skills_manager.skills_for_config(&per_turn_config)); let turn_context = Session::make_turn_context( @@ -8345,6 +8367,7 @@ mod tests { model_info, None, "turn_id".to_string(), + Arc::clone(&ps_repl), Arc::clone(&js_repl), skills_outcome, ); @@ -8359,6 +8382,7 @@ mod tests { conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), services, + ps_repl, js_repl, next_internal_sub_id: AtomicU64::new(0), }; @@ -8502,6 +8526,7 @@ mod tests { config.js_repl_node_path.clone(), config.js_repl_node_module_dirs.clone(), )); + let ps_repl = Arc::new(PsReplHandle::with_pwsh_path(config.ps_repl_path.clone())); let skills_outcome = Arc::new(services.skills_manager.skills_for_config(&per_turn_config)); let turn_context = Arc::new(Session::make_turn_context( @@ -8513,6 +8538,7 @@ mod tests { model_info, None, "turn_id".to_string(), + Arc::clone(&ps_repl), Arc::clone(&js_repl), skills_outcome, )); @@ -8527,6 +8553,7 @@ mod tests { conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), services, + ps_repl, js_repl, next_internal_sub_id: AtomicU64::new(0), }); diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index b5520a2f08..78dceedb06 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -394,6 +394,9 @@ pub struct Config { /// code via [`ConfigOverrides`]. pub main_execve_wrapper_exe: Option, + /// Optional absolute path to the PowerShell 7 runtime used by `ps_repl`. + pub ps_repl_path: Option, + /// Optional absolute path to the Node runtime used by `js_repl`. pub js_repl_node_path: Option, @@ -1121,6 +1124,9 @@ pub struct ConfigToml { /// Default: `300000` (5 minutes). pub background_terminal_max_timeout: Option, + /// Optional absolute path to the PowerShell 7 runtime used by `ps_repl`. + pub ps_repl_path: Option, + /// Optional absolute path to the Node runtime used by `js_repl`. pub js_repl_node_path: Option, @@ -1553,6 +1559,7 @@ pub struct ConfigOverrides { pub config_profile: Option, pub codex_linux_sandbox_exe: Option, pub main_execve_wrapper_exe: Option, + pub ps_repl_path: Option, pub js_repl_node_path: Option, pub js_repl_node_module_dirs: Option>, pub zsh_path: Option, @@ -1683,6 +1690,7 @@ impl Config { config_profile: config_profile_key, codex_linux_sandbox_exe, main_execve_wrapper_exe, + ps_repl_path: ps_repl_path_override, js_repl_node_path: js_repl_node_path_override, js_repl_node_module_dirs: js_repl_node_module_dirs_override, zsh_path: zsh_path_override, @@ -1981,6 +1989,9 @@ impl Config { "experimental compact prompt file", )?; let compact_prompt = compact_prompt.or(file_compact_prompt); + let ps_repl_path = ps_repl_path_override + .or(config_profile.ps_repl_path.map(Into::into)) + .or(cfg.ps_repl_path.map(Into::into)); let js_repl_node_path = js_repl_node_path_override .or(config_profile.js_repl_node_path.map(Into::into)) .or(cfg.js_repl_node_path.map(Into::into)); @@ -2150,6 +2161,7 @@ impl Config { file_opener: cfg.file_opener.unwrap_or(UriBasedFileOpener::VsCode), codex_linux_sandbox_exe, main_execve_wrapper_exe, + ps_repl_path, js_repl_node_path, js_repl_node_module_dirs, zsh_path, @@ -4909,6 +4921,7 @@ model_verbosity = "high" file_opener: UriBasedFileOpener::VsCode, codex_linux_sandbox_exe: None, main_execve_wrapper_exe: None, + ps_repl_path: None, js_repl_node_path: None, js_repl_node_module_dirs: Vec::new(), zsh_path: None, @@ -5037,6 +5050,7 @@ model_verbosity = "high" file_opener: UriBasedFileOpener::VsCode, codex_linux_sandbox_exe: None, main_execve_wrapper_exe: None, + ps_repl_path: None, js_repl_node_path: None, js_repl_node_module_dirs: Vec::new(), zsh_path: None, @@ -5163,6 +5177,7 @@ model_verbosity = "high" file_opener: UriBasedFileOpener::VsCode, codex_linux_sandbox_exe: None, main_execve_wrapper_exe: None, + ps_repl_path: None, js_repl_node_path: None, js_repl_node_module_dirs: Vec::new(), zsh_path: None, @@ -5275,6 +5290,7 @@ model_verbosity = "high" file_opener: UriBasedFileOpener::VsCode, codex_linux_sandbox_exe: None, main_execve_wrapper_exe: None, + ps_repl_path: None, js_repl_node_path: None, js_repl_node_module_dirs: Vec::new(), zsh_path: None, diff --git a/codex-rs/core/src/config/profile.rs b/codex-rs/core/src/config/profile.rs index 61b02adcd0..5daa211b88 100644 --- a/codex-rs/core/src/config/profile.rs +++ b/codex-rs/core/src/config/profile.rs @@ -33,6 +33,8 @@ pub struct ConfigProfile { pub chatgpt_base_url: Option, /// Optional path to a file containing model instructions. pub model_instructions_file: Option, + /// Optional absolute path to the PowerShell 7 runtime used by `ps_repl`. + pub ps_repl_path: Option, pub js_repl_node_path: Option, /// Ordered list of directories to search for Node modules in `js_repl`. pub js_repl_node_module_dirs: Option>, diff --git a/codex-rs/core/src/features.rs b/codex-rs/core/src/features.rs index cd8af6aa18..f249394869 100644 --- a/codex-rs/core/src/features.rs +++ b/codex-rs/core/src/features.rs @@ -78,6 +78,8 @@ pub enum Feature { ShellTool, // Experimental + /// Enable PowerShell REPL tools backed by a persistent pwsh kernel. + PsRepl, /// Enable JavaScript REPL tools backed by a persistent Node kernel. JsRepl, /// Only expose js_repl tools directly to the model. @@ -460,6 +462,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Stable, default_enabled: true, }, + FeatureSpec { + id: Feature::PsRepl, + key: "ps_repl", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::JsRepl, key: "js_repl", @@ -825,6 +833,12 @@ mod tests { assert_eq!(Feature::JsRepl.default_enabled(), false); } + #[test] + fn ps_repl_is_under_development_and_disabled_by_default() { + assert_eq!(Feature::PsRepl.stage(), Stage::UnderDevelopment); + assert_eq!(Feature::PsRepl.default_enabled(), false); + } + #[test] fn collab_is_legacy_alias_for_multi_agent() { assert_eq!(feature_for_key("multi_agent"), Some(Feature::Collab)); diff --git a/codex-rs/core/src/project_doc.rs b/codex-rs/core/src/project_doc.rs index 8323b86a1b..02eec2fe7c 100644 --- a/codex-rs/core/src/project_doc.rs +++ b/codex-rs/core/src/project_doc.rs @@ -69,6 +69,25 @@ fn render_js_repl_instructions(config: &Config) -> Option { Some(section) } +fn render_ps_repl_instructions(config: &Config) -> Option { + if !config.features.enabled(Feature::PsRepl) { + return None; + } + + let mut section = String::from("## PowerShell REPL (pwsh)\n"); + section.push_str( + "- Use `ps_repl` for PowerShell-backed automation in a persistent `pwsh` kernel.\n", + ); + section.push_str("- `ps_repl` is a freeform/custom tool. Direct `ps_repl` calls must send raw PowerShell input (optionally with first-line `# codex-ps-repl: timeout_ms=15000`). Do not wrap code in JSON (for example `{\"code\":\"...\"}`), quotes, or markdown code fences.\n"); + section.push_str("- Helpers: `$CodexTmpDir`, `Invoke-CodexTool -Name -Arguments `, `$Codex.TmpDir`, and `$Codex.Tool(, )`.\n"); + section.push_str("- `Invoke-CodexTool` returns the raw tool output object. Use it for built-in tools, dynamic tools, and MCP tools.\n"); + section.push_str("- PowerShell session state persists across calls, including variables, functions, aliases, imported modules, environment changes, and `$LASTEXITCODE`. Reset the kernel with `ps_repl_reset` when needed.\n"); + section.push_str("- To share generated images with the model, write a file under `$CodexTmpDir`, call `Invoke-CodexTool -Name view_image -Arguments @{ path = \"/absolute/path\" }`, then delete the file.\n"); + section.push_str("- Avoid direct `[Console]::Write*`, raw StdOut/StdErr writes, or other host-level output that bypasses PowerShell streams; they can corrupt the JSON line protocol. Use pipeline output, `Write-Output`, `Write-Host`, `Write-Verbose`, or `Write-Warning` instead."); + + Some(section) +} + /// Combines `Config::instructions` and `AGENTS.md` (if present) into a single /// string of instructions. pub(crate) async fn get_user_instructions( @@ -103,6 +122,13 @@ pub(crate) async fn get_user_instructions( output.push_str(&js_repl_section); } + if let Some(ps_repl_section) = render_ps_repl_instructions(config) { + if !output.is_empty() { + output.push_str("\n\n"); + } + output.push_str(&ps_repl_section); + } + let skills_section = skills.and_then(render_skills_section); if let Some(skills_section) = skills_section { if !output.is_empty() { @@ -492,6 +518,19 @@ mod tests { assert_eq!(res, expected); } + #[tokio::test] + async fn ps_repl_instructions_are_appended_when_enabled() { + let tmp = tempfile::tempdir().expect("tempdir"); + let mut cfg = make_config(&tmp, 4096, None).await; + cfg.features.enable(Feature::PsRepl); + + let res = get_user_instructions(&cfg, None) + .await + .expect("ps_repl instructions expected"); + let expected = "## PowerShell REPL (pwsh)\n- Use `ps_repl` for PowerShell-backed automation in a persistent `pwsh` kernel.\n- `ps_repl` is a freeform/custom tool. Direct `ps_repl` calls must send raw PowerShell input (optionally with first-line `# codex-ps-repl: timeout_ms=15000`). Do not wrap code in JSON (for example `{\"code\":\"...\"}`), quotes, or markdown code fences.\n- Helpers: `$CodexTmpDir`, `Invoke-CodexTool -Name -Arguments `, `$Codex.TmpDir`, and `$Codex.Tool(, )`.\n- `Invoke-CodexTool` returns the raw tool output object. Use it for built-in tools, dynamic tools, and MCP tools.\n- PowerShell session state persists across calls, including variables, functions, aliases, imported modules, environment changes, and `$LASTEXITCODE`. Reset the kernel with `ps_repl_reset` when needed.\n- To share generated images with the model, write a file under `$CodexTmpDir`, call `Invoke-CodexTool -Name view_image -Arguments @{ path = \"/absolute/path\" }`, then delete the file.\n- Avoid direct `[Console]::Write*`, raw StdOut/StdErr writes, or other host-level output that bypasses PowerShell streams; they can corrupt the JSON line protocol. Use pipeline output, `Write-Output`, `Write-Host`, `Write-Verbose`, or `Write-Warning` instead."; + assert_eq!(res, expected); + } + /// When both system instructions *and* a project doc are present the two /// should be concatenated with the separator. #[tokio::test] diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index 0700b4d013..29ac9dffb6 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -20,6 +20,7 @@ pub type SharedTurnDiffTracker = Arc>; pub enum ToolCallSource { Direct, JsRepl, + PsRepl, } #[derive(Clone)] diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index 89e35fa227..132abb98fb 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -8,6 +8,7 @@ mod mcp; mod mcp_resource; pub(crate) mod multi_agents; mod plan; +mod ps_repl; mod read_file; mod request_user_input; mod search_tool_bm25; @@ -38,6 +39,8 @@ pub use mcp::McpHandler; pub use mcp_resource::McpResourceHandler; pub use multi_agents::MultiAgentHandler; pub use plan::PlanHandler; +pub use ps_repl::PsReplHandler; +pub use ps_repl::PsReplResetHandler; pub use read_file::ReadFileHandler; pub use request_user_input::RequestUserInputHandler; pub(crate) use request_user_input::request_user_input_tool_description; diff --git a/codex-rs/core/src/tools/handlers/ps_repl.rs b/codex-rs/core/src/tools/handlers/ps_repl.rs new file mode 100644 index 0000000000..82fe78e51f --- /dev/null +++ b/codex-rs/core/src/tools/handlers/ps_repl.rs @@ -0,0 +1,363 @@ +use async_trait::async_trait; +use serde_json::Value as JsonValue; +use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; + +use crate::exec::ExecToolCallOutput; +use crate::exec::StreamOutput; +use crate::features::Feature; +use crate::function_tool::FunctionCallError; +use crate::protocol::ExecCommandSource; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolOutput; +use crate::tools::context::ToolPayload; +use crate::tools::events::ToolEmitter; +use crate::tools::events::ToolEventCtx; +use crate::tools::events::ToolEventFailure; +use crate::tools::events::ToolEventStage; +use crate::tools::handlers::parse_arguments; +use crate::tools::ps_repl::PS_REPL_PRAGMA_PREFIX; +use crate::tools::ps_repl::PsExecResult; +use crate::tools::ps_repl::PsReplArgs; +use crate::tools::registry::ToolHandler; +use crate::tools::registry::ToolKind; +use codex_protocol::models::FunctionCallOutputBody; +use codex_protocol::models::FunctionCallOutputContentItem; + +pub struct PsReplHandler; +pub struct PsReplResetHandler; + +fn join_outputs(stdout: &str, stderr: &str) -> String { + if stdout.is_empty() { + stderr.to_string() + } else if stderr.is_empty() { + stdout.to_string() + } else { + format!("{stdout}\n{stderr}") + } +} + +fn build_ps_repl_exec_output( + output: &str, + error: Option<&str>, + duration: Duration, +) -> ExecToolCallOutput { + let stdout = output.to_string(); + let stderr = error.unwrap_or("").to_string(); + let aggregated_output = join_outputs(&stdout, &stderr); + ExecToolCallOutput { + exit_code: if error.is_some() { 1 } else { 0 }, + stdout: StreamOutput::new(stdout), + stderr: StreamOutput::new(stderr), + aggregated_output: StreamOutput::new(aggregated_output), + duration, + timed_out: false, + } +} + +async fn emit_ps_repl_exec_begin( + session: &crate::codex::Session, + turn: &crate::codex::TurnContext, + call_id: &str, +) { + let emitter = ToolEmitter::shell( + vec!["ps_repl".to_string()], + turn.cwd.clone(), + ExecCommandSource::Agent, + false, + ); + let ctx = ToolEventCtx::new(session, turn, call_id, None); + emitter.emit(ctx, ToolEventStage::Begin).await; +} + +async fn emit_ps_repl_exec_end( + session: &crate::codex::Session, + turn: &crate::codex::TurnContext, + call_id: &str, + output: &str, + error: Option<&str>, + duration: Duration, +) { + let exec_output = build_ps_repl_exec_output(output, error, duration); + let emitter = ToolEmitter::shell( + vec!["ps_repl".to_string()], + turn.cwd.clone(), + ExecCommandSource::Agent, + false, + ); + let ctx = ToolEventCtx::new(session, turn, call_id, None); + let stage = if error.is_some() { + ToolEventStage::Failure(ToolEventFailure::Output(exec_output)) + } else { + ToolEventStage::Success(exec_output) + }; + emitter.emit(ctx, stage).await; +} + +#[async_trait] +impl ToolHandler for PsReplHandler { + fn kind(&self) -> ToolKind { + ToolKind::Function + } + + fn matches_kind(&self, payload: &ToolPayload) -> bool { + matches!( + payload, + ToolPayload::Function { .. } | ToolPayload::Custom { .. } + ) + } + + async fn handle(&self, invocation: ToolInvocation) -> Result { + let ToolInvocation { + session, + turn, + tracker, + payload, + call_id, + .. + } = invocation; + + if !session.features().enabled(Feature::PsRepl) { + return Err(FunctionCallError::RespondToModel( + "ps_repl is disabled by feature flag".to_string(), + )); + } + + let args = match payload { + ToolPayload::Function { arguments } => parse_arguments(&arguments)?, + ToolPayload::Custom { input } => parse_freeform_args(&input)?, + _ => { + return Err(FunctionCallError::RespondToModel( + "ps_repl expects custom or function payload".to_string(), + )); + } + }; + let manager = turn.ps_repl.manager().await?; + let started_at = Instant::now(); + emit_ps_repl_exec_begin(session.as_ref(), turn.as_ref(), &call_id).await; + let result = manager + .execute(Arc::clone(&session), Arc::clone(&turn), tracker, args) + .await; + let result = match result { + Ok(result) => result, + Err(err) => { + let message = err.to_string(); + emit_ps_repl_exec_end( + session.as_ref(), + turn.as_ref(), + &call_id, + "", + Some(&message), + started_at.elapsed(), + ) + .await; + return Err(err); + } + }; + + emit_ps_repl_exec_end( + session.as_ref(), + turn.as_ref(), + &call_id, + &result.output, + None, + started_at.elapsed(), + ) + .await; + + Ok(build_tool_output(result)) + } +} + +fn build_tool_output(result: PsExecResult) -> ToolOutput { + let PsExecResult { + output, + content_items, + } = result; + let mut items = Vec::with_capacity(content_items.len() + 1); + if !output.is_empty() { + items.push(FunctionCallOutputContentItem::InputText { + text: output.clone(), + }); + } + items.extend(content_items); + + ToolOutput::Function { + body: if items.is_empty() { + FunctionCallOutputBody::Text(output) + } else { + FunctionCallOutputBody::ContentItems(items) + }, + success: Some(true), + } +} + +#[async_trait] +impl ToolHandler for PsReplResetHandler { + fn kind(&self) -> ToolKind { + ToolKind::Function + } + + async fn handle(&self, invocation: ToolInvocation) -> Result { + if !invocation.session.features().enabled(Feature::PsRepl) { + return Err(FunctionCallError::RespondToModel( + "ps_repl is disabled by feature flag".to_string(), + )); + } + let manager = invocation.turn.ps_repl.manager().await?; + manager.reset().await?; + Ok(ToolOutput::Function { + body: FunctionCallOutputBody::Text("ps_repl kernel reset".to_string()), + success: Some(true), + }) + } +} + +fn parse_freeform_args(input: &str) -> Result { + if input.trim().is_empty() { + return Err(FunctionCallError::RespondToModel( + "ps_repl expects raw PowerShell tool input (non-empty). Provide PowerShell source text, optionally with first-line `# codex-ps-repl: ...`." + .to_string(), + )); + } + + let mut args = PsReplArgs { + code: input.to_string(), + timeout_ms: None, + }; + + let mut lines = input.splitn(2, '\n'); + let first_line = lines.next().unwrap_or_default(); + let rest = lines.next().unwrap_or_default(); + let trimmed = first_line.trim_start(); + let Some(pragma) = trimmed.strip_prefix(PS_REPL_PRAGMA_PREFIX) else { + reject_json_or_quoted_source(&args.code)?; + return Ok(args); + }; + + let mut timeout_ms: Option = None; + let directive = pragma.trim(); + if !directive.is_empty() { + for token in directive.split_whitespace() { + let (key, value) = token.split_once('=').ok_or_else(|| { + FunctionCallError::RespondToModel(format!( + "ps_repl pragma expects space-separated key=value pairs (supported keys: timeout_ms); got `{token}`" + )) + })?; + match key { + "timeout_ms" => { + if timeout_ms.is_some() { + return Err(FunctionCallError::RespondToModel( + "ps_repl pragma specifies timeout_ms more than once".to_string(), + )); + } + let parsed = value.parse::().map_err(|_| { + FunctionCallError::RespondToModel(format!( + "ps_repl pragma timeout_ms must be an integer; got `{value}`" + )) + })?; + timeout_ms = Some(parsed); + } + _ => { + return Err(FunctionCallError::RespondToModel(format!( + "ps_repl pragma only supports timeout_ms; got `{key}`" + ))); + } + } + } + } + + if rest.trim().is_empty() { + return Err(FunctionCallError::RespondToModel( + "ps_repl pragma must be followed by PowerShell source on subsequent lines".to_string(), + )); + } + + reject_json_or_quoted_source(rest)?; + args.code = rest.to_string(); + args.timeout_ms = timeout_ms; + Ok(args) +} + +fn reject_json_or_quoted_source(code: &str) -> Result<(), FunctionCallError> { + let trimmed = code.trim(); + if trimmed.starts_with("```") { + return Err(FunctionCallError::RespondToModel( + "ps_repl expects raw PowerShell source, not markdown code fences. Resend plain PowerShell only (optional first line `# codex-ps-repl: ...`)." + .to_string(), + )); + } + if is_quoted_source(trimmed) { + return Err(FunctionCallError::RespondToModel( + "ps_repl is a freeform tool and expects raw PowerShell source. Resend plain PowerShell only (optional first line `# codex-ps-repl: ...`); do not send quoted code or markdown fences." + .to_string(), + )); + } + let Ok(value) = serde_json::from_str::(trimmed) else { + return Ok(()); + }; + match value { + JsonValue::Object(_) | JsonValue::String(_) => Err(FunctionCallError::RespondToModel( + "ps_repl is a freeform tool and expects raw PowerShell source. Resend plain PowerShell only (optional first line `# codex-ps-repl: ...`); do not send JSON (`{\"code\":...}`), quoted code, or markdown fences." + .to_string(), + )), + _ => Ok(()), + } +} + +fn is_quoted_source(input: &str) -> bool { + input.len() >= 2 + && ((input.starts_with('"') && input.ends_with('"')) + || (input.starts_with('\'') && input.ends_with('\''))) +} + +#[cfg(test)] +mod tests { + use super::parse_freeform_args; + use pretty_assertions::assert_eq; + + #[test] + fn parse_freeform_args_without_pragma() { + let args = parse_freeform_args("Write-Output 'ok'").expect("parse args"); + assert_eq!(args.code, "Write-Output 'ok'"); + assert_eq!(args.timeout_ms, None); + } + + #[test] + fn parse_freeform_args_with_pragma() { + let input = "# codex-ps-repl: timeout_ms=15000\nWrite-Output 'ok'"; + let args = parse_freeform_args(input).expect("parse args"); + assert_eq!(args.code, "Write-Output 'ok'"); + assert_eq!(args.timeout_ms, Some(15_000)); + } + + #[test] + fn parse_freeform_args_rejects_unknown_key() { + let err = parse_freeform_args("# codex-ps-repl: nope=1\nWrite-Output 'ok'") + .expect_err("expected error"); + assert_eq!( + err.to_string(), + "ps_repl pragma only supports timeout_ms; got `nope`" + ); + } + + #[test] + fn parse_freeform_args_rejects_json_wrapped_code() { + let err = + parse_freeform_args(r#"{"code":"Write-Output 'ok'"}"#).expect_err("expected error"); + assert_eq!( + err.to_string(), + "ps_repl is a freeform tool and expects raw PowerShell source. Resend plain PowerShell only (optional first line `# codex-ps-repl: ...`); do not send JSON (`{\"code\":...}`), quoted code, or markdown fences." + ); + } + + #[test] + fn parse_freeform_args_rejects_quoted_source() { + let err = parse_freeform_args("'Write-Output ok'").expect_err("expected error"); + assert_eq!( + err.to_string(), + "ps_repl is a freeform tool and expects raw PowerShell source. Resend plain PowerShell only (optional first line `# codex-ps-repl: ...`); do not send quoted code or markdown fences." + ); + } +} diff --git a/codex-rs/core/src/tools/mod.rs b/codex-rs/core/src/tools/mod.rs index 27fe515126..d3b426a832 100644 --- a/codex-rs/core/src/tools/mod.rs +++ b/codex-rs/core/src/tools/mod.rs @@ -5,6 +5,7 @@ pub mod js_repl; pub(crate) mod network_approval; pub mod orchestrator; pub mod parallel; +pub mod ps_repl; pub mod registry; pub mod router; pub mod runtimes; diff --git a/codex-rs/core/src/tools/ps_repl/kernel.ps1 b/codex-rs/core/src/tools/ps_repl/kernel.ps1 new file mode 100644 index 0000000000..fe7bade5fb --- /dev/null +++ b/codex-rs/core/src/tools/ps_repl/kernel.ps1 @@ -0,0 +1,187 @@ +$ErrorActionPreference = 'Stop' +$ProgressPreference = 'SilentlyContinue' +[Console]::InputEncoding = [System.Text.Encoding]::UTF8 +[Console]::OutputEncoding = [System.Text.Encoding]::UTF8 + +$script:CodexTmpDir = if ($env:CODEX_PS_TMP_DIR) { + $env:CODEX_PS_TMP_DIR +} else { + (Get-Location).Path +} +$script:ToolCounter = 0 +$script:ActiveExecId = $null + +function Send-KernelMessage { + param( + [Parameter(Mandatory)] + [hashtable]$Message + ) + + $json = $Message | ConvertTo-Json -Compress -Depth 100 + [Console]::Out.WriteLine($json) + [Console]::Out.Flush() +} + +function Read-KernelMessage { + while ($true) { + $line = [Console]::In.ReadLine() + if ($null -eq $line) { + return $null + } + if ([string]::IsNullOrWhiteSpace($line)) { + continue + } + + try { + return $line | ConvertFrom-Json -AsHashtable -Depth 100 + } catch { + continue + } + } +} + +function Format-StreamItem { + param($Item) + + if ($null -eq $Item) { + return $null + } + if ($Item -is [string]) { + return $Item.TrimEnd("`r", "`n") + } + + $text = $Item | Out-String -Width 4096 + $trimmed = $text.TrimEnd("`r", "`n") + if ([string]::IsNullOrEmpty($trimmed)) { + return $null + } + $trimmed +} + +function Wait-ToolResult { + param( + [Parameter(Mandatory)] + [string]$Id + ) + + while ($true) { + $message = Read-KernelMessage + if ($null -eq $message) { + throw "ps_repl kernel closed while waiting for tool result" + } + if ($message.type -ne 'run_tool_result') { + throw "ps_repl kernel received unexpected message while waiting for tool result: $($message.type)" + } + if ($message.id -ne $Id) { + throw "ps_repl kernel received mismatched tool result: expected $Id, got $($message.id)" + } + return $message + } +} + +function Invoke-CodexTool { + [CmdletBinding()] + param( + [Parameter(Mandatory)] + [string]$Name, + [object]$Arguments + ) + + if ([string]::IsNullOrWhiteSpace($Name)) { + throw "Invoke-CodexTool expects a non-empty tool name" + } + if ($null -eq $script:ActiveExecId) { + throw "Invoke-CodexTool can only be used while a ps_repl exec is running" + } + + $toolId = "{0}-tool-{1}" -f $script:ActiveExecId, $script:ToolCounter + $script:ToolCounter += 1 + + $argumentsJson = '{}' + if ($PSBoundParameters.ContainsKey('Arguments')) { + if ($Arguments -is [string]) { + $argumentsJson = $Arguments + } else { + $argumentsJson = $Arguments | ConvertTo-Json -Compress -Depth 100 + } + } + + Send-KernelMessage @{ + type = 'run_tool' + id = $toolId + exec_id = $script:ActiveExecId + tool_name = $Name + arguments = $argumentsJson + } + + $result = Wait-ToolResult -Id $toolId + if (-not $result.ok) { + if ($null -ne $result.error -and -not [string]::IsNullOrWhiteSpace([string]$result.error)) { + throw [System.Exception]::new([string]$result.error) + } + throw [System.Exception]::new('tool failed') + } + + $result.response +} + +$script:Codex = [pscustomobject]@{ + TmpDir = $script:CodexTmpDir +} +$null = $script:Codex | Add-Member -MemberType ScriptMethod -Name Tool -Value { + param($Name, $Arguments) + + if ($PSBoundParameters.ContainsKey('Arguments')) { + Invoke-CodexTool -Name $Name -Arguments $Arguments + } else { + Invoke-CodexTool -Name $Name + } +} +Set-Variable -Name Codex -Scope Script -Value $script:Codex +Set-Variable -Name CodexTmpDir -Scope Script -Value $script:CodexTmpDir + +while ($true) { + $message = Read-KernelMessage + if ($null -eq $message) { + break + } + if ($message.type -ne 'exec') { + continue + } + + $script:ActiveExecId = [string]$message.id + + try { + $scriptBlock = [scriptblock]::Create([string]$message.code) + $items = @(. $scriptBlock *>&1) + $outputLines = foreach ($item in $items) { + $formatted = Format-StreamItem -Item $item + if ($null -ne $formatted -and $formatted -ne '') { + $formatted + } + } + $output = [string]::Join("`n", @($outputLines)) + Send-KernelMessage @{ + type = 'exec_result' + id = [string]$message.id + ok = $true + output = $output + error = $null + } + } catch { + $errorMessage = if ($_.Exception -and $_.Exception.Message) { + [string]$_.Exception.Message + } else { + [string]$_ + } + Send-KernelMessage @{ + type = 'exec_result' + id = [string]$message.id + ok = $false + output = '' + error = $errorMessage + } + } finally { + $script:ActiveExecId = $null + } +} diff --git a/codex-rs/core/src/tools/ps_repl/mod.rs b/codex-rs/core/src/tools/ps_repl/mod.rs new file mode 100644 index 0000000000..f3b2aca2bc --- /dev/null +++ b/codex-rs/core/src/tools/ps_repl/mod.rs @@ -0,0 +1,1686 @@ +use std::collections::HashMap; +use std::collections::VecDeque; +use std::fmt; +#[cfg(unix)] +use std::os::unix::process::ExitStatusExt; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Duration; + +use codex_protocol::ThreadId; +use codex_protocol::models::ContentItem; +use codex_protocol::models::FunctionCallOutputContentItem; +use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::ResponseInputItem; +use serde::Deserialize; +use serde::Serialize; +use serde_json::Value as JsonValue; +use tokio::io::AsyncBufReadExt; +use tokio::io::AsyncWriteExt; +use tokio::io::BufReader; +use tokio::process::Child; +use tokio::process::ChildStdin; +use tokio::sync::Mutex; +use tokio::sync::Notify; +use tokio::sync::OnceCell; +use tokio_util::sync::CancellationToken; +use tracing::info; +use tracing::trace; +use tracing::warn; +use uuid::Uuid; + +use crate::client_common::tools::ToolSpec; +use crate::codex::Session; +use crate::codex::TurnContext; +use crate::exec::ExecExpiration; +use crate::exec_env::create_env; +use crate::function_tool::FunctionCallError; +use crate::sandboxing::CommandSpec; +use crate::sandboxing::SandboxManager; +use crate::sandboxing::SandboxPermissions; +use crate::tools::ToolRouter; +use crate::tools::context::SharedTurnDiffTracker; +use crate::tools::sandboxing::SandboxablePreference; +use crate::truncate::TruncationPolicy; +use crate::truncate::truncate_text; + +pub(crate) const PS_REPL_PRAGMA_PREFIX: &str = "# codex-ps-repl:"; +const KERNEL_SOURCE: &str = include_str!("kernel.ps1"); +const PS_REPL_MIN_PWSH_MAJOR: u64 = 7; +const PS_REPL_STDERR_TAIL_LINE_LIMIT: usize = 20; +const PS_REPL_STDERR_TAIL_LINE_MAX_BYTES: usize = 512; +const PS_REPL_STDERR_TAIL_MAX_BYTES: usize = 4_096; +const PS_REPL_STDERR_TAIL_SEPARATOR: &str = " | "; +const PS_REPL_EXEC_ID_LOG_LIMIT: usize = 8; +const PS_REPL_MODEL_DIAG_STDERR_MAX_BYTES: usize = 1_024; +const PS_REPL_MODEL_DIAG_ERROR_MAX_BYTES: usize = 256; +const PS_REPL_TOOL_RESPONSE_TEXT_PREVIEW_MAX_BYTES: usize = 512; + +/// Per-task ps_repl handle stored on the turn context. +pub(crate) struct PsReplHandle { + pwsh_path: Option, + cell: OnceCell>, +} + +impl fmt::Debug for PsReplHandle { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("PsReplHandle").finish_non_exhaustive() + } +} + +impl PsReplHandle { + pub(crate) fn with_pwsh_path(pwsh_path: Option) -> Self { + Self { + pwsh_path, + cell: OnceCell::new(), + } + } + + pub(crate) async fn manager(&self) -> Result, FunctionCallError> { + self.cell + .get_or_try_init(|| async { PsReplManager::new(self.pwsh_path.clone()).await }) + .await + .cloned() + } +} + +#[derive(Clone, Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PsReplArgs { + pub code: String, + #[serde(default)] + pub timeout_ms: Option, +} + +#[derive(Clone, Debug)] +pub struct PsExecResult { + pub output: String, + pub content_items: Vec, +} + +struct KernelState { + child: Arc>, + recent_stderr: Arc>>, + stdin: Arc>, + pending_execs: Arc>>>, + exec_contexts: Arc>>, + shutdown: CancellationToken, +} + +#[derive(Clone)] +struct ExecContext { + session: Arc, + turn: Arc, + tracker: SharedTurnDiffTracker, +} + +#[derive(Default)] +struct ExecToolCalls { + in_flight: usize, + content_items: Vec, + notify: Arc, + cancel: CancellationToken, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[allow(clippy::enum_variant_names)] +enum PsReplToolCallPayloadKind { + MessageContent, + FunctionText, + FunctionContentItems, + CustomText, + CustomContentItems, + McpResult, + McpErrorResult, + Error, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +struct PsReplToolCallResponseSummary { + response_type: Option, + payload_kind: Option, + payload_text_preview: Option, + payload_text_length: Option, + payload_item_count: Option, + text_item_count: Option, + image_item_count: Option, + structured_content_present: Option, + result_is_error: Option, +} + +enum KernelStreamEnd { + Shutdown, + StdoutEof, + StdoutReadError(String), +} + +impl KernelStreamEnd { + fn reason(&self) -> &'static str { + match self { + Self::Shutdown => "shutdown", + Self::StdoutEof => "stdout_eof", + Self::StdoutReadError(_) => "stdout_read_error", + } + } + + fn error(&self) -> Option<&str> { + match self { + Self::StdoutReadError(err) => Some(err), + _ => None, + } + } +} + +struct KernelDebugSnapshot { + pid: Option, + status: String, + stderr_tail: String, +} + +fn format_exit_status(status: std::process::ExitStatus) -> String { + if let Some(code) = status.code() { + return format!("code={code}"); + } + #[cfg(unix)] + if let Some(signal) = status.signal() { + return format!("signal={signal}"); + } + "unknown".to_string() +} + +fn format_stderr_tail(lines: &VecDeque) -> String { + if lines.is_empty() { + return "".to_string(); + } + lines + .iter() + .cloned() + .collect::>() + .join(PS_REPL_STDERR_TAIL_SEPARATOR) +} + +fn truncate_utf8_prefix_by_bytes(input: &str, max_bytes: usize) -> String { + if input.len() <= max_bytes { + return input.to_string(); + } + if max_bytes == 0 { + return String::new(); + } + let mut end = max_bytes; + while end > 0 && !input.is_char_boundary(end) { + end -= 1; + } + input[..end].to_string() +} + +fn stderr_tail_formatted_bytes(lines: &VecDeque) -> usize { + if lines.is_empty() { + return 0; + } + let payload_bytes: usize = lines.iter().map(String::len).sum(); + let separator_bytes = PS_REPL_STDERR_TAIL_SEPARATOR.len() * (lines.len() - 1); + payload_bytes + separator_bytes +} + +fn stderr_tail_bytes_with_candidate(lines: &VecDeque, line: &str) -> usize { + if lines.is_empty() { + return line.len(); + } + stderr_tail_formatted_bytes(lines) + PS_REPL_STDERR_TAIL_SEPARATOR.len() + line.len() +} + +fn push_stderr_tail_line(lines: &mut VecDeque, line: &str) -> String { + let max_line_bytes = PS_REPL_STDERR_TAIL_LINE_MAX_BYTES.min(PS_REPL_STDERR_TAIL_MAX_BYTES); + let bounded_line = truncate_utf8_prefix_by_bytes(line, max_line_bytes); + if bounded_line.is_empty() { + return bounded_line; + } + + while !lines.is_empty() + && (lines.len() >= PS_REPL_STDERR_TAIL_LINE_LIMIT + || stderr_tail_bytes_with_candidate(lines, &bounded_line) + > PS_REPL_STDERR_TAIL_MAX_BYTES) + { + lines.pop_front(); + } + + lines.push_back(bounded_line.clone()); + bounded_line +} + +fn is_kernel_status_exited(status: &str) -> bool { + status.starts_with("exited(") +} + +fn should_include_model_diagnostics_for_write_error( + err_message: &str, + snapshot: &KernelDebugSnapshot, +) -> bool { + is_kernel_status_exited(&snapshot.status) + || err_message.to_ascii_lowercase().contains("broken pipe") +} + +fn format_model_kernel_failure_details( + reason: &str, + stream_error: Option<&str>, + snapshot: &KernelDebugSnapshot, +) -> String { + let payload = serde_json::json!({ + "reason": reason, + "stream_error": stream_error + .map(|err| truncate_utf8_prefix_by_bytes(err, PS_REPL_MODEL_DIAG_ERROR_MAX_BYTES)), + "kernel_pid": snapshot.pid, + "kernel_status": snapshot.status, + "kernel_stderr_tail": truncate_utf8_prefix_by_bytes( + &snapshot.stderr_tail, + PS_REPL_MODEL_DIAG_STDERR_MAX_BYTES, + ), + }); + let encoded = serde_json::to_string(&payload) + .unwrap_or_else(|err| format!(r#"{{"reason":"serialization_error","error":"{err}"}}"#)); + format!("ps_repl diagnostics: {encoded}") +} + +fn with_model_kernel_failure_message( + base_message: &str, + reason: &str, + stream_error: Option<&str>, + snapshot: &KernelDebugSnapshot, +) -> String { + format!( + "{base_message}\n\n{}", + format_model_kernel_failure_details(reason, stream_error, snapshot) + ) +} + +pub struct PsReplManager { + pwsh_path: Option, + tmp_dir: tempfile::TempDir, + kernel: Arc>>, + exec_lock: Arc, + exec_tool_calls: Arc>>, +} + +impl PsReplManager { + async fn new(pwsh_path: Option) -> Result, FunctionCallError> { + let tmp_dir = tempfile::tempdir().map_err(|err| { + FunctionCallError::RespondToModel(format!("failed to create ps_repl temp dir: {err}")) + })?; + + let manager = Arc::new(Self { + pwsh_path, + tmp_dir, + kernel: Arc::new(Mutex::new(None)), + exec_lock: Arc::new(tokio::sync::Semaphore::new(1)), + exec_tool_calls: Arc::new(Mutex::new(HashMap::new())), + }); + + Ok(manager) + } + + async fn register_exec_tool_calls(&self, exec_id: &str) { + self.exec_tool_calls + .lock() + .await + .insert(exec_id.to_string(), ExecToolCalls::default()); + } + + async fn clear_exec_tool_calls(&self, exec_id: &str) { + if let Some(state) = self.exec_tool_calls.lock().await.remove(exec_id) { + state.cancel.cancel(); + state.notify.notify_waiters(); + } + } + + async fn wait_for_exec_tool_calls(&self, exec_id: &str) { + loop { + let notified = { + let calls = self.exec_tool_calls.lock().await; + calls + .get(exec_id) + .filter(|state| state.in_flight > 0) + .map(|state| Arc::clone(&state.notify).notified_owned()) + }; + match notified { + Some(notified) => notified.await, + None => return, + } + } + } + + async fn begin_exec_tool_call( + exec_tool_calls: &Arc>>, + exec_id: &str, + ) -> Option { + let mut calls = exec_tool_calls.lock().await; + let state = calls.get_mut(exec_id)?; + state.in_flight += 1; + Some(state.cancel.clone()) + } + + async fn record_exec_tool_call_content_items( + exec_tool_calls: &Arc>>, + exec_id: &str, + content_items: Vec, + ) { + if content_items.is_empty() { + return; + } + + let mut calls = exec_tool_calls.lock().await; + if let Some(state) = calls.get_mut(exec_id) { + state.content_items.extend(content_items); + } + } + + async fn finish_exec_tool_call( + exec_tool_calls: &Arc>>, + exec_id: &str, + ) { + let notify = { + let mut calls = exec_tool_calls.lock().await; + let Some(state) = calls.get_mut(exec_id) else { + return; + }; + if state.in_flight == 0 { + return; + } + state.in_flight -= 1; + if state.in_flight == 0 { + Some(Arc::clone(&state.notify)) + } else { + None + } + }; + if let Some(notify) = notify { + notify.notify_waiters(); + } + } + + async fn wait_for_exec_tool_calls_map( + exec_tool_calls: &Arc>>, + exec_id: &str, + ) { + loop { + let notified = { + let calls = exec_tool_calls.lock().await; + calls + .get(exec_id) + .filter(|state| state.in_flight > 0) + .map(|state| Arc::clone(&state.notify).notified_owned()) + }; + match notified { + Some(notified) => notified.await, + None => return, + } + } + } + + async fn clear_exec_tool_calls_map( + exec_tool_calls: &Arc>>, + exec_id: &str, + ) { + if let Some(state) = exec_tool_calls.lock().await.remove(exec_id) { + state.cancel.cancel(); + state.notify.notify_waiters(); + } + } + + async fn clear_all_exec_tool_calls_map( + exec_tool_calls: &Arc>>, + ) { + let states = { + let mut calls = exec_tool_calls.lock().await; + calls.drain().map(|(_, state)| state).collect::>() + }; + for state in states { + state.cancel.cancel(); + state.notify.notify_waiters(); + } + } + + fn log_tool_call_response( + req: &RunToolRequest, + ok: bool, + summary: &PsReplToolCallResponseSummary, + response: Option<&JsonValue>, + error: Option<&str>, + ) { + info!( + exec_id = %req.exec_id, + tool_call_id = %req.id, + tool_name = %req.tool_name, + ok, + summary = ?summary, + "ps_repl nested tool call completed" + ); + if let Some(response) = response { + trace!( + exec_id = %req.exec_id, + tool_call_id = %req.id, + tool_name = %req.tool_name, + response_json = %response, + "ps_repl nested tool call raw response" + ); + } + if let Some(error) = error { + trace!( + exec_id = %req.exec_id, + tool_call_id = %req.id, + tool_name = %req.tool_name, + error = %error, + "ps_repl nested tool call raw error" + ); + } + } + + fn summarize_text_payload( + response_type: Option<&str>, + payload_kind: PsReplToolCallPayloadKind, + text: &str, + ) -> PsReplToolCallResponseSummary { + PsReplToolCallResponseSummary { + response_type: response_type.map(str::to_owned), + payload_kind: Some(payload_kind), + payload_text_preview: (!text.is_empty()).then(|| { + truncate_text( + text, + TruncationPolicy::Bytes(PS_REPL_TOOL_RESPONSE_TEXT_PREVIEW_MAX_BYTES), + ) + }), + payload_text_length: Some(text.len()), + ..Default::default() + } + } + + fn summarize_function_output_payload( + response_type: &str, + payload_kind: PsReplToolCallPayloadKind, + output: &FunctionCallOutputPayload, + ) -> PsReplToolCallResponseSummary { + let (payload_item_count, text_item_count, image_item_count) = + if let Some(items) = output.content_items() { + let text_item_count = items + .iter() + .filter(|item| matches!(item, FunctionCallOutputContentItem::InputText { .. })) + .count(); + let image_item_count = items.len().saturating_sub(text_item_count); + ( + Some(items.len()), + Some(text_item_count), + Some(image_item_count), + ) + } else { + (None, None, None) + }; + let payload_text = output.body.to_text(); + PsReplToolCallResponseSummary { + response_type: Some(response_type.to_string()), + payload_kind: Some(payload_kind), + payload_text_preview: payload_text.as_deref().and_then(|text| { + (!text.is_empty()).then(|| { + truncate_text( + text, + TruncationPolicy::Bytes(PS_REPL_TOOL_RESPONSE_TEXT_PREVIEW_MAX_BYTES), + ) + }) + }), + payload_text_length: payload_text.as_ref().map(String::len), + payload_item_count, + text_item_count, + image_item_count, + ..Default::default() + } + } + + fn summarize_message_payload(content: &[ContentItem]) -> PsReplToolCallResponseSummary { + let text_item_count = content + .iter() + .filter(|item| { + matches!( + item, + ContentItem::InputText { .. } | ContentItem::OutputText { .. } + ) + }) + .count(); + let image_item_count = content.len().saturating_sub(text_item_count); + let payload_text = content + .iter() + .filter_map(|item| match item { + ContentItem::InputText { text } | ContentItem::OutputText { text } + if !text.trim().is_empty() => + { + Some(text.as_str()) + } + ContentItem::InputText { .. } + | ContentItem::InputImage { .. } + | ContentItem::OutputText { .. } => None, + }) + .collect::>(); + let payload_text = if payload_text.is_empty() { + None + } else { + Some(payload_text.join("\n")) + }; + PsReplToolCallResponseSummary { + response_type: Some("message".to_string()), + payload_kind: Some(PsReplToolCallPayloadKind::MessageContent), + payload_text_preview: payload_text.as_deref().and_then(|text| { + (!text.is_empty()).then(|| { + truncate_text( + text, + TruncationPolicy::Bytes(PS_REPL_TOOL_RESPONSE_TEXT_PREVIEW_MAX_BYTES), + ) + }) + }), + payload_text_length: payload_text.as_ref().map(String::len), + payload_item_count: Some(content.len()), + text_item_count: Some(text_item_count), + image_item_count: Some(image_item_count), + ..Default::default() + } + } + + fn summarize_tool_call_response(response: &ResponseInputItem) -> PsReplToolCallResponseSummary { + match response { + ResponseInputItem::Message { content, .. } => Self::summarize_message_payload(content), + ResponseInputItem::FunctionCallOutput { output, .. } => { + let payload_kind = if output.content_items().is_some() { + PsReplToolCallPayloadKind::FunctionContentItems + } else { + PsReplToolCallPayloadKind::FunctionText + }; + Self::summarize_function_output_payload( + "function_call_output", + payload_kind, + output, + ) + } + ResponseInputItem::CustomToolCallOutput { output, .. } => { + let payload_kind = if output.content_items().is_some() { + PsReplToolCallPayloadKind::CustomContentItems + } else { + PsReplToolCallPayloadKind::CustomText + }; + Self::summarize_function_output_payload( + "custom_tool_call_output", + payload_kind, + output, + ) + } + ResponseInputItem::McpToolCallOutput { result, .. } => match result { + Ok(result) => { + let output = FunctionCallOutputPayload::from(result); + let mut summary = Self::summarize_function_output_payload( + "mcp_tool_call_output", + PsReplToolCallPayloadKind::McpResult, + &output, + ); + summary.payload_item_count = Some(result.content.len()); + summary.structured_content_present = Some(result.structured_content.is_some()); + summary.result_is_error = Some(result.is_error.unwrap_or(false)); + summary + } + Err(error) => { + let mut summary = Self::summarize_text_payload( + Some("mcp_tool_call_output"), + PsReplToolCallPayloadKind::McpErrorResult, + error, + ); + summary.result_is_error = Some(true); + summary + } + }, + } + } + + fn summarize_tool_call_error(error: &str) -> PsReplToolCallResponseSummary { + Self::summarize_text_payload(None, PsReplToolCallPayloadKind::Error, error) + } + + pub async fn reset(&self) -> Result<(), FunctionCallError> { + let _permit = self.exec_lock.clone().acquire_owned().await.map_err(|_| { + FunctionCallError::RespondToModel("ps_repl execution unavailable".to_string()) + })?; + self.reset_kernel().await; + Self::clear_all_exec_tool_calls_map(&self.exec_tool_calls).await; + Ok(()) + } + + async fn reset_kernel(&self) { + let state = { + let mut guard = self.kernel.lock().await; + guard.take() + }; + if let Some(state) = state { + state.shutdown.cancel(); + Self::kill_kernel_child(&state.child, "reset").await; + } + } + + pub async fn execute( + &self, + session: Arc, + turn: Arc, + tracker: SharedTurnDiffTracker, + args: PsReplArgs, + ) -> Result { + let _permit = self.exec_lock.clone().acquire_owned().await.map_err(|_| { + FunctionCallError::RespondToModel("ps_repl execution unavailable".to_string()) + })?; + + let (stdin, pending_execs, exec_contexts, child, recent_stderr) = { + let mut kernel = self.kernel.lock().await; + if kernel.is_none() { + let state = self + .start_kernel(Arc::clone(&turn), Some(session.conversation_id)) + .await + .map_err(FunctionCallError::RespondToModel)?; + *kernel = Some(state); + } + + let state = match kernel.as_ref() { + Some(state) => state, + None => { + return Err(FunctionCallError::RespondToModel( + "ps_repl kernel unavailable".to_string(), + )); + } + }; + ( + Arc::clone(&state.stdin), + Arc::clone(&state.pending_execs), + Arc::clone(&state.exec_contexts), + Arc::clone(&state.child), + Arc::clone(&state.recent_stderr), + ) + }; + + let (req_id, rx) = { + let req_id = Uuid::new_v4().to_string(); + let mut pending = pending_execs.lock().await; + let (tx, rx) = tokio::sync::oneshot::channel(); + pending.insert(req_id.clone(), tx); + exec_contexts.lock().await.insert( + req_id.clone(), + ExecContext { + session: Arc::clone(&session), + turn: Arc::clone(&turn), + tracker, + }, + ); + (req_id, rx) + }; + self.register_exec_tool_calls(&req_id).await; + + let payload = HostToKernel::Exec { + id: req_id.clone(), + code: args.code, + timeout_ms: args.timeout_ms, + }; + + if let Err(err) = Self::write_message(&stdin, &payload).await { + pending_execs.lock().await.remove(&req_id); + exec_contexts.lock().await.remove(&req_id); + self.clear_exec_tool_calls(&req_id).await; + let snapshot = Self::kernel_debug_snapshot(&child, &recent_stderr).await; + let err_message = err.to_string(); + warn!( + exec_id = %req_id, + error = %err_message, + kernel_pid = ?snapshot.pid, + kernel_status = %snapshot.status, + kernel_stderr_tail = %snapshot.stderr_tail, + "failed to submit ps_repl exec request to kernel" + ); + let message = + if should_include_model_diagnostics_for_write_error(&err_message, &snapshot) { + with_model_kernel_failure_message( + &err_message, + "write_failed", + Some(&err_message), + &snapshot, + ) + } else { + err_message + }; + return Err(FunctionCallError::RespondToModel(message)); + } + + let timeout_ms = args.timeout_ms.unwrap_or(30_000); + let response = match tokio::time::timeout(Duration::from_millis(timeout_ms), rx).await { + Ok(Ok(msg)) => msg, + Ok(Err(_)) => { + let mut pending = pending_execs.lock().await; + pending.remove(&req_id); + exec_contexts.lock().await.remove(&req_id); + self.wait_for_exec_tool_calls(&req_id).await; + self.clear_exec_tool_calls(&req_id).await; + let snapshot = Self::kernel_debug_snapshot(&child, &recent_stderr).await; + let message = if is_kernel_status_exited(&snapshot.status) { + with_model_kernel_failure_message( + "ps_repl kernel closed unexpectedly", + "response_channel_closed", + None, + &snapshot, + ) + } else { + "ps_repl kernel closed unexpectedly".to_string() + }; + return Err(FunctionCallError::RespondToModel(message)); + } + Err(_) => { + self.reset_kernel().await; + self.wait_for_exec_tool_calls(&req_id).await; + self.exec_tool_calls.lock().await.clear(); + return Err(FunctionCallError::RespondToModel( + "ps_repl execution timed out; kernel reset, rerun your request".to_string(), + )); + } + }; + + match response { + ExecResultMessage::Ok { content_items } => { + let (output, content_items) = split_exec_result_content_items(content_items); + Ok(PsExecResult { + output, + content_items, + }) + } + ExecResultMessage::Err { message } => Err(FunctionCallError::RespondToModel(message)), + } + } + + async fn start_kernel( + &self, + turn: Arc, + thread_id: Option, + ) -> Result { + let pwsh_path = resolve_compatible_pwsh(self.pwsh_path.as_deref()).await?; + + let kernel_path = self + .write_kernel_script() + .await + .map_err(|err| err.to_string())?; + + let mut env = create_env(&turn.shell_environment_policy, thread_id); + env.insert( + "CODEX_PS_TMP_DIR".to_string(), + self.tmp_dir.path().to_string_lossy().to_string(), + ); + + let spec = CommandSpec { + program: pwsh_path.to_string_lossy().to_string(), + args: vec![ + "-NoLogo".to_string(), + "-NoProfile".to_string(), + "-File".to_string(), + kernel_path.to_string_lossy().to_string(), + ], + cwd: turn.cwd.clone(), + env, + expiration: ExecExpiration::DefaultTimeout, + sandbox_permissions: SandboxPermissions::UseDefault, + additional_permissions: None, + justification: None, + }; + + let sandbox = SandboxManager::new(); + let has_managed_network_requirements = turn + .config + .config_layer_stack + .requirements_toml() + .network + .is_some(); + let sandbox_type = sandbox.select_initial( + &turn.sandbox_policy, + SandboxablePreference::Auto, + turn.windows_sandbox_level, + has_managed_network_requirements, + ); + let exec_env = sandbox + .transform(crate::sandboxing::SandboxTransformRequest { + spec, + policy: &turn.sandbox_policy, + sandbox: sandbox_type, + enforce_managed_network: has_managed_network_requirements, + network: None, + sandbox_policy_cwd: &turn.cwd, + #[cfg(target_os = "macos")] + macos_seatbelt_profile_extensions: None, + codex_linux_sandbox_exe: turn.codex_linux_sandbox_exe.as_ref(), + use_linux_sandbox_bwrap: turn + .features + .enabled(crate::features::Feature::UseLinuxSandboxBwrap), + windows_sandbox_level: turn.windows_sandbox_level, + }) + .map_err(|err| format!("failed to configure sandbox for ps_repl: {err}"))?; + + let mut cmd = + tokio::process::Command::new(exec_env.command.first().cloned().unwrap_or_default()); + if exec_env.command.len() > 1 { + cmd.args(&exec_env.command[1..]); + } + #[cfg(unix)] + cmd.arg0( + exec_env + .arg0 + .clone() + .unwrap_or_else(|| exec_env.command.first().cloned().unwrap_or_default()), + ); + cmd.current_dir(&exec_env.cwd); + cmd.env_clear(); + cmd.envs(exec_env.env); + cmd.stdin(std::process::Stdio::piped()) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .kill_on_drop(true); + + let mut child = cmd + .spawn() + .map_err(|err| format!("failed to start PowerShell runtime: {err}"))?; + let stdout = child + .stdout + .take() + .ok_or_else(|| "ps_repl kernel missing stdout".to_string())?; + let stderr = child.stderr.take(); + let stdin = child + .stdin + .take() + .ok_or_else(|| "ps_repl kernel missing stdin".to_string())?; + + let shutdown = CancellationToken::new(); + let pending_execs: Arc< + Mutex>>, + > = Arc::new(Mutex::new(HashMap::new())); + let exec_contexts: Arc>> = + Arc::new(Mutex::new(HashMap::new())); + let stdin_arc = Arc::new(Mutex::new(stdin)); + let child = Arc::new(Mutex::new(child)); + let recent_stderr = Arc::new(Mutex::new(VecDeque::with_capacity( + PS_REPL_STDERR_TAIL_LINE_LIMIT, + ))); + + tokio::spawn(Self::read_stdout( + stdout, + Arc::clone(&child), + Arc::clone(&self.kernel), + Arc::clone(&recent_stderr), + Arc::clone(&pending_execs), + Arc::clone(&exec_contexts), + Arc::clone(&self.exec_tool_calls), + Arc::clone(&stdin_arc), + shutdown.clone(), + )); + if let Some(stderr) = stderr { + tokio::spawn(Self::read_stderr( + stderr, + Arc::clone(&recent_stderr), + shutdown.clone(), + )); + } else { + warn!("ps_repl kernel missing stderr"); + } + + Ok(KernelState { + child, + recent_stderr, + stdin: stdin_arc, + pending_execs, + exec_contexts, + shutdown, + }) + } + + async fn write_kernel_script(&self) -> Result { + let dir = self.tmp_dir.path(); + let kernel_path = dir.join("ps_repl_kernel.ps1"); + tokio::fs::write(&kernel_path, KERNEL_SOURCE).await?; + Ok(kernel_path) + } + + async fn write_message( + stdin: &Arc>, + msg: &HostToKernel, + ) -> Result<(), FunctionCallError> { + let encoded = serde_json::to_string(msg).map_err(|err| { + FunctionCallError::RespondToModel(format!("failed to serialize kernel message: {err}")) + })?; + let mut guard = stdin.lock().await; + guard.write_all(encoded.as_bytes()).await.map_err(|err| { + FunctionCallError::RespondToModel(format!("failed to write to kernel: {err}")) + })?; + guard.write_all(b"\n").await.map_err(|err| { + FunctionCallError::RespondToModel(format!("failed to flush kernel message: {err}")) + })?; + Ok(()) + } + + async fn kernel_stderr_tail_snapshot(recent_stderr: &Arc>>) -> String { + let tail = recent_stderr.lock().await; + format_stderr_tail(&tail) + } + + async fn kernel_debug_snapshot( + child: &Arc>, + recent_stderr: &Arc>>, + ) -> KernelDebugSnapshot { + let (pid, status) = { + let mut guard = child.lock().await; + let pid = guard.id(); + let status = match guard.try_wait() { + Ok(Some(status)) => format!("exited({})", format_exit_status(status)), + Ok(None) => "running".to_string(), + Err(err) => format!("unknown ({err})"), + }; + (pid, status) + }; + let stderr_tail = { + let tail = recent_stderr.lock().await; + format_stderr_tail(&tail) + }; + KernelDebugSnapshot { + pid, + status, + stderr_tail, + } + } + + async fn kill_kernel_child(child: &Arc>, reason: &'static str) { + let mut guard = child.lock().await; + let pid = guard.id(); + match guard.try_wait() { + Ok(Some(_)) => return, + Ok(None) => {} + Err(err) => { + warn!( + kernel_pid = ?pid, + kill_reason = reason, + error = %err, + "failed to inspect ps_repl kernel before kill" + ); + } + } + + if let Err(err) = guard.start_kill() { + warn!( + kernel_pid = ?pid, + kill_reason = reason, + error = %err, + "failed to send kill signal to ps_repl kernel" + ); + return; + } + + match tokio::time::timeout(Duration::from_secs(2), guard.wait()).await { + Ok(Ok(_status)) => {} + Ok(Err(err)) => { + warn!( + kernel_pid = ?pid, + kill_reason = reason, + error = %err, + "failed while waiting for ps_repl kernel exit" + ); + } + Err(_) => { + warn!( + kernel_pid = ?pid, + kill_reason = reason, + "timed out waiting for ps_repl kernel to exit after kill" + ); + } + } + } + + fn truncate_id_list(ids: &[String]) -> Vec { + if ids.len() <= PS_REPL_EXEC_ID_LOG_LIMIT { + return ids.to_vec(); + } + let mut output = ids[..PS_REPL_EXEC_ID_LOG_LIMIT].to_vec(); + output.push(format!("...+{}", ids.len() - PS_REPL_EXEC_ID_LOG_LIMIT)); + output + } + + #[allow(clippy::too_many_arguments)] + async fn read_stdout( + stdout: tokio::process::ChildStdout, + child: Arc>, + manager_kernel: Arc>>, + recent_stderr: Arc>>, + pending_execs: Arc>>>, + exec_contexts: Arc>>, + exec_tool_calls: Arc>>, + stdin: Arc>, + shutdown: CancellationToken, + ) { + let mut reader = BufReader::new(stdout).lines(); + let end_reason = loop { + let line = tokio::select! { + _ = shutdown.cancelled() => break KernelStreamEnd::Shutdown, + res = reader.next_line() => match res { + Ok(Some(line)) => line, + Ok(None) => break KernelStreamEnd::StdoutEof, + Err(err) => break KernelStreamEnd::StdoutReadError(err.to_string()), + }, + }; + + let parsed: Result = serde_json::from_str(&line); + let msg = match parsed { + Ok(m) => m, + Err(err) => { + warn!("ps_repl kernel sent invalid json: {err} (line: {line})"); + continue; + } + }; + + match msg { + KernelToHost::ExecResult { + id, + ok, + output, + error, + } => { + PsReplManager::wait_for_exec_tool_calls_map(&exec_tool_calls, &id).await; + let content_items = { + let calls = exec_tool_calls.lock().await; + calls + .get(&id) + .map(|state| state.content_items.clone()) + .unwrap_or_default() + }; + let mut pending = pending_execs.lock().await; + if let Some(tx) = pending.remove(&id) { + let payload = if ok { + ExecResultMessage::Ok { + content_items: build_exec_result_content_items( + output, + content_items, + ), + } + } else { + ExecResultMessage::Err { + message: error + .unwrap_or_else(|| "ps_repl execution failed".to_string()), + } + }; + let _ = tx.send(payload); + } + exec_contexts.lock().await.remove(&id); + PsReplManager::clear_exec_tool_calls_map(&exec_tool_calls, &id).await; + } + KernelToHost::RunTool(req) => { + let Some(reset_cancel) = + PsReplManager::begin_exec_tool_call(&exec_tool_calls, &req.exec_id).await + else { + let exec_id = req.exec_id.clone(); + let tool_call_id = req.id.clone(); + let payload = HostToKernel::RunToolResult(RunToolResult { + id: req.id, + ok: false, + response: None, + error: Some("ps_repl exec context not found".to_string()), + }); + if let Err(err) = PsReplManager::write_message(&stdin, &payload).await { + let snapshot = + PsReplManager::kernel_debug_snapshot(&child, &recent_stderr).await; + warn!( + exec_id = %exec_id, + tool_call_id = %tool_call_id, + error = %err, + kernel_pid = ?snapshot.pid, + kernel_status = %snapshot.status, + kernel_stderr_tail = %snapshot.stderr_tail, + "failed to reply to kernel run_tool request" + ); + } + continue; + }; + let stdin_clone = Arc::clone(&stdin); + let exec_contexts = Arc::clone(&exec_contexts); + let exec_tool_calls_for_task = Arc::clone(&exec_tool_calls); + let recent_stderr = Arc::clone(&recent_stderr); + tokio::spawn(async move { + let exec_id = req.exec_id.clone(); + let tool_call_id = req.id.clone(); + let tool_name = req.tool_name.clone(); + let context = { exec_contexts.lock().await.get(&exec_id).cloned() }; + let result = match context { + Some(ctx) => { + tokio::select! { + _ = reset_cancel.cancelled() => RunToolResult { + id: tool_call_id.clone(), + ok: false, + response: None, + error: Some("ps_repl execution reset".to_string()), + }, + result = PsReplManager::run_tool_request( + ctx, + req, + Arc::clone(&exec_tool_calls_for_task), + ) => result, + } + } + None => RunToolResult { + id: tool_call_id.clone(), + ok: false, + response: None, + error: Some("ps_repl exec context not found".to_string()), + }, + }; + PsReplManager::finish_exec_tool_call(&exec_tool_calls_for_task, &exec_id) + .await; + let payload = HostToKernel::RunToolResult(result); + if let Err(err) = PsReplManager::write_message(&stdin_clone, &payload).await + { + let stderr_tail = + PsReplManager::kernel_stderr_tail_snapshot(&recent_stderr).await; + warn!( + exec_id = %exec_id, + tool_call_id = %tool_call_id, + tool_name = %tool_name, + error = %err, + kernel_stderr_tail = %stderr_tail, + "failed to reply to kernel run_tool request" + ); + } + }); + } + } + }; + + let exec_ids = { + let mut contexts = exec_contexts.lock().await; + let ids = contexts.keys().cloned().collect::>(); + contexts.clear(); + ids + }; + for exec_id in exec_ids { + PsReplManager::wait_for_exec_tool_calls_map(&exec_tool_calls, &exec_id).await; + PsReplManager::clear_exec_tool_calls_map(&exec_tool_calls, &exec_id).await; + } + let unexpected_snapshot = if matches!(end_reason, KernelStreamEnd::Shutdown) { + None + } else { + Some(Self::kernel_debug_snapshot(&child, &recent_stderr).await) + }; + let kernel_failure_message = unexpected_snapshot.as_ref().map(|snapshot| { + with_model_kernel_failure_message( + "ps_repl kernel exited unexpectedly", + end_reason.reason(), + end_reason.error(), + snapshot, + ) + }); + let kernel_exit_message = kernel_failure_message + .clone() + .unwrap_or_else(|| "ps_repl kernel exited unexpectedly".to_string()); + + { + let mut kernel = manager_kernel.lock().await; + let should_clear = kernel + .as_ref() + .is_some_and(|state| Arc::ptr_eq(&state.child, &child)); + if should_clear { + kernel.take(); + } + } + + let mut pending = pending_execs.lock().await; + let pending_exec_ids = pending.keys().cloned().collect::>(); + for (_id, tx) in pending.drain() { + let _ = tx.send(ExecResultMessage::Err { + message: kernel_exit_message.clone(), + }); + } + drop(pending); + + if !matches!(end_reason, KernelStreamEnd::Shutdown) { + let mut pending_exec_ids = pending_exec_ids; + pending_exec_ids.sort_unstable(); + let snapshot = Self::kernel_debug_snapshot(&child, &recent_stderr).await; + warn!( + reason = %end_reason.reason(), + stream_error = %end_reason.error().unwrap_or(""), + kernel_pid = ?snapshot.pid, + kernel_status = %snapshot.status, + pending_exec_count = pending_exec_ids.len(), + pending_exec_ids = ?Self::truncate_id_list(&pending_exec_ids), + kernel_stderr_tail = %snapshot.stderr_tail, + "ps_repl kernel terminated unexpectedly" + ); + } + } + + async fn run_tool_request( + exec: ExecContext, + req: RunToolRequest, + exec_tool_calls: Arc>>, + ) -> RunToolResult { + if is_ps_repl_internal_tool(&req.tool_name) { + let error = "ps_repl cannot invoke itself".to_string(); + let summary = Self::summarize_tool_call_error(&error); + Self::log_tool_call_response(&req, false, &summary, None, Some(&error)); + return RunToolResult { + id: req.id, + ok: false, + response: None, + error: Some(error), + }; + } + + let mcp_tools = exec + .session + .services + .mcp_connection_manager + .read() + .await + .list_all_tools() + .await; + + let router = ToolRouter::from_config( + &exec.turn.tools_config, + Some( + mcp_tools + .into_iter() + .map(|(name, tool)| (name, tool.tool)) + .collect(), + ), + None, + exec.turn.dynamic_tools.as_slice(), + ); + + let payload = + if let Some((server, tool)) = exec.session.parse_mcp_tool_name(&req.tool_name).await { + crate::tools::context::ToolPayload::Mcp { + server, + tool, + raw_arguments: req.arguments.clone(), + } + } else if is_freeform_tool(&router.specs(), &req.tool_name) { + crate::tools::context::ToolPayload::Custom { + input: req.arguments.clone(), + } + } else { + crate::tools::context::ToolPayload::Function { + arguments: req.arguments.clone(), + } + }; + + let tool_name = req.tool_name.clone(); + let call = crate::tools::router::ToolCall { + tool_name: tool_name.clone(), + call_id: req.id.clone(), + payload, + }; + + let session = Arc::clone(&exec.session); + let turn = Arc::clone(&exec.turn); + let tracker = Arc::clone(&exec.tracker); + + match router + .dispatch_tool_call( + session.clone(), + turn, + tracker, + call, + crate::tools::router::ToolCallSource::PsRepl, + ) + .await + { + Ok(response) => { + if let Some(items) = response_content_items(&response) { + Self::record_exec_tool_call_content_items( + &exec_tool_calls, + &req.exec_id, + items, + ) + .await; + } + + let summary = Self::summarize_tool_call_response(&response); + match serde_json::to_value(response) { + Ok(value) => { + Self::log_tool_call_response(&req, true, &summary, Some(&value), None); + RunToolResult { + id: req.id, + ok: true, + response: Some(value), + error: None, + } + } + Err(err) => { + let error = format!("failed to serialize tool output: {err}"); + let summary = Self::summarize_tool_call_error(&error); + Self::log_tool_call_response(&req, false, &summary, None, Some(&error)); + RunToolResult { + id: req.id, + ok: false, + response: None, + error: Some(error), + } + } + } + } + Err(err) => { + let error = err.to_string(); + let summary = Self::summarize_tool_call_error(&error); + Self::log_tool_call_response(&req, false, &summary, None, Some(&error)); + RunToolResult { + id: req.id, + ok: false, + response: None, + error: Some(error), + } + } + } + } + + async fn read_stderr( + stderr: tokio::process::ChildStderr, + recent_stderr: Arc>>, + shutdown: CancellationToken, + ) { + let mut reader = BufReader::new(stderr).lines(); + + loop { + let line = tokio::select! { + _ = shutdown.cancelled() => break, + res = reader.next_line() => match res { + Ok(Some(line)) => line, + Ok(None) => break, + Err(err) => { + warn!("ps_repl kernel stderr ended: {err}"); + break; + } + }, + }; + let trimmed = line.trim(); + if !trimmed.is_empty() { + let bounded_line = { + let mut tail = recent_stderr.lock().await; + push_stderr_tail_line(&mut tail, trimmed) + }; + if bounded_line.is_empty() { + continue; + } + warn!("ps_repl stderr: {bounded_line}"); + } + } + } +} + +fn response_content_items( + response: &ResponseInputItem, +) -> Option> { + match response { + ResponseInputItem::FunctionCallOutput { output, .. } + | ResponseInputItem::CustomToolCallOutput { output, .. } => output + .content_items() + .map(<[FunctionCallOutputContentItem]>::to_vec), + ResponseInputItem::McpToolCallOutput { result, .. } => match result { + Ok(result) => FunctionCallOutputPayload::from(result) + .content_items() + .map(<[FunctionCallOutputContentItem]>::to_vec), + Err(_) => None, + }, + ResponseInputItem::Message { .. } => None, + } +} + +fn build_exec_result_content_items( + output: String, + content_items: Vec, +) -> Vec { + let mut all_content_items = Vec::with_capacity(content_items.len() + 1); + all_content_items.push(FunctionCallOutputContentItem::InputText { text: output }); + all_content_items.extend(content_items); + all_content_items +} + +fn split_exec_result_content_items( + mut content_items: Vec, +) -> (String, Vec) { + match content_items.first() { + Some(FunctionCallOutputContentItem::InputText { .. }) => { + let FunctionCallOutputContentItem::InputText { text } = content_items.remove(0) else { + unreachable!("first content item should be input_text"); + }; + (text, content_items) + } + Some(FunctionCallOutputContentItem::InputImage { .. }) | None => { + (String::new(), content_items) + } + } +} + +fn is_freeform_tool(specs: &[ToolSpec], name: &str) -> bool { + specs + .iter() + .any(|spec| spec.name() == name && matches!(spec, ToolSpec::Freeform(_))) +} + +fn is_ps_repl_internal_tool(name: &str) -> bool { + matches!(name, "ps_repl" | "ps_repl_reset") +} + +#[derive(Clone, Debug, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +enum KernelToHost { + ExecResult { + id: String, + ok: bool, + output: String, + #[serde(default)] + error: Option, + }, + RunTool(RunToolRequest), +} + +#[derive(Clone, Debug, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +enum HostToKernel { + Exec { + id: String, + code: String, + #[serde(default)] + timeout_ms: Option, + }, + RunToolResult(RunToolResult), +} + +#[derive(Clone, Debug, Deserialize)] +struct RunToolRequest { + id: String, + exec_id: String, + tool_name: String, + arguments: String, +} + +#[derive(Clone, Debug, Serialize)] +struct RunToolResult { + id: String, + ok: bool, + #[serde(default)] + response: Option, + #[serde(default)] + error: Option, +} + +#[derive(Debug)] +enum ExecResultMessage { + Ok { + content_items: Vec, + }, + Err { + message: String, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] +struct PwshVersion { + major: u64, + minor: u64, + patch: u64, +} + +impl fmt::Display for PwshVersion { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}.{}.{}", self.major, self.minor, self.patch) + } +} + +impl PwshVersion { + fn parse(input: &str) -> Result { + let trimmed = input + .trim() + .trim_start_matches("PowerShell") + .trim() + .trim_start_matches('v'); + let mut parts = trimmed.split(['.', '-', '+']); + let major = parts + .next() + .ok_or_else(|| "missing major version".to_string())? + .parse::() + .map_err(|err| format!("invalid major version: {err}"))?; + let minor = parts + .next() + .ok_or_else(|| "missing minor version".to_string())? + .parse::() + .map_err(|err| format!("invalid minor version: {err}"))?; + let patch = parts + .next() + .ok_or_else(|| "missing patch version".to_string())? + .parse::() + .map_err(|err| format!("invalid patch version: {err}"))?; + Ok(Self { + major, + minor, + patch, + }) + } +} + +async fn read_pwsh_version(pwsh_path: &Path) -> Result { + let output = tokio::process::Command::new(pwsh_path) + .arg("--version") + .output() + .await + .map_err(|err| format!("failed to execute pwsh: {err}"))?; + + if !output.status.success() { + let mut details = String::new(); + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + let stdout = stdout.trim(); + let stderr = stderr.trim(); + if !stdout.is_empty() { + details.push_str(" stdout: "); + details.push_str(stdout); + } + if !stderr.is_empty() { + details.push_str(" stderr: "); + details.push_str(stderr); + } + let details = if details.is_empty() { + String::new() + } else { + format!(" ({details})") + }; + return Err(format!( + "failed to read pwsh version (status {status}){details}", + status = output.status + )); + } + + let stdout = String::from_utf8_lossy(&output.stdout); + let stdout = stdout.trim(); + PwshVersion::parse(stdout) + .map_err(|err| format!("failed to parse pwsh version output `{stdout}`: {err}")) +} + +async fn ensure_pwsh_version(pwsh_path: &Path) -> Result<(), String> { + let found = read_pwsh_version(pwsh_path).await?; + if found.major < PS_REPL_MIN_PWSH_MAJOR { + return Err(format!( + "PowerShell runtime too old for ps_repl (resolved {pwsh_path}): found v{found}, requires >= v{required}. Install/update PowerShell or set ps_repl_path to a newer runtime.", + pwsh_path = pwsh_path.display(), + required = PS_REPL_MIN_PWSH_MAJOR, + )); + } + Ok(()) +} + +pub(crate) async fn resolve_compatible_pwsh(config_path: Option<&Path>) -> Result { + let pwsh_path = resolve_pwsh(config_path).ok_or_else(|| { + "PowerShell 7 runtime not found; install pwsh or set CODEX_PS_REPL_PATH".to_string() + })?; + ensure_pwsh_version(&pwsh_path).await?; + Ok(pwsh_path) +} + +pub(crate) fn resolve_pwsh(config_path: Option<&Path>) -> Option { + if let Some(path) = std::env::var_os("CODEX_PS_REPL_PATH") { + let p = PathBuf::from(path); + if p.exists() { + return Some(p); + } + } + + if let Some(path) = config_path + && path.exists() + { + return Some(path.to_path_buf()); + } + + #[cfg(windows)] + if let Some(path) = crate::powershell::try_find_pwsh_executable_blocking() { + return Some(path.as_path().to_path_buf()); + } + + if let Ok(path) = which::which("pwsh") { + return Some(path); + } + + None +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + + #[test] + fn pwsh_version_parses_leading_label() { + let version = PwshVersion::parse("PowerShell 7.5.4").unwrap(); + assert_eq!( + version, + PwshVersion { + major: 7, + minor: 5, + patch: 4, + } + ); + } + + #[test] + fn pwsh_version_parses_v_prefix_and_suffix() { + let version = PwshVersion::parse("v7.5.4-preview.1").unwrap(); + assert_eq!( + version, + PwshVersion { + major: 7, + minor: 5, + patch: 4, + } + ); + } + + #[test] + fn ps_repl_internal_tool_guard_matches_expected_names() { + assert!(is_ps_repl_internal_tool("ps_repl")); + assert!(is_ps_repl_internal_tool("ps_repl_reset")); + assert!(!is_ps_repl_internal_tool("shell_command")); + assert!(!is_ps_repl_internal_tool("list_mcp_resources")); + } +} diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 76730d0890..a4ba2afde9 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -52,6 +52,7 @@ pub(crate) struct ToolsConfig { pub agent_roles: BTreeMap, pub search_tool: bool, pub request_permission_enabled: bool, + pub ps_repl_enabled: bool, pub js_repl_enabled: bool, pub js_repl_tools_only: bool, pub collab_tools: bool, @@ -77,6 +78,7 @@ impl ToolsConfig { session_source, } = params; let include_apply_patch_tool = features.enabled(Feature::ApplyPatchFreeform); + let include_ps_repl = features.enabled(Feature::PsRepl); let include_js_repl = features.enabled(Feature::JsRepl); let include_js_repl_tools_only = include_js_repl && features.enabled(Feature::JsReplToolsOnly); @@ -136,6 +138,7 @@ impl ToolsConfig { agent_roles: BTreeMap::new(), search_tool: include_search_tool, request_permission_enabled, + ps_repl_enabled: include_ps_repl, js_repl_enabled: include_js_repl, js_repl_tools_only: include_js_repl_tools_only, collab_tools: include_collab_tools, @@ -1329,6 +1332,47 @@ JS_SOURCE: /(?:\s*)(?:[^\s{\"`]|`[^`]|``[^`])[\s\S]*/ }) } +fn create_ps_repl_tool() -> ToolSpec { + const PS_REPL_FREEFORM_GRAMMAR: &str = r#" +start: pragma_source | plain_source + +pragma_source: PRAGMA_LINE NEWLINE ps_source +plain_source: PS_SOURCE + +ps_source: PS_SOURCE + +PRAGMA_LINE: /[ \t]*# codex-ps-repl:[^\r\n]*/ +NEWLINE: /\r?\n/ +PS_SOURCE: /(?:\s*)(?:[^\s{\"'`]|#[^\r\n]|`[^`])[\s\S]*/ +"#; + + ToolSpec::Freeform(FreeformTool { + name: "ps_repl".to_string(), + description: "Runs PowerShell in a persistent pwsh kernel. This is a freeform tool: send raw PowerShell source text, optionally with a first-line pragma like `# codex-ps-repl: timeout_ms=15000`; do not send JSON/quotes/markdown fences." + .to_string(), + format: FreeformToolFormat { + r#type: "grammar".to_string(), + syntax: "lark".to_string(), + definition: PS_REPL_FREEFORM_GRAMMAR.to_string(), + }, + }) +} + +fn create_ps_repl_reset_tool() -> ToolSpec { + ToolSpec::Function(ResponsesApiTool { + name: "ps_repl_reset".to_string(), + description: + "Restarts the ps_repl kernel for this run and clears persisted PowerShell session state." + .to_string(), + strict: false, + parameters: JsonSchema::Object { + properties: BTreeMap::new(), + required: None, + additional_properties: Some(false.into()), + }, + }) +} + fn create_js_repl_reset_tool() -> ToolSpec { ToolSpec::Function(ResponsesApiTool { name: "js_repl_reset".to_string(), @@ -1658,6 +1702,8 @@ pub(crate) fn build_specs( use crate::tools::handlers::McpResourceHandler; use crate::tools::handlers::MultiAgentHandler; use crate::tools::handlers::PlanHandler; + use crate::tools::handlers::PsReplHandler; + use crate::tools::handlers::PsReplResetHandler; use crate::tools::handlers::ReadFileHandler; use crate::tools::handlers::RequestUserInputHandler; use crate::tools::handlers::SearchToolBm25Handler; @@ -1683,6 +1729,8 @@ pub(crate) fn build_specs( default_mode_request_user_input: config.default_mode_request_user_input, }); let search_tool_handler = Arc::new(SearchToolBm25Handler); + let ps_repl_handler = Arc::new(PsReplHandler); + let ps_repl_reset_handler = Arc::new(PsReplResetHandler); let js_repl_handler = Arc::new(JsReplHandler); let js_repl_reset_handler = Arc::new(JsReplResetHandler); let request_permission_enabled = config.request_permission_enabled; @@ -1737,6 +1785,13 @@ pub(crate) fn build_specs( builder.push_spec(PLAN_TOOL.clone()); builder.register_handler("update_plan", plan_handler); + if config.ps_repl_enabled { + builder.push_spec(create_ps_repl_tool()); + builder.push_spec(create_ps_repl_reset_tool()); + builder.register_handler("ps_repl", ps_repl_handler); + builder.register_handler("ps_repl_reset", ps_repl_reset_handler); + } + if config.js_repl_enabled { builder.push_spec(create_js_repl_tool()); builder.push_spec(create_js_repl_reset_tool()); @@ -2234,6 +2289,62 @@ mod tests { ); } + #[test] + fn ps_repl_requires_feature_flag() { + let config = test_config(); + let model_info = + ModelsManager::construct_model_info_offline_for_tests("gpt-5-codex", &config); + let features = Features::with_defaults(); + + let tools_config = ToolsConfig::new(&ToolsConfigParams { + model_info: &model_info, + features: &features, + web_search_mode: Some(WebSearchMode::Cached), + session_source: SessionSource::Cli, + }); + let (tools, _) = build_specs(&tools_config, None, None, &[]).build(); + + assert!( + !tools.iter().any(|tool| tool.spec.name() == "ps_repl"), + "ps_repl should be disabled when the feature is off" + ); + assert!( + !tools.iter().any(|tool| tool.spec.name() == "ps_repl_reset"), + "ps_repl_reset should be disabled when the feature is off" + ); + } + + #[test] + fn ps_repl_enabled_adds_tools() { + let config = test_config(); + let model_info = + ModelsManager::construct_model_info_offline_for_tests("gpt-5-codex", &config); + let mut features = Features::with_defaults(); + features.enable(Feature::PsRepl); + + let tools_config = ToolsConfig::new(&ToolsConfigParams { + model_info: &model_info, + features: &features, + web_search_mode: Some(WebSearchMode::Cached), + session_source: SessionSource::Cli, + }); + let (tools, _) = build_specs(&tools_config, None, None, &[]).build(); + assert_contains_tool_names(&tools, &["ps_repl", "ps_repl_reset"]); + } + + #[test] + fn ps_repl_freeform_grammar_mentions_pragma_and_ps_source() { + let ToolSpec::Freeform(FreeformTool { format, .. }) = create_ps_repl_tool() else { + panic!("ps_repl should use a freeform tool spec"); + }; + + assert_eq!(format.syntax, "lark"); + assert!(format.definition.contains("PRAGMA_LINE")); + assert!(format.definition.contains("PS_SOURCE")); + assert!(format.definition.contains("codex-ps-repl:")); + assert!(!format.definition.contains("(?!")); + } + #[test] fn js_repl_enabled_adds_tools() { let config = test_config(); diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index e9867caa76..f3734116eb 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -97,6 +97,7 @@ mod personality; mod personality_migration; mod plugins; mod prompt_caching; +mod ps_repl; mod quota_exceeded; mod read_file; mod realtime_conversation; diff --git a/codex-rs/core/tests/suite/ps_repl.rs b/codex-rs/core/tests/suite/ps_repl.rs new file mode 100644 index 0000000000..a88d2f9b08 --- /dev/null +++ b/codex-rs/core/tests/suite/ps_repl.rs @@ -0,0 +1,484 @@ +#![allow(clippy::expect_used, clippy::unwrap_used)] + +use anyhow::Result; +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use codex_core::features::Feature; +use codex_protocol::protocol::EventMsg; +use core_test_support::responses; +use core_test_support::responses::ResponseMock; +use core_test_support::responses::ResponsesRequest; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_custom_tool_call; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::sse; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event_match; +use serde_json::Value; +use std::fs; +#[cfg(unix)] +use std::os::unix::fs::PermissionsExt; +use std::path::Path; +use tempfile::tempdir; +use wiremock::MockServer; + +fn custom_tool_output_text_and_success( + req: &ResponsesRequest, + call_id: &str, +) -> (String, Option) { + let (output, success) = req + .custom_tool_call_output_content_and_success(call_id) + .expect("custom tool output should be present"); + (output.unwrap_or_default(), success) +} + +fn tool_names(body: &serde_json::Value) -> Vec { + body["tools"] + .as_array() + .expect("tools array should be present") + .iter() + .map(|tool| { + tool.get("name") + .and_then(|value| value.as_str()) + .or_else(|| tool.get("type").and_then(|value| value.as_str())) + .expect("tool should have a name or type") + .to_string() + }) + .collect() +} + +fn write_too_old_pwsh_script(dir: &Path) -> Result { + #[cfg(windows)] + { + let path = dir.join("old-pwsh.cmd"); + fs::write(&path, "@echo off\r\necho PowerShell 5.1.0\r\n")?; + Ok(path) + } + + #[cfg(unix)] + { + let path = dir.join("old-pwsh.sh"); + fs::write(&path, "#!/bin/sh\necho PowerShell 5.1.0\n")?; + let mut permissions = fs::metadata(&path)?.permissions(); + permissions.set_mode(0o755); + fs::set_permissions(&path, permissions)?; + Ok(path) + } + + #[cfg(not(any(unix, windows)))] + { + anyhow::bail!("unsupported platform for ps_repl test fixture"); + } +} + +fn write_test_png(dir: &Path) -> Result { + let path = dir.join("dot.png"); + let png_bytes = BASE64_STANDARD.decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==", + )?; + fs::write(&path, png_bytes)?; + Ok(path) +} + +fn ps_single_quote(input: &Path) -> String { + input.display().to_string().replace('\'', "''") +} + +async fn run_ps_repl_turn( + server: &MockServer, + prompt: &str, + calls: &[(&str, &str)], +) -> Result { + let test = test_codex() + .with_config(|config| { + config.features.enable(Feature::PsRepl); + }) + .build(server) + .await?; + + let mut first_events = vec![ev_response_created("resp-1")]; + for (call_id, ps_input) in calls { + first_events.push(ev_custom_tool_call(call_id, "ps_repl", ps_input)); + } + first_events.push(ev_completed("resp-1")); + responses::mount_sse_once(server, sse(first_events)).await; + + let second_mock = responses::mount_sse_once( + server, + sse(vec![ + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]), + ) + .await; + + test.submit_turn(prompt).await?; + Ok(second_mock) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn ps_repl_is_not_advertised_when_startup_pwsh_is_incompatible() -> Result<()> { + skip_if_no_network!(Ok(())); + if std::env::var_os("CODEX_PS_REPL_PATH").is_some() { + return Ok(()); + } + + let server = responses::start_mock_server().await; + let temp = tempdir()?; + let old_pwsh = write_too_old_pwsh_script(temp.path())?; + + let test = test_codex() + .with_config(move |config| { + config.features.enable(Feature::PsRepl); + config.ps_repl_path = Some(old_pwsh); + }) + .build(&server) + .await?; + let warning = wait_for_event_match(&test.codex, |event| match event { + EventMsg::Warning(ev) if ev.message.contains("Disabled `ps_repl` for this session") => { + Some(ev.message.clone()) + } + _ => None, + }) + .await; + assert!( + warning.contains("PowerShell runtime"), + "warning should explain the PowerShell compatibility issue: {warning}" + ); + + let request_mock = responses::mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("msg-1", "done"), + ev_completed("resp-1"), + ]), + ) + .await; + + test.submit_turn("hello").await?; + + let body = request_mock.single_request().body_json(); + let tools = tool_names(&body); + assert!( + !tools.iter().any(|tool| tool == "ps_repl"), + "ps_repl should be omitted when startup validation fails: {tools:?}" + ); + assert!( + !tools.iter().any(|tool| tool == "ps_repl_reset"), + "ps_repl_reset should be omitted when startup validation fails: {tools:?}" + ); + let instructions = body["instructions"].as_str().unwrap_or_default(); + assert!( + !instructions.contains("## PowerShell REPL (pwsh)"), + "startup instructions should not mention ps_repl when it is disabled: {instructions}" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn ps_repl_persists_variables_functions_and_modules() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let test = test_codex() + .with_config(|config| { + config.features.enable(Feature::PsRepl); + }) + .build(&server) + .await?; + + responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_custom_tool_call( + "call-1", + "ps_repl", + r#" +$x = 41 +New-Module -Name CodexPsTestModule -ScriptBlock { + function Get-CodexValue { 42 } +} | Import-Module +function Add-One { + param([int]$Value) + $Value + 1 +} +Write-Output "state-ready" +"#, + ), + ev_completed("resp-1"), + ]), + ) + .await; + let second_mock = responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-2"), + ev_custom_tool_call( + "call-2", + "ps_repl", + r#" +Write-Output ($x + 1) +Write-Output (Add-One -Value 1) +Write-Output (Get-CodexValue) +"#, + ), + ev_completed("resp-2"), + ]), + ) + .await; + let third_mock = responses::mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("msg-1", "done"), + ev_completed("resp-3"), + ]), + ) + .await; + + test.submit_turn("run ps_repl twice").await?; + + let req2 = second_mock.single_request(); + let (first_output, first_success) = custom_tool_output_text_and_success(&req2, "call-1"); + assert_ne!( + first_success, + Some(false), + "first ps_repl call failed unexpectedly: {first_output}" + ); + assert!(first_output.contains("state-ready")); + + let req3 = third_mock.single_request(); + let (second_output, second_success) = custom_tool_output_text_and_success(&req3, "call-2"); + assert_ne!( + second_success, + Some(false), + "second ps_repl call failed unexpectedly: {second_output}" + ); + let lines = second_output + .lines() + .map(str::trim) + .filter(|line| !line.is_empty()) + .collect::>(); + assert!( + lines.iter().filter(|line| **line == "42").count() >= 2, + "expected persisted variable and module output, got: {second_output}" + ); + assert!( + lines.contains(&"2"), + "expected persisted function output, got: {second_output}" + ); + assert!( + !second_output.contains("Get-CodexValue"), + "unexpected formatting leak: {second_output}" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn ps_repl_can_invoke_builtin_tools() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let mock = run_ps_repl_turn( + &server, + "use ps_repl to call a tool", + &[( + "call-1", + "$toolOut = Invoke-CodexTool -Name list_mcp_resources -Arguments @{}; Write-Output $toolOut.type", + )], + ) + .await?; + + let req = mock.single_request(); + let (output, success) = custom_tool_output_text_and_success(&req, "call-1"); + assert_ne!( + success, + Some(false), + "ps_repl call failed unexpectedly: {output}" + ); + assert!(output.contains("function_call_output")); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn ps_repl_tool_call_rejects_recursive_ps_repl_invocation() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let mock = run_ps_repl_turn( + &server, + "use ps_repl recursively", + &[( + "call-1", + r#" +try { + Invoke-CodexTool -Name ps_repl -Arguments "Write-Output 'recursive'" | Out-Null + Write-Output "unexpected-success" +} catch { + Write-Output $_.Exception.Message +} +"#, + )], + ) + .await?; + + let req = mock.single_request(); + let (output, success) = custom_tool_output_text_and_success(&req, "call-1"); + assert_ne!( + success, + Some(false), + "ps_repl call failed unexpectedly: {output}" + ); + assert!( + output.contains("ps_repl cannot invoke itself"), + "expected recursion guard message, got output: {output}" + ); + assert!( + !output.contains("unexpected-success"), + "recursive ps_repl call unexpectedly succeeded: {output}" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn ps_repl_resets_after_timeout_and_accepts_followup_execution() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let test = test_codex() + .with_config(|config| { + config.features.enable(Feature::PsRepl); + }) + .build(&server) + .await?; + + responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_custom_tool_call( + "call-1", + "ps_repl", + "# codex-ps-repl: timeout_ms=50\nStart-Sleep -Milliseconds 500", + ), + ev_completed("resp-1"), + ]), + ) + .await; + let second_mock = responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-2"), + ev_custom_tool_call("call-2", "ps_repl", "Write-Output 'healthy'"), + ev_completed("resp-2"), + ]), + ) + .await; + let third_mock = responses::mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("msg-1", "done"), + ev_completed("resp-3"), + ]), + ) + .await; + + test.submit_turn("run ps_repl after timeout").await?; + + let req2 = second_mock.single_request(); + let (first_output, first_success) = custom_tool_output_text_and_success(&req2, "call-1"); + assert_ne!( + first_success, + Some(true), + "timeout should not report success: {first_output}" + ); + assert!( + first_output.contains("ps_repl execution timed out"), + "expected timeout output, got: {first_output}" + ); + + let req3 = third_mock.single_request(); + let (second_output, second_success) = custom_tool_output_text_and_success(&req3, "call-2"); + assert_ne!( + second_success, + Some(false), + "ps_repl follow-up execution failed unexpectedly: {second_output}" + ); + assert!(second_output.contains("healthy")); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn ps_repl_captures_standard_powershell_output_streams() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let mock = run_ps_repl_turn( + &server, + "capture powershell output", + &[( + "call-1", + "Write-Output 'stdout'; Write-Warning 'warn-stream'", + )], + ) + .await?; + + let req = mock.single_request(); + let (output, success) = custom_tool_output_text_and_success(&req, "call-1"); + assert_ne!( + success, + Some(false), + "ps_repl call failed unexpectedly: {output}" + ); + assert!(output.contains("stdout")); + assert!(output.contains("warn-stream")); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn ps_repl_view_image_propagates_content_items() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let temp = tempdir()?; + let png_path = write_test_png(temp.path())?; + let png_path = ps_single_quote(&png_path); + let script = + format!("$null = Invoke-CodexTool -Name view_image -Arguments @{{ path = '{png_path}' }}"); + + let mock = run_ps_repl_turn( + &server, + "render an image via ps_repl", + &[("call-1", &script)], + ) + .await?; + + let req = mock.single_request(); + let custom_output = req.custom_tool_call_output("call-1"); + let output_items = custom_output + .get("output") + .and_then(Value::as_array) + .expect("custom_tool_call_output should be a content item array"); + let image_url = output_items + .iter() + .find_map(|item| { + (item.get("type").and_then(Value::as_str) == Some("input_image")) + .then(|| item.get("image_url").and_then(Value::as_str)) + .flatten() + }) + .expect("image_url present in ps_repl custom tool output"); + assert!( + image_url.starts_with("data:image/png;base64,"), + "expected png data URL, got {image_url}" + ); + + Ok(()) +} diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index c0a559e067..9a7a1bbbb1 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -277,6 +277,7 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result model_provider: model_provider.clone(), codex_linux_sandbox_exe: arg0_paths.codex_linux_sandbox_exe.clone(), main_execve_wrapper_exe: arg0_paths.main_execve_wrapper_exe.clone(), + ps_repl_path: None, js_repl_node_path: None, js_repl_node_module_dirs: None, zsh_path: None, diff --git a/docs/ps_repl.md b/docs/ps_repl.md new file mode 100644 index 0000000000..3a5ccc29d7 --- /dev/null +++ b/docs/ps_repl.md @@ -0,0 +1,68 @@ +# PowerShell REPL (`ps_repl`) + +`ps_repl` runs PowerShell in a persistent `pwsh`-backed kernel. + +## Feature gate + +`ps_repl` is disabled by default and only appears when: + +```toml +[features] +ps_repl = true +``` + +The initial rollout stage is under development. + +## PowerShell runtime + +`ps_repl` requires PowerShell 7 or newer. + +Runtime resolution order: + +1. `CODEX_PS_REPL_PATH` environment variable +2. `ps_repl_path` in config/profile +3. `pwsh` discovered on `PATH` + +You can configure an explicit runtime path: + +```toml +ps_repl_path = "/absolute/path/to/pwsh" +``` + +If only Windows PowerShell (`powershell.exe`) is available, `ps_repl` stays disabled and emits a startup warning telling you to install PowerShell 7. + +## Usage + +- `ps_repl` is a freeform tool: send raw PowerShell source text. +- Optional first-line pragma: + - `# codex-ps-repl: timeout_ms=15000` +- Variables, functions, aliases, imported modules, environment changes, and `$LASTEXITCODE` persist across calls. +- Use `ps_repl_reset` to clear the kernel state. + +## Helper APIs inside the kernel + +`ps_repl` exposes these helpers: + +- `$CodexTmpDir`: per-session scratch directory path. +- `Invoke-CodexTool -Name -Arguments `: executes a normal Codex tool call from inside `ps_repl`. +- `$Codex.TmpDir`: alias for the scratch directory. +- `$Codex.Tool(, )`: thin alias to `Invoke-CodexTool`. + +`Invoke-CodexTool` returns the raw tool output object. Nested tool calls can also return multimodal content such as `view_image` results. + +To share generated images with the model, write a file under `$CodexTmpDir`, call: + +```powershell +Invoke-CodexTool -Name view_image -Arguments @{ path = "/absolute/path" } +``` + +Then delete the file. + +## Output and transport + +`ps_repl` uses a JSON-line transport over stdio. + +- Safe output forms: pipeline output, `Write-Output`, `Write-Host`, `Write-Verbose`, `Write-Warning` +- Avoid: direct `[Console]::Write*`, raw StdOut/StdErr writes, or other host-level output that bypasses PowerShell streams + +Bypassing PowerShell streams can corrupt the transport between the Rust host and the persistent kernel.