diff --git a/.github/workflows/issue-deduplicator.yml b/.github/workflows/issue-deduplicator.yml index eed4d46167..2bac731e61 100644 --- a/.github/workflows/issue-deduplicator.yml +++ b/.github/workflows/issue-deduplicator.yml @@ -105,6 +105,10 @@ jobs: const issues = Array.isArray(parsed?.issues) ? parsed.issues : []; const currentIssueNumber = String(context.payload.issue.number); + + console.log(`Current issue number: ${currentIssueNumber}`); + console.log(issues); + const filteredIssues = issues.filter((value) => String(value) !== currentIssueNumber); if (filteredIssues.length === 0) { @@ -113,7 +117,8 @@ jobs: } const lines = [ - 'Potential duplicates detected:', + 'Potential duplicates detected. Please review them and close your issue if it is a duplicate.', + '', ...filteredIssues.map((value) => `- #${String(value)}`), '', '*Powered by [Codex Action](https://github.com/openai/codex-action)*']; diff --git a/AGENTS.md b/AGENTS.md index accfe447f7..832f1d6543 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -73,3 +73,28 @@ If you don’t have the tool: ### Test assertions - Tests should use pretty_assertions::assert_eq for clearer diffs. Import this at the top of the test module if it isn't already. + +### Integration tests (core) + +- Prefer the utilities in `core_test_support::responses` when writing end-to-end Codex tests. + +- All `mount_sse*` helpers return a `ResponseMock`; hold onto it so you can assert against outbound `/responses` POST bodies. +- Use `ResponseMock::single_request()` when a test should only issue one POST, or `ResponseMock::requests()` to inspect every captured `ResponsesRequest`. +- `ResponsesRequest` exposes helpers (`body_json`, `input`, `function_call_output`, `custom_tool_call_output`, `call_output`, `header`, `path`, `query_param`) so assertions can target structured payloads instead of manual JSON digging. +- Build SSE payloads with the provided `ev_*` constructors and the `sse(...)`. + +- Typical pattern: + + ```rust + let mock = responses::mount_sse_once(&server, responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call(call_id, "shell", &serde_json::to_string(&args)?), + responses::ev_completed("resp-1"), + ])).await; + + codex.submit(Op::UserTurn { ... }).await?; + + // Assert request body if needed. + let request = mock.single_request(); + // assert using request.function_call_output(call_id) or request.json_body() or other helpers. + ``` diff --git a/README.md b/README.md index f17f605e18..c9cc4e1fd4 100644 --- a/README.md +++ b/README.md @@ -61,7 +61,7 @@ You can also use Codex with an API key, but this requires [additional setup](./d ### Model Context Protocol (MCP) -Codex CLI supports [MCP servers](./docs/advanced.md#model-context-protocol-mcp). Enable by adding an `mcp_servers` section to your `~/.codex/config.toml`. +Codex can access MCP servers. To configure them, refer to the [config docs](./docs/config.md#mcp_servers). ### Configuration diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 350773a3c5..8c992a6f8e 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -992,7 +992,7 @@ dependencies = [ "tokio-stream", "tracing", "tracing-subscriber", - "unicode-width 0.1.14", + "unicode-width 0.2.1", ] [[package]] @@ -1424,6 +1424,8 @@ dependencies = [ "tracing", "tracing-appender", "tracing-subscriber", + "tree-sitter-bash", + "tree-sitter-highlight", "unicode-segmentation", "unicode-width 0.2.1", "url", @@ -4770,9 +4772,9 @@ dependencies = [ [[package]] name = "rmcp" -version = "0.8.0" +version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "583d060e99feb3a3683fb48a1e4bf5f8d4a50951f429726f330ee5ff548837f8" +checksum = "6f35acda8f89fca5fd8c96cae3c6d5b4c38ea0072df4c8030915f3b5ff469c1c" dependencies = [ "base64", "bytes", @@ -4804,9 +4806,9 @@ dependencies = [ [[package]] name = "rmcp-macros" -version = "0.8.0" +version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "421d8b0ba302f479214889486f9550e63feca3af310f1190efcf6e2016802693" +checksum = "c9f1d5220aaa23b79c3d02e18f7a554403b3ccea544bbb6c69d6bcb3e854a274" dependencies = [ "darling 0.21.3", "proc-macro2", @@ -6019,6 +6021,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] @@ -6260,9 +6263,9 @@ dependencies = [ [[package]] name = "tree-sitter" -version = "0.25.9" +version = "0.25.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ccd2a058a86cfece0bf96f7cce1021efef9c8ed0e892ab74639173e5ed7a34fa" +checksum = "78f873475d258561b06f1c595d93308a7ed124d9977cb26b148c2084a4a3cc87" dependencies = [ "cc", "regex", @@ -6282,6 +6285,18 @@ dependencies = [ "tree-sitter-language", ] +[[package]] +name = "tree-sitter-highlight" +version = "0.25.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adc5f880ad8d8f94e88cb81c3557024cf1a8b75e3b504c50481ed4f5a6006ff3" +dependencies = [ + "regex", + "streaming-iterator", + "thiserror 2.0.16", + "tree-sitter", +] + [[package]] name = "tree-sitter-language" version = "0.1.5" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index bcbe844509..ca40b1a536 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -175,8 +175,9 @@ tracing = "0.1.41" tracing-appender = "0.2.3" tracing-subscriber = "0.3.20" tracing-test = "0.2.5" -tree-sitter = "0.25.9" -tree-sitter-bash = "0.25.0" +tree-sitter = "0.25.10" +tree-sitter-bash = "0.25" +tree-sitter-highlight = "0.25.10" ts-rs = "11" unicode-segmentation = "1.12.0" unicode-width = "0.2" diff --git a/codex-rs/README.md b/codex-rs/README.md index f5a5b838fc..a05a8d48b1 100644 --- a/codex-rs/README.md +++ b/codex-rs/README.md @@ -23,9 +23,15 @@ Codex supports a rich set of configuration options. Note that the Rust CLI uses ### Model Context Protocol Support -Codex CLI functions as an MCP client that can connect to MCP servers on startup. See the [`mcp_servers`](../docs/config.md#mcp_servers) section in the configuration documentation for details. +#### MCP client -It is still experimental, but you can also launch Codex as an MCP _server_ by running `codex mcp-server`. Use the [`@modelcontextprotocol/inspector`](https://github.com/modelcontextprotocol/inspector) to try it out: +Codex CLI functions as an MCP client that allows the Codex CLI and IDE extension to connect to MCP servers on startup. See the [`configuration documentation`](../docs/config.md#mcp_servers) for details. + +#### MCP server (experimental) + +Codex can be launched as an MCP _server_ by running `codex mcp-server`. This allows _other_ MCP clients to use Codex as a tool for another agent. + +Use the [`@modelcontextprotocol/inspector`](https://github.com/modelcontextprotocol/inspector) to try it out: ```shell npx @modelcontextprotocol/inspector codex mcp-server diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 04a7ff88f3..d889823298 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -157,9 +157,7 @@ struct LoginCommand { )] api_key: Option, - /// EXPERIMENTAL: Use device code flow (not yet supported) - /// This feature is experimental and may changed in future releases. - #[arg(long = "experimental_use-device-code", hide = true)] + #[arg(long = "use-device-code")] use_device_code: bool, /// EXPERIMENTAL: Use custom OAuth issuer base URL (advanced) @@ -294,7 +292,8 @@ async fn cli_main(codex_linux_sandbox_exe: Option) -> anyhow::Result<() last, config_overrides, ); - codex_tui::run_main(interactive, codex_linux_sandbox_exe).await?; + let exit_info = codex_tui::run_main(interactive, codex_linux_sandbox_exe).await?; + print_exit_messages(exit_info); } Some(Subcommand::Login(mut login_cli)) => { prepend_config_flags( diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index a105dbbac1..0cf3c0e228 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -4,6 +4,7 @@ use anyhow::Context; use anyhow::Result; use anyhow::anyhow; use anyhow::bail; +use clap::ArgGroup; use codex_common::CliConfigOverrides; use codex_core::config::Config; use codex_core::config::ConfigOverrides; @@ -77,13 +78,61 @@ pub struct AddArgs { /// Name for the MCP server configuration. pub name: String, - /// Environment variables to set when launching the server. - #[arg(long, value_parser = parse_env_pair, value_name = "KEY=VALUE")] - pub env: Vec<(String, String)>, + #[command(flatten)] + pub transport_args: AddMcpTransportArgs, +} +#[derive(Debug, clap::Args)] +#[command( + group( + ArgGroup::new("transport") + .args(["command", "url"]) + .required(true) + .multiple(false) + ) +)] +pub struct AddMcpTransportArgs { + #[command(flatten)] + pub stdio: Option, + + #[command(flatten)] + pub streamable_http: Option, +} + +#[derive(Debug, clap::Args)] +pub struct AddMcpStdioArgs { /// Command to launch the MCP server. - #[arg(trailing_var_arg = true, num_args = 1..)] + /// Use --url for a streamable HTTP server. + #[arg( + trailing_var_arg = true, + num_args = 0.., + )] pub command: Vec, + + /// Environment variables to set when launching the server. + /// Only valid with stdio servers. + #[arg( + long, + value_parser = parse_env_pair, + value_name = "KEY=VALUE", + )] + pub env: Vec<(String, String)>, +} + +#[derive(Debug, clap::Args)] +pub struct AddMcpStreamableHttpArgs { + /// URL for a streamable HTTP MCP server. + #[arg(long)] + pub url: String, + + /// Optional environment variable to read for a bearer token. + /// Only valid with streamable HTTP servers. + #[arg( + long = "bearer-token-env-var", + value_name = "ENV_VAR", + requires = "url" + )] + pub bearer_token_env_var: Option, } #[derive(Debug, clap::Parser)] @@ -140,37 +189,51 @@ async fn run_add(config_overrides: &CliConfigOverrides, add_args: AddArgs) -> Re // Validate any provided overrides even though they are not currently applied. config_overrides.parse_overrides().map_err(|e| anyhow!(e))?; - let AddArgs { name, env, command } = add_args; + let AddArgs { + name, + transport_args, + } = add_args; validate_server_name(&name)?; - let mut command_parts = command.into_iter(); - let command_bin = command_parts - .next() - .ok_or_else(|| anyhow!("command is required"))?; - let command_args: Vec = command_parts.collect(); - - let env_map = if env.is_empty() { - None - } else { - let mut map = HashMap::new(); - for (key, value) in env { - map.insert(key, value); - } - Some(map) - }; - let codex_home = find_codex_home().context("failed to resolve CODEX_HOME")?; let mut servers = load_global_mcp_servers(&codex_home) .await .with_context(|| format!("failed to load MCP servers from {}", codex_home.display()))?; - let new_entry = McpServerConfig { - transport: McpServerTransportConfig::Stdio { - command: command_bin, - args: command_args, - env: env_map, + let transport = match transport_args { + AddMcpTransportArgs { + stdio: Some(stdio), .. + } => { + let mut command_parts = stdio.command.into_iter(); + let command_bin = command_parts + .next() + .ok_or_else(|| anyhow!("command is required"))?; + let command_args: Vec = command_parts.collect(); + + let env_map = if stdio.env.is_empty() { + None + } else { + Some(stdio.env.into_iter().collect::>()) + }; + McpServerTransportConfig::Stdio { + command: command_bin, + args: command_args, + env: env_map, + } + } + AddMcpTransportArgs { + streamable_http: Some(streamable_http), + .. + } => McpServerTransportConfig::StreamableHttp { + url: streamable_http.url, + bearer_token_env_var: streamable_http.bearer_token_env_var, }, + AddMcpTransportArgs { .. } => bail!("exactly one of --command or --url must be provided"), + }; + + let new_entry = McpServerConfig { + transport, startup_timeout_sec: None, tool_timeout_sec: None, }; @@ -236,7 +299,7 @@ async fn run_login(config_overrides: &CliConfigOverrides, login_args: LoginArgs) _ => bail!("OAuth login is only supported for streamable HTTP servers."), }; - perform_oauth_login(&name, &url).await?; + perform_oauth_login(&name, &url, config.mcp_oauth_credentials_store_mode).await?; println!("Successfully logged in to MCP server '{name}'."); Ok(()) } @@ -259,7 +322,7 @@ async fn run_logout(config_overrides: &CliConfigOverrides, logout_args: LogoutAr _ => bail!("OAuth logout is only supported for streamable_http transports."), }; - match delete_oauth_tokens(&name, &url) { + match delete_oauth_tokens(&name, &url, config.mcp_oauth_credentials_store_mode) { Ok(true) => println!("Removed OAuth credentials for '{name}'."), Ok(false) => println!("No OAuth credentials stored for '{name}'."), Err(err) => return Err(anyhow!("failed to delete OAuth credentials: {err}")), @@ -288,11 +351,14 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> "args": args, "env": env, }), - McpServerTransportConfig::StreamableHttp { url, bearer_token } => { + McpServerTransportConfig::StreamableHttp { + url, + bearer_token_env_var, + } => { serde_json::json!({ "type": "streamable_http", "url": url, - "bearer_token": bearer_token, + "bearer_token_env_var": bearer_token_env_var, }) } }; @@ -345,13 +411,15 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> }; stdio_rows.push([name.clone(), command.clone(), args_display, env_display]); } - McpServerTransportConfig::StreamableHttp { url, bearer_token } => { - let has_bearer = if bearer_token.is_some() { - "True" - } else { - "False" - }; - http_rows.push([name.clone(), url.clone(), has_bearer.into()]); + McpServerTransportConfig::StreamableHttp { + url, + bearer_token_env_var, + } => { + http_rows.push([ + name.clone(), + url.clone(), + bearer_token_env_var.clone().unwrap_or("-".to_string()), + ]); } } } @@ -396,7 +464,7 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> } if !http_rows.is_empty() { - let mut widths = ["Name".len(), "Url".len(), "Has Bearer Token".len()]; + let mut widths = ["Name".len(), "Url".len(), "Bearer Token Env Var".len()]; for row in &http_rows { for (i, cell) in row.iter().enumerate() { widths[i] = widths[i].max(cell.len()); @@ -407,7 +475,7 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> "{: Re "args": args, "env": env, }), - McpServerTransportConfig::StreamableHttp { url, bearer_token } => serde_json::json!({ + McpServerTransportConfig::StreamableHttp { + url, + bearer_token_env_var, + } => serde_json::json!({ "type": "streamable_http", "url": url, - "bearer_token": bearer_token, + "bearer_token_env_var": bearer_token_env_var, }), }; let output = serde_json::to_string_pretty(&serde_json::json!({ @@ -493,11 +564,14 @@ async fn run_get(config_overrides: &CliConfigOverrides, get_args: GetArgs) -> Re }; println!(" env: {env_display}"); } - McpServerTransportConfig::StreamableHttp { url, bearer_token } => { + McpServerTransportConfig::StreamableHttp { + url, + bearer_token_env_var, + } => { println!(" transport: streamable_http"); println!(" url: {url}"); - let bearer = bearer_token.as_deref().unwrap_or("-"); - println!(" bearer_token: {bearer}"); + let env_var = bearer_token_env_var.as_deref().unwrap_or("-"); + println!(" bearer_token_env_var: {env_var}"); } } if let Some(timeout) = server.startup_timeout_sec { diff --git a/codex-rs/cli/tests/mcp_add_remove.rs b/codex-rs/cli/tests/mcp_add_remove.rs index 6530760e4f..9ccb9f73d0 100644 --- a/codex-rs/cli/tests/mcp_add_remove.rs +++ b/codex-rs/cli/tests/mcp_add_remove.rs @@ -93,3 +93,116 @@ async fn add_with_env_preserves_key_order_and_values() -> Result<()> { Ok(()) } + +#[tokio::test] +async fn add_streamable_http_without_manual_token() -> Result<()> { + let codex_home = TempDir::new()?; + + let mut add_cmd = codex_command(codex_home.path())?; + add_cmd + .args(["mcp", "add", "github", "--url", "https://example.com/mcp"]) + .assert() + .success(); + + let servers = load_global_mcp_servers(codex_home.path()).await?; + let github = servers.get("github").expect("github server should exist"); + match &github.transport { + McpServerTransportConfig::StreamableHttp { + url, + bearer_token_env_var, + } => { + assert_eq!(url, "https://example.com/mcp"); + assert!(bearer_token_env_var.is_none()); + } + other => panic!("unexpected transport: {other:?}"), + } + + assert!(!codex_home.path().join(".credentials.json").exists()); + assert!(!codex_home.path().join(".env").exists()); + + Ok(()) +} + +#[tokio::test] +async fn add_streamable_http_with_custom_env_var() -> Result<()> { + let codex_home = TempDir::new()?; + + let mut add_cmd = codex_command(codex_home.path())?; + add_cmd + .args([ + "mcp", + "add", + "issues", + "--url", + "https://example.com/issues", + "--bearer-token-env-var", + "GITHUB_TOKEN", + ]) + .assert() + .success(); + + let servers = load_global_mcp_servers(codex_home.path()).await?; + let issues = servers.get("issues").expect("issues server should exist"); + match &issues.transport { + McpServerTransportConfig::StreamableHttp { + url, + bearer_token_env_var, + } => { + assert_eq!(url, "https://example.com/issues"); + assert_eq!(bearer_token_env_var.as_deref(), Some("GITHUB_TOKEN")); + } + other => panic!("unexpected transport: {other:?}"), + } + Ok(()) +} + +#[tokio::test] +async fn add_streamable_http_rejects_removed_flag() -> Result<()> { + let codex_home = TempDir::new()?; + + let mut add_cmd = codex_command(codex_home.path())?; + add_cmd + .args([ + "mcp", + "add", + "github", + "--url", + "https://example.com/mcp", + "--with-bearer-token", + ]) + .assert() + .failure() + .stderr(contains("--with-bearer-token")); + + let servers = load_global_mcp_servers(codex_home.path()).await?; + assert!(servers.is_empty()); + + Ok(()) +} + +#[tokio::test] +async fn add_cant_add_command_and_url() -> Result<()> { + let codex_home = TempDir::new()?; + + let mut add_cmd = codex_command(codex_home.path())?; + add_cmd + .args([ + "mcp", + "add", + "github", + "--url", + "https://example.com/mcp", + "--command", + "--", + "echo", + "hello", + ]) + .assert() + .failure() + .stderr(contains("unexpected argument '--command' found")); + + let servers = load_global_mcp_servers(codex_home.path()).await?; + assert!(servers.is_empty()); + + Ok(()) +} diff --git a/codex-rs/cloud-tasks/Cargo.toml b/codex-rs/cloud-tasks/Cargo.toml index 13d7252760..46044fbb8d 100644 --- a/codex-rs/cloud-tasks/Cargo.toml +++ b/codex-rs/cloud-tasks/Cargo.toml @@ -11,10 +11,10 @@ path = "src/lib.rs" workspace = true [dependencies] -anyhow = "1" -base64 = "0.22" -chrono = { version = "0.4", features = ["serde"] } -clap = { version = "4", features = ["derive"] } +anyhow = { workspace = true } +base64 = { workspace = true } +chrono = { workspace = true, features = ["serde"] } +clap = { workspace = true, features = ["derive"] } codex-cloud-tasks-client = { path = "../cloud-tasks-client", features = [ "mock", "online", @@ -23,16 +23,16 @@ codex-common = { path = "../common", features = ["cli"] } codex-core = { path = "../core" } codex-login = { path = "../login" } codex-tui = { path = "../tui" } -crossterm = { version = "0.28.1", features = ["event-stream"] } -ratatui = { version = "0.29.0" } -reqwest = { version = "0.12", features = ["json"] } -serde = { version = "1", features = ["derive"] } -serde_json = "1" -tokio = { version = "1", features = ["macros", "rt-multi-thread"] } -tokio-stream = "0.1.17" -tracing = { version = "0.1.41", features = ["log"] } -tracing-subscriber = { version = "0.3.19", features = ["env-filter"] } -unicode-width = "0.1" +crossterm = { workspace = true, features = ["event-stream"] } +ratatui = { workspace = true } +reqwest = { workspace = true, features = ["json"] } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } +tokio-stream = { workspace = true } +tracing = { workspace = true, features = ["log"] } +tracing-subscriber = { workspace = true, features = ["env-filter"] } +unicode-width = { workspace = true } [dev-dependencies] -async-trait = "0.1" +async-trait = { workspace = true } diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 367ccbcef2..4259e64fc4 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -61,7 +61,7 @@ tokio = { workspace = true, features = [ "rt-multi-thread", "signal", ] } -tokio-util = { workspace = true } +tokio-util = { workspace = true, features = ["rt"] } toml = { workspace = true } toml_edit = { workspace = true } tracing = { workspace = true, features = ["log"] } diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 9cdb1d4c01..a62ec547e1 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -23,7 +23,9 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::TaskStartedEvent; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnContextItem; +use futures::future::BoxFuture; use futures::prelude::*; +use futures::stream::FuturesOrdered; use mcp_types::CallToolResult; use serde_json; use serde_json::Value; @@ -363,6 +365,7 @@ impl Session { let mcp_fut = McpConnectionManager::new( config.mcp_servers.clone(), config.use_experimental_use_rmcp_client, + config.mcp_oauth_credentials_store_mode, ); let default_shell_fut = shell::default_user_shell(); let history_meta_fut = crate::message_history::history_metadata(&config); @@ -1594,7 +1597,7 @@ async fn spawn_review_thread( // Seed the child task with the review prompt as the initial user message. let input: Vec = vec![InputItem::Text { - text: format!("{base_instructions}\n\n---\n\nNow, here's your task: {review_prompt}"), + text: review_prompt, }]; let tc = Arc::new(review_turn_context); @@ -2111,14 +2114,15 @@ async fn try_run_turn( sess.persist_rollout_items(&[rollout_item]).await; let mut stream = turn_context.client.clone().stream(&prompt).await?; - let mut output = Vec::new(); - let mut tool_runtime = ToolCallRuntime::new( + let tool_runtime = ToolCallRuntime::new( Arc::clone(&router), Arc::clone(&sess), Arc::clone(&turn_context), Arc::clone(&turn_diff_tracker), sub_id.to_string(), ); + let mut output: FuturesOrdered>> = + FuturesOrdered::new(); loop { // Poll the next item from the model stream. We must inspect *both* Ok and Err @@ -2126,9 +2130,8 @@ async fn try_run_turn( // `response.completed`) bubble up and trigger the caller's retry logic. let event = stream.next().await; let event = match event { - Some(event) => event, + Some(res) => res?, None => { - tool_runtime.abort_all(); return Err(CodexErr::Stream( "stream closed before response.completed".into(), None, @@ -2136,14 +2139,8 @@ async fn try_run_turn( } }; - let event = match event { - Ok(ev) => ev, - Err(e) => { - tool_runtime.abort_all(); - // Propagate the underlying stream error to the caller (run_turn), which - // will apply the configured `stream_max_retries` policy. - return Err(e); - } + let add_completed = &mut |response_item: ProcessedResponseItem| { + output.push_back(future::ready(Ok(response_item)).boxed()); }; match event { @@ -2153,14 +2150,18 @@ async fn try_run_turn( Ok(Some(call)) => { let payload_preview = call.payload.log_payload().into_owned(); tracing::info!("ToolCall: {} {}", call.tool_name, payload_preview); - let index = output.len(); - output.push(ProcessedResponseItem { - item, - response: None, - }); - tool_runtime - .handle_tool_call(call, index, output.as_mut_slice()) - .await?; + + let response = tool_runtime.handle_tool_call(call); + + output.push_back( + async move { + Ok(ProcessedResponseItem { + item, + response: Some(response.await?), + }) + } + .boxed(), + ); } Ok(None) => { let response = handle_non_tool_response_item( @@ -2170,7 +2171,7 @@ async fn try_run_turn( item.clone(), ) .await?; - output.push(ProcessedResponseItem { item, response }); + add_completed(ProcessedResponseItem { item, response }); } Err(FunctionCallError::MissingLocalShellCallId) => { let msg = "LocalShellCall without call_id or id"; @@ -2187,7 +2188,7 @@ async fn try_run_turn( success: None, }, }; - output.push(ProcessedResponseItem { + add_completed(ProcessedResponseItem { item, response: Some(response), }); @@ -2200,7 +2201,7 @@ async fn try_run_turn( success: None, }, }; - output.push(ProcessedResponseItem { + add_completed(ProcessedResponseItem { item, response: Some(response), }); @@ -2231,7 +2232,7 @@ async fn try_run_turn( sess.update_token_usage_info(sub_id, turn_context.as_ref(), token_usage.as_ref()) .await; - tool_runtime.resolve_pending(output.as_mut_slice()).await?; + let processed_items: Vec = output.try_collect().await?; let unified_diff = { let mut tracker = turn_diff_tracker.lock().await; @@ -2247,7 +2248,7 @@ async fn try_run_turn( } let result = TurnRunResult { - processed_items: output, + processed_items, total_token_usage: token_usage.clone(), }; diff --git a/codex-rs/core/src/config.rs b/codex-rs/core/src/config.rs index 28ad84ba7a..d0437163fc 100644 --- a/codex-rs/core/src/config.rs +++ b/codex-rs/core/src/config.rs @@ -33,12 +33,15 @@ use codex_protocol::config_types::ReasoningEffort; use codex_protocol::config_types::ReasoningSummary; use codex_protocol::config_types::SandboxMode; use codex_protocol::config_types::Verbosity; +use codex_rmcp_client::OAuthCredentialsStoreMode; use dirs::home_dir; use serde::Deserialize; use std::collections::BTreeMap; use std::collections::HashMap; +use std::io::ErrorKind; use std::path::Path; use std::path::PathBuf; + use tempfile::NamedTempFile; use toml::Value as TomlValue; use toml_edit::Array as TomlArray; @@ -142,6 +145,15 @@ pub struct Config { /// Definition for MCP servers that Codex can reach out to for tool calls. pub mcp_servers: HashMap, + /// Preferred store for MCP OAuth credentials. + /// keyring: Use an OS-specific keyring service. + /// Credentials stored in the keyring will only be readable by Codex unless the user explicitly grants access via OS-level keyring access. + /// https://github.com/openai/codex/blob/main/codex-rs/rmcp-client/src/oauth.rs#L2 + /// file: CODEX_HOME/.credentials.json + /// This file will be readable to Codex and other applications running as the same user. + /// auto (default): keyring if available, otherwise file. + pub mcp_oauth_credentials_store_mode: OAuthCredentialsStoreMode, + /// Combined provider map (defaults merged with user-defined overrides). pub model_providers: HashMap, @@ -301,12 +313,35 @@ pub async fn load_global_mcp_servers( return Ok(BTreeMap::new()); }; + ensure_no_inline_bearer_tokens(servers_value)?; + servers_value .clone() .try_into() .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e)) } +/// We briefly allowed plain text bearer_token fields in MCP server configs. +/// We want to warn people who recently added these fields but can remove this after a few months. +fn ensure_no_inline_bearer_tokens(value: &TomlValue) -> std::io::Result<()> { + let Some(servers_table) = value.as_table() else { + return Ok(()); + }; + + for (server_name, server_value) in servers_table { + if let Some(server_table) = server_value.as_table() + && server_table.contains_key("bearer_token") + { + let message = format!( + "mcp_servers.{server_name} uses unsupported `bearer_token`; set `bearer_token_env_var`." + ); + return Err(std::io::Error::new(ErrorKind::InvalidData, message)); + } + } + + Ok(()) +} + pub fn write_global_mcp_servers( codex_home: &Path, servers: &BTreeMap, @@ -355,10 +390,13 @@ pub fn write_global_mcp_servers( entry["env"] = TomlItem::Table(env_table); } } - McpServerTransportConfig::StreamableHttp { url, bearer_token } => { + McpServerTransportConfig::StreamableHttp { + url, + bearer_token_env_var, + } => { entry["url"] = toml_edit::value(url.clone()); - if let Some(token) = bearer_token { - entry["bearer_token"] = toml_edit::value(token.clone()); + if let Some(env_var) = bearer_token_env_var { + entry["bearer_token_env_var"] = toml_edit::value(env_var.clone()); } } } @@ -694,6 +732,14 @@ pub struct ConfigToml { #[serde(default)] pub mcp_servers: HashMap, + /// Preferred backend for storing MCP OAuth credentials. + /// keyring: Use an OS-specific keyring service. + /// https://github.com/openai/codex/blob/main/codex-rs/rmcp-client/src/oauth.rs#L2 + /// file: Use a file in the Codex home directory. + /// auto (default): Use the OS-specific keyring service if available, otherwise use a file. + #[serde(default)] + pub mcp_oauth_credentials_store: Option, + /// User-defined provider entries that extend/override the built-in list. #[serde(default)] pub model_providers: HashMap, @@ -1074,6 +1120,9 @@ impl Config { user_instructions, base_instructions, mcp_servers: cfg.mcp_servers, + // The config.toml omits "_mode" because it's a config file. However, "_mode" + // is important in code to differentiate the mode from the store implementation. + mcp_oauth_credentials_store_mode: cfg.mcp_oauth_credentials_store.unwrap_or_default(), model_providers, project_doc_max_bytes: cfg.project_doc_max_bytes.unwrap_or(PROJECT_DOC_MAX_BYTES), project_doc_fallback_filenames: cfg @@ -1364,6 +1413,85 @@ exclude_slash_tmp = true ); } + #[test] + fn config_defaults_to_auto_oauth_store_mode() -> std::io::Result<()> { + let codex_home = TempDir::new()?; + let cfg = ConfigToml::default(); + + let config = Config::load_from_base_config_with_overrides( + cfg, + ConfigOverrides::default(), + codex_home.path().to_path_buf(), + )?; + + assert_eq!( + config.mcp_oauth_credentials_store_mode, + OAuthCredentialsStoreMode::Auto, + ); + + Ok(()) + } + + #[test] + fn config_honors_explicit_file_oauth_store_mode() -> std::io::Result<()> { + let codex_home = TempDir::new()?; + let cfg = ConfigToml { + mcp_oauth_credentials_store: Some(OAuthCredentialsStoreMode::File), + ..Default::default() + }; + + let config = Config::load_from_base_config_with_overrides( + cfg, + ConfigOverrides::default(), + codex_home.path().to_path_buf(), + )?; + + assert_eq!( + config.mcp_oauth_credentials_store_mode, + OAuthCredentialsStoreMode::File, + ); + + Ok(()) + } + + #[tokio::test] + async fn managed_config_overrides_oauth_store_mode() -> anyhow::Result<()> { + let codex_home = TempDir::new()?; + let managed_path = codex_home.path().join("managed_config.toml"); + let config_path = codex_home.path().join(CONFIG_TOML_FILE); + + std::fs::write(&config_path, "mcp_oauth_credentials_store = \"file\"\n")?; + std::fs::write(&managed_path, "mcp_oauth_credentials_store = \"keyring\"\n")?; + + let overrides = crate::config_loader::LoaderOverrides { + managed_config_path: Some(managed_path.clone()), + #[cfg(target_os = "macos")] + managed_preferences_base64: None, + }; + + let root_value = load_resolved_config(codex_home.path(), Vec::new(), overrides).await?; + let cfg: ConfigToml = root_value.try_into().map_err(|e| { + tracing::error!("Failed to deserialize overridden config: {e}"); + std::io::Error::new(std::io::ErrorKind::InvalidData, e) + })?; + assert_eq!( + cfg.mcp_oauth_credentials_store, + Some(OAuthCredentialsStoreMode::Keyring), + ); + + let final_config = Config::load_from_base_config_with_overrides( + cfg, + ConfigOverrides::default(), + codex_home.path().to_path_buf(), + )?; + assert_eq!( + final_config.mcp_oauth_credentials_store_mode, + OAuthCredentialsStoreMode::Keyring, + ); + + Ok(()) + } + #[tokio::test] async fn load_global_mcp_servers_returns_empty_if_missing() -> anyhow::Result<()> { let codex_home = TempDir::new()?; @@ -1471,6 +1599,31 @@ startup_timeout_ms = 2500 Ok(()) } + #[tokio::test] + async fn load_global_mcp_servers_rejects_inline_bearer_token() -> anyhow::Result<()> { + let codex_home = TempDir::new()?; + let config_path = codex_home.path().join(CONFIG_TOML_FILE); + + std::fs::write( + &config_path, + r#" +[mcp_servers.docs] +url = "https://example.com/mcp" +bearer_token = "secret" +"#, + )?; + + let err = load_global_mcp_servers(codex_home.path()) + .await + .expect_err("bearer_token entries should be rejected"); + + assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); + assert!(err.to_string().contains("bearer_token")); + assert!(err.to_string().contains("bearer_token_env_var")); + + Ok(()) + } + #[tokio::test] async fn write_global_mcp_servers_serializes_env_sorted() -> anyhow::Result<()> { let codex_home = TempDir::new()?; @@ -1534,7 +1687,7 @@ ZIG_VAR = "3" McpServerConfig { transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), - bearer_token: Some("secret-token".to_string()), + bearer_token_env_var: Some("MCP_TOKEN".to_string()), }, startup_timeout_sec: Some(Duration::from_secs(2)), tool_timeout_sec: None, @@ -1549,7 +1702,7 @@ ZIG_VAR = "3" serialized, r#"[mcp_servers.docs] url = "https://example.com/mcp" -bearer_token = "secret-token" +bearer_token_env_var = "MCP_TOKEN" startup_timeout_sec = 2.0 "# ); @@ -1557,9 +1710,12 @@ startup_timeout_sec = 2.0 let loaded = load_global_mcp_servers(codex_home.path()).await?; let docs = loaded.get("docs").expect("docs entry"); match &docs.transport { - McpServerTransportConfig::StreamableHttp { url, bearer_token } => { + McpServerTransportConfig::StreamableHttp { + url, + bearer_token_env_var, + } => { assert_eq!(url, "https://example.com/mcp"); - assert_eq!(bearer_token.as_deref(), Some("secret-token")); + assert_eq!(bearer_token_env_var.as_deref(), Some("MCP_TOKEN")); } other => panic!("unexpected transport {other:?}"), } @@ -1570,7 +1726,7 @@ startup_timeout_sec = 2.0 McpServerConfig { transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), - bearer_token: None, + bearer_token_env_var: None, }, startup_timeout_sec: None, tool_timeout_sec: None, @@ -1589,9 +1745,12 @@ url = "https://example.com/mcp" let loaded = load_global_mcp_servers(codex_home.path()).await?; let docs = loaded.get("docs").expect("docs entry"); match &docs.transport { - McpServerTransportConfig::StreamableHttp { url, bearer_token } => { + McpServerTransportConfig::StreamableHttp { + url, + bearer_token_env_var, + } => { assert_eq!(url, "https://example.com/mcp"); - assert!(bearer_token.is_none()); + assert!(bearer_token_env_var.is_none()); } other => panic!("unexpected transport {other:?}"), } @@ -1896,6 +2055,7 @@ model_verbosity = "high" notify: None, cwd: fixture.cwd(), mcp_servers: HashMap::new(), + mcp_oauth_credentials_store_mode: Default::default(), model_providers: fixture.model_provider_map.clone(), project_doc_max_bytes: PROJECT_DOC_MAX_BYTES, project_doc_fallback_filenames: Vec::new(), @@ -1958,6 +2118,7 @@ model_verbosity = "high" notify: None, cwd: fixture.cwd(), mcp_servers: HashMap::new(), + mcp_oauth_credentials_store_mode: Default::default(), model_providers: fixture.model_provider_map.clone(), project_doc_max_bytes: PROJECT_DOC_MAX_BYTES, project_doc_fallback_filenames: Vec::new(), @@ -2035,6 +2196,7 @@ model_verbosity = "high" notify: None, cwd: fixture.cwd(), mcp_servers: HashMap::new(), + mcp_oauth_credentials_store_mode: Default::default(), model_providers: fixture.model_provider_map.clone(), project_doc_max_bytes: PROJECT_DOC_MAX_BYTES, project_doc_fallback_filenames: Vec::new(), @@ -2098,6 +2260,7 @@ model_verbosity = "high" notify: None, cwd: fixture.cwd(), mcp_servers: HashMap::new(), + mcp_oauth_credentials_store_mode: Default::default(), model_providers: fixture.model_provider_map.clone(), project_doc_max_bytes: PROJECT_DOC_MAX_BYTES, project_doc_fallback_filenames: Vec::new(), diff --git a/codex-rs/core/src/config_types.rs b/codex-rs/core/src/config_types.rs index 1b8f3ac067..8940ac5c08 100644 --- a/codex-rs/core/src/config_types.rs +++ b/codex-rs/core/src/config_types.rs @@ -48,6 +48,7 @@ impl<'de> Deserialize<'de> for McpServerConfig { url: Option, bearer_token: Option, + bearer_token_env_var: Option, #[serde(default)] startup_timeout_sec: Option, @@ -86,11 +87,15 @@ impl<'de> Deserialize<'de> for McpServerConfig { args, env, url, - bearer_token, + bearer_token_env_var, .. } => { throw_if_set("stdio", "url", url.as_ref())?; - throw_if_set("stdio", "bearer_token", bearer_token.as_ref())?; + throw_if_set( + "stdio", + "bearer_token_env_var", + bearer_token_env_var.as_ref(), + )?; McpServerTransportConfig::Stdio { command, args: args.unwrap_or_default(), @@ -100,6 +105,7 @@ impl<'de> Deserialize<'de> for McpServerConfig { RawMcpServerConfig { url: Some(url), bearer_token, + bearer_token_env_var, command, args, env, @@ -108,7 +114,11 @@ impl<'de> Deserialize<'de> for McpServerConfig { throw_if_set("streamable_http", "command", command.as_ref())?; throw_if_set("streamable_http", "args", args.as_ref())?; throw_if_set("streamable_http", "env", env.as_ref())?; - McpServerTransportConfig::StreamableHttp { url, bearer_token } + throw_if_set("streamable_http", "bearer_token", bearer_token.as_ref())?; + McpServerTransportConfig::StreamableHttp { + url, + bearer_token_env_var, + } } _ => return Err(SerdeError::custom("invalid transport")), }; @@ -135,11 +145,11 @@ pub enum McpServerTransportConfig { /// https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#streamable-http StreamableHttp { url: String, - /// A plain text bearer token to use for authentication. - /// This bearer token will be included in the HTTP request header as an `Authorization: Bearer ` header. - /// This should be used with caution because it lives on disk in clear text. + /// Name of the environment variable to read for an HTTP bearer token. + /// When set, requests will include the token via `Authorization: Bearer `. + /// The actual secret value must be provided via the environment. #[serde(default, skip_serializing_if = "Option::is_none")] - bearer_token: Option, + bearer_token_env_var: Option, }, } @@ -506,17 +516,17 @@ mod tests { cfg.transport, McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), - bearer_token: None + bearer_token_env_var: None } ); } #[test] - fn deserialize_streamable_http_server_config_with_bearer_token() { + fn deserialize_streamable_http_server_config_with_env_var() { let cfg: McpServerConfig = toml::from_str( r#" url = "https://example.com/mcp" - bearer_token = "secret" + bearer_token_env_var = "GITHUB_TOKEN" "#, ) .expect("should deserialize http config"); @@ -525,7 +535,7 @@ mod tests { cfg.transport, McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), - bearer_token: Some("secret".to_string()) + bearer_token_env_var: Some("GITHUB_TOKEN".to_string()) } ); } @@ -553,13 +563,18 @@ mod tests { } #[test] - fn deserialize_rejects_bearer_token_for_stdio_transport() { - toml::from_str::( + fn deserialize_rejects_inline_bearer_token_field() { + let err = toml::from_str::( r#" - command = "echo" + url = "https://example.com" bearer_token = "secret" "#, ) - .expect_err("should reject bearer token for stdio transport"); + .expect_err("should reject bearer_token field"); + + assert!( + err.to_string().contains("bearer_token is not supported"), + "unexpected error: {err}" + ); } } diff --git a/codex-rs/core/src/default_client.rs b/codex-rs/core/src/default_client.rs index 5a6ea8d826..42690e04cc 100644 --- a/codex-rs/core/src/default_client.rs +++ b/codex-rs/core/src/default_client.rs @@ -20,7 +20,7 @@ use std::sync::OnceLock; /// The full user agent string is returned from the mcp initialize response. /// Parenthesis will be added by Codex. This should only specify what goes inside of the parenthesis. pub static USER_AGENT_SUFFIX: LazyLock>> = LazyLock::new(|| Mutex::new(None)); - +pub const DEFAULT_ORIGINATOR: &str = "codex_cli_rs"; pub const CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR: &str = "CODEX_INTERNAL_ORIGINATOR_OVERRIDE"; #[derive(Debug, Clone)] pub struct Originator { @@ -35,10 +35,11 @@ pub enum SetOriginatorError { AlreadyInitialized, } -fn init_originator_from_env() -> Originator { - let default = "codex_cli_rs"; +fn get_originator_value(provided: Option) -> Originator { let value = std::env::var(CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR) - .unwrap_or_else(|_| default.to_string()); + .ok() + .or(provided) + .unwrap_or(DEFAULT_ORIGINATOR.to_string()); match HeaderValue::from_str(&value) { Ok(header_value) => Originator { @@ -48,31 +49,22 @@ fn init_originator_from_env() -> Originator { Err(e) => { tracing::error!("Unable to turn originator override {value} into header value: {e}"); Originator { - value: default.to_string(), - header_value: HeaderValue::from_static(default), + value: DEFAULT_ORIGINATOR.to_string(), + header_value: HeaderValue::from_static(DEFAULT_ORIGINATOR), } } } } -fn build_originator(value: String) -> Result { - let header_value = - HeaderValue::from_str(&value).map_err(|_| SetOriginatorError::InvalidHeaderValue)?; - Ok(Originator { - value, - header_value, - }) -} - -pub fn set_default_originator(value: &str) -> Result<(), SetOriginatorError> { - let originator = build_originator(value.to_string())?; +pub fn set_default_originator(value: String) -> Result<(), SetOriginatorError> { + let originator = get_originator_value(Some(value)); ORIGINATOR .set(originator) .map_err(|_| SetOriginatorError::AlreadyInitialized) } pub fn originator() -> &'static Originator { - ORIGINATOR.get_or_init(init_originator_from_env) + ORIGINATOR.get_or_init(|| get_originator_value(None)) } pub fn get_codex_user_agent() -> String { diff --git a/codex-rs/core/src/mcp_connection_manager.rs b/codex-rs/core/src/mcp_connection_manager.rs index 76738a03ae..0c2fdc6301 100644 --- a/codex-rs/core/src/mcp_connection_manager.rs +++ b/codex-rs/core/src/mcp_connection_manager.rs @@ -8,6 +8,7 @@ use std::collections::HashMap; use std::collections::HashSet; +use std::env; use std::ffi::OsString; use std::sync::Arc; use std::time::Duration; @@ -16,6 +17,7 @@ use anyhow::Context; use anyhow::Result; use anyhow::anyhow; use codex_mcp_client::McpClient; +use codex_rmcp_client::OAuthCredentialsStoreMode; use codex_rmcp_client::RmcpClient; use mcp_types::ClientCapabilities; use mcp_types::Implementation; @@ -125,9 +127,11 @@ impl McpClientAdapter { bearer_token: Option, params: mcp_types::InitializeRequestParams, startup_timeout: Duration, + store_mode: OAuthCredentialsStoreMode, ) -> Result { let client = Arc::new( - RmcpClient::new_streamable_http_client(&server_name, &url, bearer_token).await?, + RmcpClient::new_streamable_http_client(&server_name, &url, bearer_token, store_mode) + .await?, ); client.initialize(params, Some(startup_timeout)).await?; Ok(McpClientAdapter::Rmcp(client)) @@ -182,6 +186,7 @@ impl McpConnectionManager { pub async fn new( mcp_servers: HashMap, use_rmcp_client: bool, + store_mode: OAuthCredentialsStoreMode, ) -> Result<(Self, ClientStartErrors)> { // Early exit if no servers are configured. if mcp_servers.is_empty() { @@ -205,6 +210,14 @@ impl McpConnectionManager { let startup_timeout = cfg.startup_timeout_sec.unwrap_or(DEFAULT_STARTUP_TIMEOUT); let tool_timeout = cfg.tool_timeout_sec.unwrap_or(DEFAULT_TOOL_TIMEOUT); + let resolved_bearer_token = match &cfg.transport { + McpServerTransportConfig::StreamableHttp { + bearer_token_env_var, + .. + } => resolve_bearer_token(&server_name, bearer_token_env_var.as_deref()), + _ => Ok(None), + }; + join_set.spawn(async move { let McpServerConfig { transport, .. } = cfg; let params = mcp_types::InitializeRequestParams { @@ -242,13 +255,14 @@ impl McpConnectionManager { ) .await } - McpServerTransportConfig::StreamableHttp { url, bearer_token } => { + McpServerTransportConfig::StreamableHttp { url, .. } => { McpClientAdapter::new_streamable_http_client( server_name.clone(), url, - bearer_token, + resolved_bearer_token.unwrap_or_default(), params, startup_timeout, + store_mode, ) .await } @@ -336,6 +350,33 @@ impl McpConnectionManager { } } +fn resolve_bearer_token( + server_name: &str, + bearer_token_env_var: Option<&str>, +) -> Result> { + let Some(env_var) = bearer_token_env_var else { + return Ok(None); + }; + + match env::var(env_var) { + Ok(value) => { + if value.is_empty() { + Err(anyhow!( + "Environment variable {env_var} for MCP server '{server_name}' is empty" + )) + } else { + Ok(Some(value)) + } + } + Err(env::VarError::NotPresent) => Err(anyhow!( + "Environment variable {env_var} for MCP server '{server_name}' is not set" + )), + Err(env::VarError::NotUnicode(_)) => Err(anyhow!( + "Environment variable {env_var} for MCP server '{server_name}' contains invalid Unicode" + )), + } +} + /// Query every server for its available tools and return a single map that /// contains **all** tools. Each key is the fully-qualified name for the tool. async fn list_all_tools(clients: &HashMap) -> Result> { diff --git a/codex-rs/core/src/model_family.rs b/codex-rs/core/src/model_family.rs index 96faaf47ad..388650e32c 100644 --- a/codex-rs/core/src/model_family.rs +++ b/codex-rs/core/src/model_family.rs @@ -120,8 +120,11 @@ pub fn find_family_for_model(mut slug: &str) -> Option { reasoning_summary_format: ReasoningSummaryFormat::Experimental, base_instructions: GPT_5_CODEX_INSTRUCTIONS.to_string(), experimental_supported_tools: vec![ + "grep_files".to_string(), + "list_dir".to_string(), "read_file".to_string(), "test_sync_tool".to_string(), + "test_sync_tool".to_string(), "subsession".to_string(), ], supports_parallel_tool_calls: true, @@ -135,7 +138,12 @@ pub fn find_family_for_model(mut slug: &str) -> Option { reasoning_summary_format: ReasoningSummaryFormat::Experimental, base_instructions: GPT_5_CODEX_INSTRUCTIONS.to_string(), apply_patch_tool_type: Some(ApplyPatchToolType::Freeform), - experimental_supported_tools: vec!["read_file".to_string(), "subsession".to_string()], + experimental_supported_tools: vec![ + "grep_files".to_string(), + "list_dir".to_string(), + "read_file".to_string(), + "subsession".to_string(), + ], supports_parallel_tool_calls: true, ) diff --git a/codex-rs/core/src/tools/handlers/grep_files.rs b/codex-rs/core/src/tools/handlers/grep_files.rs new file mode 100644 index 0000000000..de3cd3411c --- /dev/null +++ b/codex-rs/core/src/tools/handlers/grep_files.rs @@ -0,0 +1,272 @@ +use std::path::Path; +use std::time::Duration; + +use async_trait::async_trait; +use serde::Deserialize; +use tokio::process::Command; +use tokio::time::timeout; + +use crate::function_tool::FunctionCallError; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolOutput; +use crate::tools::context::ToolPayload; +use crate::tools::registry::ToolHandler; +use crate::tools::registry::ToolKind; + +pub struct GrepFilesHandler; + +const DEFAULT_LIMIT: usize = 100; +const MAX_LIMIT: usize = 2000; +const COMMAND_TIMEOUT: Duration = Duration::from_secs(30); + +fn default_limit() -> usize { + DEFAULT_LIMIT +} + +#[derive(Deserialize)] +struct GrepFilesArgs { + pattern: String, + #[serde(default)] + include: Option, + #[serde(default)] + path: Option, + #[serde(default = "default_limit")] + limit: usize, +} + +#[async_trait] +impl ToolHandler for GrepFilesHandler { + fn kind(&self) -> ToolKind { + ToolKind::Function + } + + async fn handle(&self, invocation: ToolInvocation) -> Result { + let ToolInvocation { payload, turn, .. } = invocation; + + let arguments = match payload { + ToolPayload::Function { arguments } => arguments, + _ => { + return Err(FunctionCallError::RespondToModel( + "grep_files handler received unsupported payload".to_string(), + )); + } + }; + + let args: GrepFilesArgs = serde_json::from_str(&arguments).map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to parse function arguments: {err:?}" + )) + })?; + + let pattern = args.pattern.trim(); + if pattern.is_empty() { + return Err(FunctionCallError::RespondToModel( + "pattern must not be empty".to_string(), + )); + } + + if args.limit == 0 { + return Err(FunctionCallError::RespondToModel( + "limit must be greater than zero".to_string(), + )); + } + + let limit = args.limit.min(MAX_LIMIT); + let search_path = turn.resolve_path(args.path.clone()); + + verify_path_exists(&search_path).await?; + + let include = args.include.as_deref().map(str::trim).and_then(|val| { + if val.is_empty() { + None + } else { + Some(val.to_string()) + } + }); + + let search_results = + run_rg_search(pattern, include.as_deref(), &search_path, limit, &turn.cwd).await?; + + if search_results.is_empty() { + Ok(ToolOutput::Function { + content: "No matches found.".to_string(), + success: Some(false), + }) + } else { + Ok(ToolOutput::Function { + content: search_results.join("\n"), + success: Some(true), + }) + } + } +} + +async fn verify_path_exists(path: &Path) -> Result<(), FunctionCallError> { + tokio::fs::metadata(path).await.map_err(|err| { + FunctionCallError::RespondToModel(format!("unable to access `{}`: {err}", path.display())) + })?; + Ok(()) +} + +async fn run_rg_search( + pattern: &str, + include: Option<&str>, + search_path: &Path, + limit: usize, + cwd: &Path, +) -> Result, FunctionCallError> { + let mut command = Command::new("rg"); + command + .current_dir(cwd) + .arg("--files-with-matches") + .arg("--sortr=modified") + .arg("--regexp") + .arg(pattern) + .arg("--no-messages"); + + if let Some(glob) = include { + command.arg("--glob").arg(glob); + } + + command.arg("--").arg(search_path); + + let output = timeout(COMMAND_TIMEOUT, command.output()) + .await + .map_err(|_| { + FunctionCallError::RespondToModel("rg timed out after 30 seconds".to_string()) + })? + .map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to launch rg: {err}. Ensure ripgrep is installed and on PATH." + )) + })?; + + match output.status.code() { + Some(0) => Ok(parse_results(&output.stdout, limit)), + Some(1) => Ok(Vec::new()), + _ => { + let stderr = String::from_utf8_lossy(&output.stderr); + Err(FunctionCallError::RespondToModel(format!( + "rg failed: {stderr}" + ))) + } + } +} + +fn parse_results(stdout: &[u8], limit: usize) -> Vec { + let mut results = Vec::new(); + for line in stdout.split(|byte| *byte == b'\n') { + if line.is_empty() { + continue; + } + if let Ok(text) = std::str::from_utf8(line) { + if text.is_empty() { + continue; + } + results.push(text.to_string()); + if results.len() == limit { + break; + } + } + } + results +} + +#[cfg(test)] +mod tests { + use super::*; + use std::process::Command as StdCommand; + use tempfile::tempdir; + + #[test] + fn parses_basic_results() { + let stdout = b"/tmp/file_a.rs\n/tmp/file_b.rs\n"; + let parsed = parse_results(stdout, 10); + assert_eq!( + parsed, + vec!["/tmp/file_a.rs".to_string(), "/tmp/file_b.rs".to_string()] + ); + } + + #[test] + fn parse_truncates_after_limit() { + let stdout = b"/tmp/file_a.rs\n/tmp/file_b.rs\n/tmp/file_c.rs\n"; + let parsed = parse_results(stdout, 2); + assert_eq!( + parsed, + vec!["/tmp/file_a.rs".to_string(), "/tmp/file_b.rs".to_string()] + ); + } + + #[tokio::test] + async fn run_search_returns_results() -> anyhow::Result<()> { + if !rg_available() { + return Ok(()); + } + let temp = tempdir().expect("create temp dir"); + let dir = temp.path(); + std::fs::write(dir.join("match_one.txt"), "alpha beta gamma").unwrap(); + std::fs::write(dir.join("match_two.txt"), "alpha delta").unwrap(); + std::fs::write(dir.join("other.txt"), "omega").unwrap(); + + let results = run_rg_search("alpha", None, dir, 10, dir).await?; + assert_eq!(results.len(), 2); + assert!(results.iter().any(|path| path.ends_with("match_one.txt"))); + assert!(results.iter().any(|path| path.ends_with("match_two.txt"))); + Ok(()) + } + + #[tokio::test] + async fn run_search_with_glob_filter() -> anyhow::Result<()> { + if !rg_available() { + return Ok(()); + } + let temp = tempdir().expect("create temp dir"); + let dir = temp.path(); + std::fs::write(dir.join("match_one.rs"), "alpha beta gamma").unwrap(); + std::fs::write(dir.join("match_two.txt"), "alpha delta").unwrap(); + + let results = run_rg_search("alpha", Some("*.rs"), dir, 10, dir).await?; + assert_eq!(results.len(), 1); + assert!(results.iter().all(|path| path.ends_with("match_one.rs"))); + Ok(()) + } + + #[tokio::test] + async fn run_search_respects_limit() -> anyhow::Result<()> { + if !rg_available() { + return Ok(()); + } + let temp = tempdir().expect("create temp dir"); + let dir = temp.path(); + std::fs::write(dir.join("one.txt"), "alpha one").unwrap(); + std::fs::write(dir.join("two.txt"), "alpha two").unwrap(); + std::fs::write(dir.join("three.txt"), "alpha three").unwrap(); + + let results = run_rg_search("alpha", None, dir, 2, dir).await?; + assert_eq!(results.len(), 2); + Ok(()) + } + + #[tokio::test] + async fn run_search_handles_no_matches() -> anyhow::Result<()> { + if !rg_available() { + return Ok(()); + } + let temp = tempdir().expect("create temp dir"); + let dir = temp.path(); + std::fs::write(dir.join("one.txt"), "omega").unwrap(); + + let results = run_rg_search("alpha", None, dir, 5, dir).await?; + assert!(results.is_empty()); + Ok(()) + } + + fn rg_available() -> bool { + StdCommand::new("rg") + .arg("--version") + .output() + .map(|output| output.status.success()) + .unwrap_or(false) + } +} diff --git a/codex-rs/core/src/tools/handlers/list_dir.rs b/codex-rs/core/src/tools/handlers/list_dir.rs new file mode 100644 index 0000000000..bcea4a756d --- /dev/null +++ b/codex-rs/core/src/tools/handlers/list_dir.rs @@ -0,0 +1,476 @@ +use std::collections::VecDeque; +use std::ffi::OsStr; +use std::fs::FileType; +use std::path::Path; +use std::path::PathBuf; + +use async_trait::async_trait; +use codex_utils_string::take_bytes_at_char_boundary; +use serde::Deserialize; +use tokio::fs; + +use crate::function_tool::FunctionCallError; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolOutput; +use crate::tools::context::ToolPayload; +use crate::tools::registry::ToolHandler; +use crate::tools::registry::ToolKind; + +pub struct ListDirHandler; + +const MAX_ENTRY_LENGTH: usize = 500; +const INDENTATION_SPACES: usize = 2; + +fn default_offset() -> usize { + 1 +} + +fn default_limit() -> usize { + 25 +} + +fn default_depth() -> usize { + 2 +} + +#[derive(Deserialize)] +struct ListDirArgs { + dir_path: String, + #[serde(default = "default_offset")] + offset: usize, + #[serde(default = "default_limit")] + limit: usize, + #[serde(default = "default_depth")] + depth: usize, +} + +#[async_trait] +impl ToolHandler for ListDirHandler { + fn kind(&self) -> ToolKind { + ToolKind::Function + } + + async fn handle(&self, invocation: ToolInvocation) -> Result { + let ToolInvocation { payload, .. } = invocation; + + let arguments = match payload { + ToolPayload::Function { arguments } => arguments, + _ => { + return Err(FunctionCallError::RespondToModel( + "list_dir handler received unsupported payload".to_string(), + )); + } + }; + + let args: ListDirArgs = serde_json::from_str(&arguments).map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to parse function arguments: {err:?}" + )) + })?; + + let ListDirArgs { + dir_path, + offset, + limit, + depth, + } = args; + + if offset == 0 { + return Err(FunctionCallError::RespondToModel( + "offset must be a 1-indexed entry number".to_string(), + )); + } + + if limit == 0 { + return Err(FunctionCallError::RespondToModel( + "limit must be greater than zero".to_string(), + )); + } + + if depth == 0 { + return Err(FunctionCallError::RespondToModel( + "depth must be greater than zero".to_string(), + )); + } + + let path = PathBuf::from(&dir_path); + if !path.is_absolute() { + return Err(FunctionCallError::RespondToModel( + "dir_path must be an absolute path".to_string(), + )); + } + + let entries = list_dir_slice(&path, offset, limit, depth).await?; + let mut output = Vec::with_capacity(entries.len() + 1); + output.push(format!("Absolute path: {}", path.display())); + output.extend(entries); + Ok(ToolOutput::Function { + content: output.join("\n"), + success: Some(true), + }) + } +} + +async fn list_dir_slice( + path: &Path, + offset: usize, + limit: usize, + depth: usize, +) -> Result, FunctionCallError> { + let mut entries = Vec::new(); + collect_entries(path, Path::new(""), depth, &mut entries).await?; + + if entries.is_empty() { + return Ok(Vec::new()); + } + + let start_index = offset - 1; + if start_index >= entries.len() { + return Err(FunctionCallError::RespondToModel( + "offset exceeds directory entry count".to_string(), + )); + } + + let remaining_entries = entries.len() - start_index; + let capped_limit = limit.min(remaining_entries); + let end_index = start_index + capped_limit; + let mut selected_entries = entries[start_index..end_index].to_vec(); + selected_entries.sort_unstable_by(|a, b| a.name.cmp(&b.name)); + let mut formatted = Vec::with_capacity(selected_entries.len()); + + for entry in &selected_entries { + formatted.push(format_entry_line(entry)); + } + + if end_index < entries.len() { + formatted.push(format!("More than {capped_limit} entries found")); + } + + Ok(formatted) +} + +async fn collect_entries( + dir_path: &Path, + relative_prefix: &Path, + depth: usize, + entries: &mut Vec, +) -> Result<(), FunctionCallError> { + let mut queue = VecDeque::new(); + queue.push_back((dir_path.to_path_buf(), relative_prefix.to_path_buf(), depth)); + + while let Some((current_dir, prefix, remaining_depth)) = queue.pop_front() { + let mut read_dir = fs::read_dir(¤t_dir).await.map_err(|err| { + FunctionCallError::RespondToModel(format!("failed to read directory: {err}")) + })?; + + let mut dir_entries = Vec::new(); + + while let Some(entry) = read_dir.next_entry().await.map_err(|err| { + FunctionCallError::RespondToModel(format!("failed to read directory: {err}")) + })? { + let file_type = entry.file_type().await.map_err(|err| { + FunctionCallError::RespondToModel(format!("failed to inspect entry: {err}")) + })?; + + let file_name = entry.file_name(); + let relative_path = if prefix.as_os_str().is_empty() { + PathBuf::from(&file_name) + } else { + prefix.join(&file_name) + }; + + let display_name = format_entry_component(&file_name); + let display_depth = prefix.components().count(); + let sort_key = format_entry_name(&relative_path); + let kind = DirEntryKind::from(&file_type); + dir_entries.push(( + entry.path(), + relative_path, + kind, + DirEntry { + name: sort_key, + display_name, + depth: display_depth, + kind, + }, + )); + } + + dir_entries.sort_unstable_by(|a, b| a.3.name.cmp(&b.3.name)); + + for (entry_path, relative_path, kind, dir_entry) in dir_entries { + if kind == DirEntryKind::Directory && remaining_depth > 1 { + queue.push_back((entry_path, relative_path, remaining_depth - 1)); + } + entries.push(dir_entry); + } + } + + Ok(()) +} + +fn format_entry_name(path: &Path) -> String { + let normalized = path.to_string_lossy().replace("\\", "/"); + if normalized.len() > MAX_ENTRY_LENGTH { + take_bytes_at_char_boundary(&normalized, MAX_ENTRY_LENGTH).to_string() + } else { + normalized + } +} + +fn format_entry_component(name: &OsStr) -> String { + let normalized = name.to_string_lossy(); + if normalized.len() > MAX_ENTRY_LENGTH { + take_bytes_at_char_boundary(&normalized, MAX_ENTRY_LENGTH).to_string() + } else { + normalized.to_string() + } +} + +fn format_entry_line(entry: &DirEntry) -> String { + let indent = " ".repeat(entry.depth * INDENTATION_SPACES); + let mut name = entry.display_name.clone(); + match entry.kind { + DirEntryKind::Directory => name.push('/'), + DirEntryKind::Symlink => name.push('@'), + DirEntryKind::Other => name.push('?'), + DirEntryKind::File => {} + } + format!("{indent}{name}") +} + +#[derive(Clone)] +struct DirEntry { + name: String, + display_name: String, + depth: usize, + kind: DirEntryKind, +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum DirEntryKind { + Directory, + File, + Symlink, + Other, +} + +impl From<&FileType> for DirEntryKind { + fn from(file_type: &FileType) -> Self { + if file_type.is_symlink() { + DirEntryKind::Symlink + } else if file_type.is_dir() { + DirEntryKind::Directory + } else if file_type.is_file() { + DirEntryKind::File + } else { + DirEntryKind::Other + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::tempdir; + + #[tokio::test] + async fn lists_directory_entries() { + let temp = tempdir().expect("create tempdir"); + let dir_path = temp.path(); + + let sub_dir = dir_path.join("nested"); + tokio::fs::create_dir(&sub_dir) + .await + .expect("create sub dir"); + + let deeper_dir = sub_dir.join("deeper"); + tokio::fs::create_dir(&deeper_dir) + .await + .expect("create deeper dir"); + + tokio::fs::write(dir_path.join("entry.txt"), b"content") + .await + .expect("write file"); + tokio::fs::write(sub_dir.join("child.txt"), b"child") + .await + .expect("write child"); + tokio::fs::write(deeper_dir.join("grandchild.txt"), b"grandchild") + .await + .expect("write grandchild"); + + #[cfg(unix)] + { + use std::os::unix::fs::symlink; + let link_path = dir_path.join("link"); + symlink(dir_path.join("entry.txt"), &link_path).expect("create symlink"); + } + + let entries = list_dir_slice(dir_path, 1, 20, 3) + .await + .expect("list directory"); + + #[cfg(unix)] + let expected = vec![ + "entry.txt".to_string(), + "link@".to_string(), + "nested/".to_string(), + " child.txt".to_string(), + " deeper/".to_string(), + " grandchild.txt".to_string(), + ]; + + #[cfg(not(unix))] + let expected = vec![ + "entry.txt".to_string(), + "nested/".to_string(), + " child.txt".to_string(), + " deeper/".to_string(), + " grandchild.txt".to_string(), + ]; + + assert_eq!(entries, expected); + } + + #[tokio::test] + async fn errors_when_offset_exceeds_entries() { + let temp = tempdir().expect("create tempdir"); + let dir_path = temp.path(); + tokio::fs::create_dir(dir_path.join("nested")) + .await + .expect("create sub dir"); + + let err = list_dir_slice(dir_path, 10, 1, 2) + .await + .expect_err("offset exceeds entries"); + assert_eq!( + err, + FunctionCallError::RespondToModel("offset exceeds directory entry count".to_string()) + ); + } + + #[tokio::test] + async fn respects_depth_parameter() { + let temp = tempdir().expect("create tempdir"); + let dir_path = temp.path(); + let nested = dir_path.join("nested"); + let deeper = nested.join("deeper"); + tokio::fs::create_dir(&nested).await.expect("create nested"); + tokio::fs::create_dir(&deeper).await.expect("create deeper"); + tokio::fs::write(dir_path.join("root.txt"), b"root") + .await + .expect("write root"); + tokio::fs::write(nested.join("child.txt"), b"child") + .await + .expect("write nested"); + tokio::fs::write(deeper.join("grandchild.txt"), b"deep") + .await + .expect("write deeper"); + + let entries_depth_one = list_dir_slice(dir_path, 1, 10, 1) + .await + .expect("list depth 1"); + assert_eq!( + entries_depth_one, + vec!["nested/".to_string(), "root.txt".to_string(),] + ); + + let entries_depth_two = list_dir_slice(dir_path, 1, 20, 2) + .await + .expect("list depth 2"); + assert_eq!( + entries_depth_two, + vec![ + "nested/".to_string(), + " child.txt".to_string(), + " deeper/".to_string(), + "root.txt".to_string(), + ] + ); + + let entries_depth_three = list_dir_slice(dir_path, 1, 30, 3) + .await + .expect("list depth 3"); + assert_eq!( + entries_depth_three, + vec![ + "nested/".to_string(), + " child.txt".to_string(), + " deeper/".to_string(), + " grandchild.txt".to_string(), + "root.txt".to_string(), + ] + ); + } + + #[tokio::test] + async fn handles_large_limit_without_overflow() { + let temp = tempdir().expect("create tempdir"); + let dir_path = temp.path(); + tokio::fs::write(dir_path.join("alpha.txt"), b"alpha") + .await + .expect("write alpha"); + tokio::fs::write(dir_path.join("beta.txt"), b"beta") + .await + .expect("write beta"); + tokio::fs::write(dir_path.join("gamma.txt"), b"gamma") + .await + .expect("write gamma"); + + let entries = list_dir_slice(dir_path, 2, usize::MAX, 1) + .await + .expect("list without overflow"); + assert_eq!( + entries, + vec!["beta.txt".to_string(), "gamma.txt".to_string(),] + ); + } + + #[tokio::test] + async fn indicates_truncated_results() { + let temp = tempdir().expect("create tempdir"); + let dir_path = temp.path(); + + for idx in 0..40 { + let file = dir_path.join(format!("file_{idx:02}.txt")); + tokio::fs::write(file, b"content") + .await + .expect("write file"); + } + + let entries = list_dir_slice(dir_path, 1, 25, 1) + .await + .expect("list directory"); + assert_eq!(entries.len(), 26); + assert_eq!( + entries.last(), + Some(&"More than 25 entries found".to_string()) + ); + } + + #[tokio::test] + async fn bfs_truncation() -> anyhow::Result<()> { + let temp = tempdir()?; + let dir_path = temp.path(); + let nested = dir_path.join("nested"); + let deeper = nested.join("deeper"); + tokio::fs::create_dir(&nested).await?; + tokio::fs::create_dir(&deeper).await?; + tokio::fs::write(dir_path.join("root.txt"), b"root").await?; + tokio::fs::write(nested.join("child.txt"), b"child").await?; + tokio::fs::write(deeper.join("grandchild.txt"), b"deep").await?; + + let entries_depth_three = list_dir_slice(dir_path, 1, 3, 3).await?; + assert_eq!( + entries_depth_three, + vec![ + "nested/".to_string(), + " child.txt".to_string(), + "root.txt".to_string(), + "More than 3 entries found".to_string() + ] + ); + + Ok(()) + } +} diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index 2dd12c3c1c..db7aff727c 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -1,5 +1,7 @@ pub mod apply_patch; mod exec_stream; +mod grep_files; +mod list_dir; mod mcp; mod plan; mod read_file; @@ -13,6 +15,8 @@ pub use plan::PLAN_TOOL; pub use apply_patch::ApplyPatchHandler; pub use exec_stream::ExecStreamHandler; +pub use grep_files::GrepFilesHandler; +pub use list_dir::ListDirHandler; pub use mcp::McpHandler; pub use plan::PlanHandler; pub use read_file::ReadFileHandler; diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index ff4104d0c1..26dfed8ecc 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -1,6 +1,8 @@ use std::sync::Arc; -use tokio::task::JoinHandle; +use tokio::sync::RwLock; +use tokio_util::either::Either; +use tokio_util::task::AbortOnDropHandle; use crate::codex::Session; use crate::codex::TurnContext; @@ -11,20 +13,13 @@ use crate::tools::router::ToolCall; use crate::tools::router::ToolRouter; use codex_protocol::models::ResponseInputItem; -use crate::codex::ProcessedResponseItem; - -struct PendingToolCall { - index: usize, - handle: JoinHandle>, -} - pub(crate) struct ToolCallRuntime { router: Arc, session: Arc, turn_context: Arc, tracker: SharedTurnDiffTracker, sub_id: String, - pending_calls: Vec, + parallel_execution: Arc>, } impl ToolCallRuntime { @@ -41,97 +36,45 @@ impl ToolCallRuntime { turn_context, tracker, sub_id, - pending_calls: Vec::new(), + parallel_execution: Arc::new(RwLock::new(())), } } - pub(crate) async fn handle_tool_call( - &mut self, + pub(crate) fn handle_tool_call( + &self, call: ToolCall, - output_index: usize, - output: &mut [ProcessedResponseItem], - ) -> Result<(), CodexErr> { + ) -> impl std::future::Future> { let supports_parallel = self.router.tool_supports_parallel(&call.tool_name); - if supports_parallel { - self.spawn_parallel(call, output_index); - } else { - self.resolve_pending(output).await?; - let response = self.dispatch_serial(call).await?; - let slot = output.get_mut(output_index).ok_or_else(|| { - CodexErr::Fatal(format!("tool output index {output_index} out of bounds")) - })?; - slot.response = Some(response); - } - Ok(()) - } - - pub(crate) fn abort_all(&mut self) { - while let Some(pending) = self.pending_calls.pop() { - pending.handle.abort(); - } - } - - pub(crate) async fn resolve_pending( - &mut self, - output: &mut [ProcessedResponseItem], - ) -> Result<(), CodexErr> { - while let Some(PendingToolCall { index, handle }) = self.pending_calls.pop() { - match handle.await { - Ok(Ok(response)) => { - if let Some(slot) = output.get_mut(index) { - slot.response = Some(response); - } - } - Ok(Err(FunctionCallError::Fatal(message))) => { - self.abort_all(); - return Err(CodexErr::Fatal(message)); - } - Ok(Err(other)) => { - self.abort_all(); - return Err(CodexErr::Fatal(other.to_string())); - } - Err(join_err) => { - self.abort_all(); - return Err(CodexErr::Fatal(format!( - "tool task failed to join: {join_err}" - ))); - } - } - } - - Ok(()) - } - - fn spawn_parallel(&mut self, call: ToolCall, index: usize) { let router = Arc::clone(&self.router); let session = Arc::clone(&self.session); let turn = Arc::clone(&self.turn_context); let tracker = Arc::clone(&self.tracker); let sub_id = self.sub_id.clone(); - let handle = tokio::spawn(async move { - router - .dispatch_tool_call(session, turn, tracker, sub_id, call) - .await - }); - self.pending_calls.push(PendingToolCall { index, handle }); - } + let lock = Arc::clone(&self.parallel_execution); - async fn dispatch_serial(&self, call: ToolCall) -> Result { - match self - .router - .dispatch_tool_call( - Arc::clone(&self.session), - Arc::clone(&self.turn_context), - Arc::clone(&self.tracker), - self.sub_id.clone(), - call, - ) - .await - { - Ok(response) => Ok(response), - Err(FunctionCallError::Fatal(message)) => Err(CodexErr::Fatal(message)), - Err(other) => Err(CodexErr::Fatal(other.to_string())), + let handle: AbortOnDropHandle> = + AbortOnDropHandle::new(tokio::spawn(async move { + let _guard = if supports_parallel { + Either::Left(lock.read().await) + } else { + Either::Right(lock.write().await) + }; + + router + .dispatch_tool_call(session, turn, tracker, sub_id, call) + .await + })); + + async move { + match handle.await { + Ok(Ok(response)) => Ok(response), + Ok(Err(FunctionCallError::Fatal(message))) => Err(CodexErr::Fatal(message)), + Ok(Err(other)) => Err(CodexErr::Fatal(other.to_string())), + Err(err) => Err(CodexErr::Fatal(format!( + "tool task failed to receive: {err:?}" + ))), + } } } } diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 9dde10a611..f5f1fac2d5 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -404,6 +404,56 @@ fn create_test_sync_tool() -> ToolSpec { }) } +fn create_grep_files_tool() -> ToolSpec { + let mut properties = BTreeMap::new(); + properties.insert( + "pattern".to_string(), + JsonSchema::String { + description: Some("Regular expression pattern to search for.".to_string()), + }, + ); + properties.insert( + "include".to_string(), + JsonSchema::String { + description: Some( + "Optional glob that limits which files are searched (e.g. \"*.rs\" or \ + \"*.{ts,tsx}\")." + .to_string(), + ), + }, + ); + properties.insert( + "path".to_string(), + JsonSchema::String { + description: Some( + "Directory or file path to search. Defaults to the session's working directory." + .to_string(), + ), + }, + ); + properties.insert( + "limit".to_string(), + JsonSchema::Number { + description: Some( + "Maximum number of file paths to return (defaults to 100).".to_string(), + ), + }, + ); + + ToolSpec::Function(ResponsesApiTool { + name: "grep_files".to_string(), + description: "Finds files whose contents match the pattern and lists them by modification \ + time." + .to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["pattern".to_string()]), + additional_properties: Some(false.into()), + }, + }) +} + fn create_read_file_tool() -> ToolSpec { let mut properties = BTreeMap::new(); properties.insert( @@ -440,6 +490,51 @@ fn create_read_file_tool() -> ToolSpec { }, }) } + +fn create_list_dir_tool() -> ToolSpec { + let mut properties = BTreeMap::new(); + properties.insert( + "dir_path".to_string(), + JsonSchema::String { + description: Some("Absolute path to the directory to list.".to_string()), + }, + ); + properties.insert( + "offset".to_string(), + JsonSchema::Number { + description: Some( + "The entry number to start listing from. Must be 1 or greater.".to_string(), + ), + }, + ); + properties.insert( + "limit".to_string(), + JsonSchema::Number { + description: Some("The maximum number of entries to return.".to_string()), + }, + ); + properties.insert( + "depth".to_string(), + JsonSchema::Number { + description: Some( + "The maximum directory depth to traverse. Must be 1 or greater.".to_string(), + ), + }, + ); + + ToolSpec::Function(ResponsesApiTool { + name: "list_dir".to_string(), + description: + "Lists entries in a local directory with 1-indexed entry numbers and simple type labels." + .to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["dir_path".to_string()]), + additional_properties: Some(false.into()), + }, + }) +} /// TODO(dylan): deprecate once we get rid of json tool #[derive(Serialize, Deserialize)] pub(crate) struct ApplyPatchToolArgs { @@ -649,6 +744,8 @@ pub(crate) fn build_specs( use crate::exec_command::create_write_stdin_tool_for_responses_api; use crate::tools::handlers::ApplyPatchHandler; use crate::tools::handlers::ExecStreamHandler; + use crate::tools::handlers::GrepFilesHandler; + use crate::tools::handlers::ListDirHandler; use crate::tools::handlers::McpHandler; use crate::tools::handlers::PlanHandler; use crate::tools::handlers::ReadFileHandler; @@ -731,8 +828,16 @@ pub(crate) fn build_specs( if config .experimental_supported_tools - .iter() - .any(|tool| tool == "read_file") + .contains(&"grep_files".to_string()) + { + let grep_files_handler = Arc::new(GrepFilesHandler); + builder.push_spec_with_parallel_support(create_grep_files_tool(), true); + builder.register_handler("grep_files", grep_files_handler); + } + + if config + .experimental_supported_tools + .contains(&"read_file".to_string()) { let read_file_handler = Arc::new(ReadFileHandler); builder.push_spec_with_parallel_support(create_read_file_tool(), true); @@ -742,7 +847,16 @@ pub(crate) fn build_specs( if config .experimental_supported_tools .iter() - .any(|tool| tool == "test_sync_tool") + .any(|tool| tool == "list_dir") + { + let list_dir_handler = Arc::new(ListDirHandler); + builder.push_spec_with_parallel_support(create_list_dir_tool(), true); + builder.register_handler("list_dir", list_dir_handler); + } + + if config + .experimental_supported_tools + .contains(&"test_sync_tool".to_string()) { let test_sync_handler = Arc::new(TestSyncHandler); builder.push_spec_with_parallel_support(create_test_sync_tool(), true); @@ -892,6 +1006,8 @@ mod tests { let (tools, _) = build_specs(&config, None).build(); assert!(!find_tool(&tools, "unified_exec").supports_parallel_tool_calls); + assert!(find_tool(&tools, "grep_files").supports_parallel_tool_calls); + assert!(find_tool(&tools, "list_dir").supports_parallel_tool_calls); assert!(find_tool(&tools, "read_file").supports_parallel_tool_calls); } @@ -920,6 +1036,12 @@ mod tests { .iter() .any(|tool| tool_name(&tool.spec) == "read_file") ); + assert!( + tools + .iter() + .any(|tool| tool_name(&tool.spec) == "grep_files") + ); + assert!(tools.iter().any(|tool| tool_name(&tool.spec) == "list_dir")); assert!( tools .iter() diff --git a/codex-rs/core/tests/common/responses.rs b/codex-rs/core/tests/common/responses.rs index 24ea824ac6..98b3eca1f4 100644 --- a/codex-rs/core/tests/common/responses.rs +++ b/codex-rs/core/tests/common/responses.rs @@ -1,11 +1,105 @@ +use std::sync::Arc; +use std::sync::Mutex; + use serde_json::Value; use wiremock::BodyPrintLimit; +use wiremock::Match; use wiremock::Mock; +use wiremock::MockBuilder; use wiremock::MockServer; use wiremock::Respond; use wiremock::ResponseTemplate; use wiremock::matchers::method; -use wiremock::matchers::path; +use wiremock::matchers::path_regex; + +#[derive(Debug, Clone)] +pub struct ResponseMock { + requests: Arc>>, +} + +impl ResponseMock { + fn new() -> Self { + Self { + requests: Arc::new(Mutex::new(Vec::new())), + } + } + + pub fn single_request(&self) -> ResponsesRequest { + let requests = self.requests.lock().unwrap(); + if requests.len() != 1 { + panic!("expected 1 request, got {}", requests.len()); + } + requests.first().unwrap().clone() + } + + pub fn requests(&self) -> Vec { + self.requests.lock().unwrap().clone() + } +} + +#[derive(Debug, Clone)] +pub struct ResponsesRequest(wiremock::Request); + +impl ResponsesRequest { + pub fn body_json(&self) -> Value { + self.0.body_json().unwrap() + } + + pub fn input(&self) -> Vec { + self.0.body_json::().unwrap()["input"] + .as_array() + .expect("input array not found in request") + .clone() + } + + pub fn function_call_output(&self, call_id: &str) -> Value { + self.call_output(call_id, "function_call_output") + } + + pub fn custom_tool_call_output(&self, call_id: &str) -> Value { + self.call_output(call_id, "custom_tool_call_output") + } + + pub fn call_output(&self, call_id: &str, call_type: &str) -> Value { + self.input() + .iter() + .find(|item| { + item.get("type").unwrap() == call_type && item.get("call_id").unwrap() == call_id + }) + .cloned() + .unwrap_or_else(|| panic!("function call output {call_id} item not found in request")) + } + + pub fn header(&self, name: &str) -> Option { + self.0 + .headers + .get(name) + .and_then(|v| v.to_str().ok()) + .map(str::to_string) + } + + pub fn path(&self) -> String { + self.0.url.path().to_string() + } + + pub fn query_param(&self, name: &str) -> Option { + self.0 + .url + .query_pairs() + .find(|(k, _)| k == name) + .map(|(_, v)| v.to_string()) + } +} + +impl Match for ResponseMock { + fn matches(&self, request: &wiremock::Request) -> bool { + self.requests + .lock() + .unwrap() + .push(ResponsesRequest(request.clone())); + true + } +} /// Build an SSE stream body from a list of JSON events. pub fn sse(events: Vec) -> String { @@ -161,34 +255,40 @@ pub fn sse_response(body: String) -> ResponseTemplate { .set_body_raw(body, "text/event-stream") } -pub async fn mount_sse_once_match(server: &MockServer, matcher: M, body: String) +fn base_mock() -> (MockBuilder, ResponseMock) { + let response_mock = ResponseMock::new(); + let mock = Mock::given(method("POST")) + .and(path_regex(".*/responses$")) + .and(response_mock.clone()); + (mock, response_mock) +} + +pub async fn mount_sse_once_match(server: &MockServer, matcher: M, body: String) -> ResponseMock where M: wiremock::Match + Send + Sync + 'static, { - Mock::given(method("POST")) - .and(path("/v1/responses")) - .and(matcher) + let (mock, response_mock) = base_mock(); + mock.and(matcher) .respond_with(sse_response(body)) .up_to_n_times(1) .mount(server) .await; + response_mock } -pub async fn mount_sse_once(server: &MockServer, body: String) { - Mock::given(method("POST")) - .and(path("/v1/responses")) - .respond_with(sse_response(body)) - .expect(1) +pub async fn mount_sse_once(server: &MockServer, body: String) -> ResponseMock { + let (mock, response_mock) = base_mock(); + mock.respond_with(sse_response(body)) + .up_to_n_times(1) .mount(server) .await; + response_mock } -pub async fn mount_sse(server: &MockServer, body: String) { - Mock::given(method("POST")) - .and(path("/v1/responses")) - .respond_with(sse_response(body)) - .mount(server) - .await; +pub async fn mount_sse(server: &MockServer, body: String) -> ResponseMock { + let (mock, response_mock) = base_mock(); + mock.respond_with(sse_response(body)).mount(server).await; + response_mock } pub async fn start_mock_server() -> MockServer { @@ -201,7 +301,7 @@ pub async fn start_mock_server() -> MockServer { /// Mounts a sequence of SSE response bodies and serves them in order for each /// POST to `/v1/responses`. Panics if more requests are received than bodies /// provided. Also asserts the exact number of expected calls. -pub async fn mount_sse_sequence(server: &MockServer, bodies: Vec) { +pub async fn mount_sse_sequence(server: &MockServer, bodies: Vec) -> ResponseMock { use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -228,10 +328,11 @@ pub async fn mount_sse_sequence(server: &MockServer, bodies: Vec) { responses: bodies, }; - Mock::given(method("POST")) - .and(path("/v1/responses")) - .respond_with(responder) + let (mock, response_mock) = base_mock(); + mock.respond_with(responder) .expect(num_calls as u64) .mount(server) .await; + + response_mock } diff --git a/codex-rs/core/tests/suite/cli_stream.rs b/codex-rs/core/tests/suite/cli_stream.rs index 8fc36772ad..f9408d5a9c 100644 --- a/codex-rs/core/tests/suite/cli_stream.rs +++ b/codex-rs/core/tests/suite/cli_stream.rs @@ -106,16 +106,12 @@ async fn exec_cli_applies_experimental_instructions_file() { "data: {\"type\":\"response.created\",\"response\":{}}\n\n", "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"r1\"}}\n\n" ); - Mock::given(method("POST")) - .and(path("/v1/responses")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_raw(sse, "text/event-stream"), - ) - .expect(1) - .mount(&server) - .await; + let resp_mock = core_test_support::responses::mount_sse_once_match( + &server, + path("/v1/responses"), + sse.to_string(), + ) + .await; // Create a temporary instructions file with a unique marker we can assert // appears in the outbound request payload. @@ -164,8 +160,8 @@ async fn exec_cli_applies_experimental_instructions_file() { // Inspect the captured request and verify our custom base instructions were // included in the `instructions` field. - let request = &server.received_requests().await.unwrap()[0]; - let body = request.body_json::().unwrap(); + let request = resp_mock.single_request(); + let body = request.body_json(); let instructions = body .get("instructions") .and_then(|v| v.as_str()) diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index c49c38e30a..eb14dabb81 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -223,15 +223,9 @@ async fn resume_includes_initial_messages_and_sends_prior_items() { // Mock server that will receive the resumed request let server = MockServer::start().await; - let first = ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_raw(sse_completed("resp1"), "text/event-stream"); - Mock::given(method("POST")) - .and(path("/v1/responses")) - .respond_with(first) - .expect(1) - .mount(&server) - .await; + let resp_mock = + responses::mount_sse_once_match(&server, path("/v1/responses"), sse_completed("resp1")) + .await; // Configure Codex to resume from our file let model_provider = ModelProviderInfo { @@ -277,8 +271,8 @@ async fn resume_includes_initial_messages_and_sends_prior_items() { .unwrap(); wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await; - let request = &server.received_requests().await.unwrap()[0]; - let request_body = request.body_json::().unwrap(); + let request = resp_mock.single_request(); + let request_body = request.body_json(); let expected_input = json!([ { "type": "message", @@ -372,18 +366,9 @@ async fn includes_base_instructions_override_in_request() { skip_if_no_network!(); // Mock server let server = MockServer::start().await; - - // First request – must NOT include `previous_response_id`. - let first = ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_raw(sse_completed("resp1"), "text/event-stream"); - - Mock::given(method("POST")) - .and(path("/v1/responses")) - .respond_with(first) - .expect(1) - .mount(&server) - .await; + let resp_mock = + responses::mount_sse_once_match(&server, path("/v1/responses"), sse_completed("resp1")) + .await; let model_provider = ModelProviderInfo { base_url: Some(format!("{}/v1", server.uri())), @@ -414,8 +399,8 @@ async fn includes_base_instructions_override_in_request() { wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await; - let request = &server.received_requests().await.unwrap()[0]; - let request_body = request.body_json::().unwrap(); + let request = resp_mock.single_request(); + let request_body = request.body_json(); assert!( request_body["instructions"] @@ -570,16 +555,9 @@ async fn includes_user_instructions_message_in_request() { skip_if_no_network!(); let server = MockServer::start().await; - let first = ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_raw(sse_completed("resp1"), "text/event-stream"); - - Mock::given(method("POST")) - .and(path("/v1/responses")) - .respond_with(first) - .expect(1) - .mount(&server) - .await; + let resp_mock = + responses::mount_sse_once_match(&server, path("/v1/responses"), sse_completed("resp1")) + .await; let model_provider = ModelProviderInfo { base_url: Some(format!("{}/v1", server.uri())), @@ -610,8 +588,8 @@ async fn includes_user_instructions_message_in_request() { wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await; - let request = &server.received_requests().await.unwrap()[0]; - let request_body = request.body_json::().unwrap(); + let request = resp_mock.single_request(); + let request_body = request.body_json(); assert!( !request_body["instructions"] diff --git a/codex-rs/core/tests/suite/grep_files.rs b/codex-rs/core/tests/suite/grep_files.rs new file mode 100644 index 0000000000..31195f7e3b --- /dev/null +++ b/codex-rs/core/tests/suite/grep_files.rs @@ -0,0 +1,237 @@ +#![cfg(not(target_os = "windows"))] + +use anyhow::Result; +use codex_core::model_family::find_family_for_model; +use codex_core::protocol::AskForApproval; +use codex_core::protocol::EventMsg; +use codex_core::protocol::InputItem; +use codex_core::protocol::Op; +use codex_core::protocol::SandboxPolicy; +use codex_protocol::config_types::ReasoningSummary; +use core_test_support::responses; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_function_call; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::sse; +use core_test_support::responses::start_mock_server; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::TestCodex; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use serde_json::Value; +use std::collections::HashSet; +use std::path::Path; +use std::process::Command as StdCommand; +use wiremock::matchers::any; + +const MODEL_WITH_TOOL: &str = "test-gpt-5-codex"; + +fn ripgrep_available() -> bool { + StdCommand::new("rg") + .arg("--version") + .output() + .map(|output| output.status.success()) + .unwrap_or(false) +} + +macro_rules! skip_if_ripgrep_missing { + ($ret:expr $(,)?) => {{ + if !ripgrep_available() { + eprintln!("rg not available in PATH; skipping test"); + return $ret; + } + }}; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn grep_files_tool_collects_matches() -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_ripgrep_missing!(Ok(())); + + let server = start_mock_server().await; + let test = build_test_codex(&server).await?; + + let search_dir = test.cwd.path().join("src"); + std::fs::create_dir_all(&search_dir)?; + let alpha = search_dir.join("alpha.rs"); + let beta = search_dir.join("beta.rs"); + let gamma = search_dir.join("gamma.txt"); + std::fs::write(&alpha, "alpha needle\n")?; + std::fs::write(&beta, "beta needle\n")?; + std::fs::write(&gamma, "needle in text but excluded\n")?; + + let call_id = "grep-files-collect"; + let arguments = serde_json::json!({ + "pattern": "needle", + "path": search_dir.to_string_lossy(), + "include": "*.rs", + }) + .to_string(); + + mount_tool_sequence(&server, call_id, &arguments, "grep_files").await; + submit_turn(&test, "please find uses of needle").await?; + + let bodies = recorded_bodies(&server).await?; + let tool_output = find_tool_output(&bodies, call_id).expect("tool output present"); + let payload = tool_output.get("output").expect("output field present"); + let (content_opt, success_opt) = extract_content_and_success(payload); + let content = content_opt.expect("content present"); + let success = success_opt.unwrap_or(true); + assert!(success, "expected success for matches, got {payload:?}"); + + let entries = collect_file_names(content); + assert_eq!(entries.len(), 2, "content: {content}"); + assert!( + entries.contains("alpha.rs"), + "missing alpha.rs in {entries:?}" + ); + assert!( + entries.contains("beta.rs"), + "missing beta.rs in {entries:?}" + ); + assert!( + !entries.contains("gamma.txt"), + "txt file should be filtered out: {entries:?}" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn grep_files_tool_reports_empty_results() -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_ripgrep_missing!(Ok(())); + + let server = start_mock_server().await; + let test = build_test_codex(&server).await?; + + let search_dir = test.cwd.path().join("logs"); + std::fs::create_dir_all(&search_dir)?; + std::fs::write(search_dir.join("output.txt"), "no hits here")?; + + let call_id = "grep-files-empty"; + let arguments = serde_json::json!({ + "pattern": "needle", + "path": search_dir.to_string_lossy(), + "limit": 5, + }) + .to_string(); + + mount_tool_sequence(&server, call_id, &arguments, "grep_files").await; + submit_turn(&test, "search again").await?; + + let bodies = recorded_bodies(&server).await?; + let tool_output = find_tool_output(&bodies, call_id).expect("tool output present"); + let payload = tool_output.get("output").expect("output field present"); + let (content_opt, success_opt) = extract_content_and_success(payload); + let content = content_opt.expect("content present"); + if let Some(success) = success_opt { + assert!(!success, "expected success=false payload: {payload:?}"); + } + assert_eq!(content, "No matches found."); + + Ok(()) +} + +#[allow(clippy::expect_used)] +async fn build_test_codex(server: &wiremock::MockServer) -> Result { + let mut builder = test_codex().with_config(|config| { + config.model = MODEL_WITH_TOOL.to_string(); + config.model_family = + find_family_for_model(MODEL_WITH_TOOL).expect("model family for test model"); + }); + builder.build(server).await +} + +async fn submit_turn(test: &TestCodex, prompt: &str) -> Result<()> { + let session_model = test.session_configured.model.clone(); + + test.codex + .submit(Op::UserTurn { + items: vec![InputItem::Text { + text: prompt.into(), + }], + final_output_json_schema: None, + cwd: test.cwd.path().to_path_buf(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::DangerFullAccess, + model: session_model, + effort: None, + summary: ReasoningSummary::Auto, + }) + .await?; + + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TaskComplete(_)) + }) + .await; + Ok(()) +} + +async fn mount_tool_sequence( + server: &wiremock::MockServer, + call_id: &str, + arguments: &str, + tool_name: &str, +) { + let first_response = sse(vec![ + ev_response_created("resp-1"), + ev_function_call(call_id, tool_name, arguments), + ev_completed("resp-1"), + ]); + responses::mount_sse_once_match(server, any(), first_response).await; + + let second_response = sse(vec![ + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]); + responses::mount_sse_once_match(server, any(), second_response).await; +} + +#[allow(clippy::expect_used)] +async fn recorded_bodies(server: &wiremock::MockServer) -> Result> { + let requests = server.received_requests().await.expect("requests recorded"); + Ok(requests + .iter() + .map(|req| req.body_json::().expect("request json")) + .collect()) +} + +fn find_tool_output<'a>(requests: &'a [Value], call_id: &str) -> Option<&'a Value> { + requests.iter().find_map(|body| { + body.get("input") + .and_then(Value::as_array) + .and_then(|items| { + items.iter().find(|item| { + item.get("type").and_then(Value::as_str) == Some("function_call_output") + && item.get("call_id").and_then(Value::as_str) == Some(call_id) + }) + }) + }) +} + +fn collect_file_names(content: &str) -> HashSet { + content + .lines() + .filter_map(|line| { + if line.trim().is_empty() { + return None; + } + Path::new(line) + .file_name() + .map(|name| name.to_string_lossy().into_owned()) + }) + .collect() +} + +fn extract_content_and_success(value: &Value) -> (Option<&str>, Option) { + match value { + Value::String(text) => (Some(text.as_str()), None), + Value::Object(obj) => ( + obj.get("content").and_then(Value::as_str), + obj.get("success").and_then(Value::as_bool), + ), + _ => (None, None), + } +} diff --git a/codex-rs/core/tests/suite/list_dir.rs b/codex-rs/core/tests/suite/list_dir.rs new file mode 100644 index 0000000000..1aa5a6484b --- /dev/null +++ b/codex-rs/core/tests/suite/list_dir.rs @@ -0,0 +1,460 @@ +#![cfg(not(target_os = "windows"))] + +use codex_core::protocol::AskForApproval; +use codex_core::protocol::EventMsg; +use codex_core::protocol::InputItem; +use codex_core::protocol::Op; +use codex_core::protocol::SandboxPolicy; +use codex_protocol::config_types::ReasoningSummary; +use core_test_support::responses; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_function_call; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::sse; +use core_test_support::responses::start_mock_server; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::TestCodex; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; +use serde_json::Value; +use wiremock::matchers::any; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[ignore = "disabled until we enable list_dir tool"] +async fn list_dir_tool_returns_entries() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + + let TestCodex { + codex, + cwd, + session_configured, + .. + } = test_codex().build(&server).await?; + + let dir_path = cwd.path().join("sample_dir"); + std::fs::create_dir(&dir_path)?; + std::fs::write(dir_path.join("alpha.txt"), "first file")?; + std::fs::create_dir(dir_path.join("nested"))?; + let dir_path = dir_path.to_string_lossy().to_string(); + + let call_id = "list-dir-call"; + let arguments = serde_json::json!({ + "dir_path": dir_path, + "offset": 1, + "limit": 2, + }) + .to_string(); + + let first_response = sse(vec![ + ev_response_created("resp-1"), + ev_function_call(call_id, "list_dir", &arguments), + ev_completed("resp-1"), + ]); + responses::mount_sse_once_match(&server, any(), first_response).await; + + let second_response = sse(vec![ + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]); + responses::mount_sse_once_match(&server, any(), second_response).await; + + let session_model = session_configured.model.clone(); + + codex + .submit(Op::UserTurn { + items: vec![InputItem::Text { + text: "list directory contents".into(), + }], + final_output_json_schema: None, + cwd: cwd.path().to_path_buf(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::DangerFullAccess, + model: session_model, + effort: None, + summary: ReasoningSummary::Auto, + }) + .await?; + + wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await; + + let requests = server.received_requests().await.expect("recorded requests"); + let request_bodies = requests + .iter() + .map(|req| req.body_json::().unwrap()) + .collect::>(); + assert!( + !request_bodies.is_empty(), + "expected at least one request body" + ); + + let tool_output_item = request_bodies + .iter() + .find_map(|body| { + body.get("input") + .and_then(Value::as_array) + .and_then(|items| { + items.iter().find(|item| { + item.get("type").and_then(Value::as_str) == Some("function_call_output") + }) + }) + }) + .unwrap_or_else(|| { + panic!("function_call_output item not found in requests: {request_bodies:#?}") + }); + + assert_eq!( + tool_output_item.get("call_id").and_then(Value::as_str), + Some(call_id) + ); + + let output_text = tool_output_item + .get("output") + .and_then(|value| match value { + Value::String(text) => Some(text.as_str()), + Value::Object(obj) => obj.get("content").and_then(Value::as_str), + _ => None, + }) + .expect("output text present"); + assert_eq!(output_text, "E1: [file] alpha.txt\nE2: [dir] nested"); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[ignore = "disabled until we enable list_dir tool"] +async fn list_dir_tool_depth_one_omits_children() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + + let TestCodex { + codex, + cwd, + session_configured, + .. + } = test_codex().build(&server).await?; + + let dir_path = cwd.path().join("depth_one"); + std::fs::create_dir(&dir_path)?; + std::fs::write(dir_path.join("alpha.txt"), "alpha")?; + std::fs::create_dir(dir_path.join("nested"))?; + std::fs::write(dir_path.join("nested").join("beta.txt"), "beta")?; + let dir_path = dir_path.to_string_lossy().to_string(); + + let call_id = "list-dir-depth1"; + let arguments = serde_json::json!({ + "dir_path": dir_path, + "offset": 1, + "limit": 10, + "depth": 1, + }) + .to_string(); + + let first_response = sse(vec![ + ev_response_created("resp-1"), + ev_function_call(call_id, "list_dir", &arguments), + ev_completed("resp-1"), + ]); + responses::mount_sse_once_match(&server, any(), first_response).await; + + let second_response = sse(vec![ + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]); + responses::mount_sse_once_match(&server, any(), second_response).await; + + let session_model = session_configured.model.clone(); + + codex + .submit(Op::UserTurn { + items: vec![InputItem::Text { + text: "list directory contents depth one".into(), + }], + final_output_json_schema: None, + cwd: cwd.path().to_path_buf(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::DangerFullAccess, + model: session_model, + effort: None, + summary: ReasoningSummary::Auto, + }) + .await?; + + wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await; + + let requests = server.received_requests().await.expect("recorded requests"); + let request_bodies = requests + .iter() + .map(|req| req.body_json::().unwrap()) + .collect::>(); + assert!( + !request_bodies.is_empty(), + "expected at least one request body" + ); + + let tool_output_item = request_bodies + .iter() + .find_map(|body| { + body.get("input") + .and_then(Value::as_array) + .and_then(|items| { + items.iter().find(|item| { + item.get("type").and_then(Value::as_str) == Some("function_call_output") + }) + }) + }) + .unwrap_or_else(|| { + panic!("function_call_output item not found in requests: {request_bodies:#?}") + }); + + assert_eq!( + tool_output_item.get("call_id").and_then(Value::as_str), + Some(call_id) + ); + + let output_text = tool_output_item + .get("output") + .and_then(|value| match value { + Value::String(text) => Some(text.as_str()), + Value::Object(obj) => obj.get("content").and_then(Value::as_str), + _ => None, + }) + .expect("output text present"); + assert_eq!(output_text, "E1: [file] alpha.txt\nE2: [dir] nested"); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[ignore = "disabled until we enable list_dir tool"] +async fn list_dir_tool_depth_two_includes_children_only() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + + let TestCodex { + codex, + cwd, + session_configured, + .. + } = test_codex().build(&server).await?; + + let dir_path = cwd.path().join("depth_two"); + std::fs::create_dir(&dir_path)?; + std::fs::write(dir_path.join("alpha.txt"), "alpha")?; + let nested = dir_path.join("nested"); + std::fs::create_dir(&nested)?; + std::fs::write(nested.join("beta.txt"), "beta")?; + let deeper = nested.join("grand"); + std::fs::create_dir(&deeper)?; + std::fs::write(deeper.join("gamma.txt"), "gamma")?; + let dir_path_string = dir_path.to_string_lossy().to_string(); + + let call_id = "list-dir-depth2"; + let arguments = serde_json::json!({ + "dir_path": dir_path_string, + "offset": 1, + "limit": 10, + "depth": 2, + }) + .to_string(); + + let first_response = sse(vec![ + serde_json::json!({ + "type": "response.created", + "response": {"id": "resp-1"} + }), + ev_function_call(call_id, "list_dir", &arguments), + ev_completed("resp-1"), + ]); + responses::mount_sse_once_match(&server, any(), first_response).await; + + let second_response = sse(vec![ + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]); + responses::mount_sse_once_match(&server, any(), second_response).await; + + let session_model = session_configured.model.clone(); + + codex + .submit(Op::UserTurn { + items: vec![InputItem::Text { + text: "list directory contents depth two".into(), + }], + final_output_json_schema: None, + cwd: cwd.path().to_path_buf(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::DangerFullAccess, + model: session_model, + effort: None, + summary: ReasoningSummary::Auto, + }) + .await?; + + wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await; + + let requests = server.received_requests().await.expect("recorded requests"); + let request_bodies = requests + .iter() + .map(|req| req.body_json::().unwrap()) + .collect::>(); + assert!( + !request_bodies.is_empty(), + "expected at least one request body" + ); + + let tool_output_item = request_bodies + .iter() + .find_map(|body| { + body.get("input") + .and_then(Value::as_array) + .and_then(|items| { + items.iter().find(|item| { + item.get("type").and_then(Value::as_str) == Some("function_call_output") + }) + }) + }) + .unwrap_or_else(|| { + panic!("function_call_output item not found in requests: {request_bodies:#?}") + }); + + assert_eq!( + tool_output_item.get("call_id").and_then(Value::as_str), + Some(call_id) + ); + + let output_text = tool_output_item + .get("output") + .and_then(|value| match value { + Value::String(text) => Some(text.as_str()), + Value::Object(obj) => obj.get("content").and_then(Value::as_str), + _ => None, + }) + .expect("output text present"); + assert_eq!( + output_text, + "E1: [file] alpha.txt\nE2: [dir] nested\nE3: [file] nested/beta.txt\nE4: [dir] nested/grand" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[ignore = "disabled until we enable list_dir tool"] +async fn list_dir_tool_depth_three_includes_grandchildren() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + + let TestCodex { + codex, + cwd, + session_configured, + .. + } = test_codex().build(&server).await?; + + let dir_path = cwd.path().join("depth_three"); + std::fs::create_dir(&dir_path)?; + std::fs::write(dir_path.join("alpha.txt"), "alpha")?; + let nested = dir_path.join("nested"); + std::fs::create_dir(&nested)?; + std::fs::write(nested.join("beta.txt"), "beta")?; + let deeper = nested.join("grand"); + std::fs::create_dir(&deeper)?; + std::fs::write(deeper.join("gamma.txt"), "gamma")?; + let dir_path_string = dir_path.to_string_lossy().to_string(); + + let call_id = "list-dir-depth3"; + let arguments = serde_json::json!({ + "dir_path": dir_path_string, + "offset": 1, + "limit": 10, + "depth": 3, + }) + .to_string(); + + let first_response = sse(vec![ + serde_json::json!({ + "type": "response.created", + "response": {"id": "resp-1"} + }), + ev_function_call(call_id, "list_dir", &arguments), + ev_completed("resp-1"), + ]); + responses::mount_sse_once_match(&server, any(), first_response).await; + + let second_response = sse(vec![ + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]); + responses::mount_sse_once_match(&server, any(), second_response).await; + + let session_model = session_configured.model.clone(); + + codex + .submit(Op::UserTurn { + items: vec![InputItem::Text { + text: "list directory contents depth three".into(), + }], + final_output_json_schema: None, + cwd: cwd.path().to_path_buf(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::DangerFullAccess, + model: session_model, + effort: None, + summary: ReasoningSummary::Auto, + }) + .await?; + + wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await; + + let requests = server.received_requests().await.expect("recorded requests"); + let request_bodies = requests + .iter() + .map(|req| req.body_json::().unwrap()) + .collect::>(); + assert!( + !request_bodies.is_empty(), + "expected at least one request body" + ); + + let tool_output_item = request_bodies + .iter() + .find_map(|body| { + body.get("input") + .and_then(Value::as_array) + .and_then(|items| { + items.iter().find(|item| { + item.get("type").and_then(Value::as_str) == Some("function_call_output") + }) + }) + }) + .unwrap_or_else(|| { + panic!("function_call_output item not found in requests: {request_bodies:#?}") + }); + + assert_eq!( + tool_output_item.get("call_id").and_then(Value::as_str), + Some(call_id) + ); + + let output_text = tool_output_item + .get("output") + .and_then(|value| match value { + Value::String(text) => Some(text.as_str()), + Value::Object(obj) => obj.get("content").and_then(Value::as_str), + _ => None, + }) + .expect("output text present"); + assert_eq!( + output_text, + "E1: [file] alpha.txt\nE2: [dir] nested\nE3: [file] nested/beta.txt\nE4: [dir] nested/grand\nE5: [file] nested/grand/gamma.txt" + ); + + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index fb8502c480..72e8cc2299 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -9,7 +9,9 @@ mod compact_resume_fork; mod exec; mod exec_stream_events; mod fork_conversation; +mod grep_files; mod json_result; +mod list_dir; mod live_cli; mod model_overrides; mod model_tools; diff --git a/codex-rs/core/tests/suite/model_tools.rs b/codex-rs/core/tests/suite/model_tools.rs index 952cfcbdec..f9c8cbd1f5 100644 --- a/codex-rs/core/tests/suite/model_tools.rs +++ b/codex-rs/core/tests/suite/model_tools.rs @@ -10,14 +10,11 @@ use codex_core::protocol::InputItem; use codex_core::protocol::Op; use core_test_support::load_default_config_for_test; use core_test_support::load_sse_fixture_with_id; +use core_test_support::responses; use core_test_support::skip_if_no_network; use core_test_support::wait_for_event; use tempfile::TempDir; -use wiremock::Mock; use wiremock::MockServer; -use wiremock::ResponseTemplate; -use wiremock::matchers::method; -use wiremock::matchers::path; fn sse_completed(id: &str) -> String { load_sse_fixture_with_id("tests/fixtures/completed_template.json", id) @@ -44,16 +41,7 @@ async fn collect_tool_identifiers_for_model(model: &str) -> Vec { let server = MockServer::start().await; let sse = sse_completed(model); - let template = ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_raw(sse, "text/event-stream"); - - Mock::given(method("POST")) - .and(path("/v1/responses")) - .respond_with(template) - .expect(1) - .mount(&server) - .await; + let resp_mock = responses::mount_sse_once_match(&server, wiremock::matchers::any(), sse).await; let model_provider = ModelProviderInfo { base_url: Some(format!("{}/v1", server.uri())), @@ -93,13 +81,7 @@ async fn collect_tool_identifiers_for_model(model: &str) -> Vec { .unwrap(); wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await; - let requests = server.received_requests().await.unwrap(); - assert_eq!( - requests.len(), - 1, - "expected a single request for model {model}" - ); - let body = requests[0].body_json::().unwrap(); + let body = resp_mock.single_request().body_json(); tool_identifiers(&body) } diff --git a/codex-rs/core/tests/suite/read_file.rs b/codex-rs/core/tests/suite/read_file.rs index a6c8a7a181..fc5a94f91b 100644 --- a/codex-rs/core/tests/suite/read_file.rs +++ b/codex-rs/core/tests/suite/read_file.rs @@ -58,7 +58,7 @@ async fn read_file_tool_returns_requested_lines() -> anyhow::Result<()> { ev_assistant_message("msg-1", "done"), ev_completed("resp-2"), ]); - responses::mount_sse_once_match(&server, any(), second_response).await; + let second_mock = responses::mount_sse_once_match(&server, any(), second_response).await; let session_model = session_configured.model.clone(); @@ -79,36 +79,12 @@ async fn read_file_tool_returns_requested_lines() -> anyhow::Result<()> { wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await; - let requests = server.received_requests().await.expect("recorded requests"); - let request_bodies = requests - .iter() - .map(|req| req.body_json::().unwrap()) - .collect::>(); - assert!( - !request_bodies.is_empty(), - "expected at least one request body" - ); - - let tool_output_item = request_bodies - .iter() - .find_map(|body| { - body.get("input") - .and_then(Value::as_array) - .and_then(|items| { - items.iter().find(|item| { - item.get("type").and_then(Value::as_str) == Some("function_call_output") - }) - }) - }) - .unwrap_or_else(|| { - panic!("function_call_output item not found in requests: {request_bodies:#?}") - }); - + let req = second_mock.single_request(); + let tool_output_item = req.function_call_output(call_id); assert_eq!( tool_output_item.get("call_id").and_then(Value::as_str), Some(call_id) ); - let output_text = tool_output_item .get("output") .and_then(|value| match value { diff --git a/codex-rs/core/tests/suite/review.rs b/codex-rs/core/tests/suite/review.rs index 66b67ae980..f3eeb3a310 100644 --- a/codex-rs/core/tests/suite/review.rs +++ b/codex-rs/core/tests/suite/review.rs @@ -445,7 +445,7 @@ async fn review_input_isolated_from_parent_history() { .await; let _complete = wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await; - // Assert the request `input` contains the environment context followed by the review prompt. + // Assert the request `input` contains the environment context followed by the user review prompt. let request = &server.received_requests().await.unwrap()[0]; let body = request.body_json::().unwrap(); let input = body["input"].as_array().expect("input array"); @@ -473,9 +473,14 @@ async fn review_input_isolated_from_parent_history() { assert_eq!(review_msg["role"].as_str().unwrap(), "user"); assert_eq!( review_msg["content"][0]["text"].as_str().unwrap(), - format!("{REVIEW_PROMPT}\n\n---\n\nNow, here's your task: Please review only this",) + review_prompt, + "user message should only contain the raw review prompt" ); + // Ensure the REVIEW_PROMPT rubric is sent via instructions. + let instructions = body["instructions"].as_str().expect("instructions string"); + assert_eq!(instructions, REVIEW_PROMPT); + // Also verify that a user interruption note was recorded in the rollout. codex.submit(Op::GetPath).await.unwrap(); let history_event = diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 212b03bd5c..a8a5d3afbf 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -232,7 +232,7 @@ async fn streamable_http_tool_call_round_trip() -> anyhow::Result<()> { McpServerConfig { transport: McpServerTransportConfig::StreamableHttp { url: server_url, - bearer_token: None, + bearer_token_env_var: None, }, startup_timeout_sec: Some(Duration::from_secs(10)), tool_timeout_sec: None, @@ -412,7 +412,7 @@ async fn streamable_http_with_oauth_round_trip() -> anyhow::Result<()> { McpServerConfig { transport: McpServerTransportConfig::StreamableHttp { url: server_url, - bearer_token: None, + bearer_token_env_var: None, }, startup_timeout_sec: Some(Duration::from_secs(10)), tool_timeout_sec: None, diff --git a/codex-rs/core/tests/suite/tool_harness.rs b/codex-rs/core/tests/suite/tool_harness.rs index 14e0e1c892..eaefe7d9dc 100644 --- a/codex-rs/core/tests/suite/tool_harness.rs +++ b/codex-rs/core/tests/suite/tool_harness.rs @@ -27,16 +27,6 @@ use serde_json::Value; use serde_json::json; use wiremock::matchers::any; -fn function_call_output(body: &Value) -> Option<&Value> { - body.get("input") - .and_then(Value::as_array) - .and_then(|items| { - items.iter().find(|item| { - item.get("type").and_then(Value::as_str) == Some("function_call_output") - }) - }) -} - fn extract_output_text(item: &Value) -> Option<&str> { item.get("output").and_then(|value| match value { Value::String(text) => Some(text.as_str()), @@ -45,12 +35,6 @@ fn extract_output_text(item: &Value) -> Option<&str> { }) } -fn find_request_with_function_call_output(requests: &[Value]) -> Option<&Value> { - requests - .iter() - .find(|body| function_call_output(body).is_some()) -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn shell_tool_executes_command_and_streams_output() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); @@ -81,7 +65,7 @@ async fn shell_tool_executes_command_and_streams_output() -> anyhow::Result<()> ev_assistant_message("msg-1", "all done"), ev_completed("resp-2"), ]); - responses::mount_sse_once_match(&server, any(), second_response).await; + let second_mock = responses::mount_sse_once_match(&server, any(), second_response).await; let session_model = session_configured.model.clone(); @@ -102,18 +86,9 @@ async fn shell_tool_executes_command_and_streams_output() -> anyhow::Result<()> wait_for_event(&codex, |event| matches!(event, EventMsg::TaskComplete(_))).await; - let requests = server.received_requests().await.expect("recorded requests"); - assert!(!requests.is_empty(), "expected at least one POST request"); - - let request_bodies = requests - .iter() - .map(|req| req.body_json::().expect("request json")) - .collect::>(); - - let body_with_tool_output = find_request_with_function_call_output(&request_bodies) - .expect("function_call_output item not found in requests"); - let output_item = function_call_output(body_with_tool_output).expect("tool output item"); - let output_text = extract_output_text(output_item).expect("output text present"); + let req = second_mock.single_request(); + let output_item = req.function_call_output(call_id); + let output_text = extract_output_text(&output_item).expect("output text present"); let exec_output: Value = serde_json::from_str(output_text)?; assert_eq!(exec_output["metadata"]["exit_code"], 0); let stdout = exec_output["output"].as_str().expect("stdout field"); @@ -159,7 +134,7 @@ async fn update_plan_tool_emits_plan_update_event() -> anyhow::Result<()> { ev_assistant_message("msg-1", "plan acknowledged"), ev_completed("resp-2"), ]); - responses::mount_sse_once_match(&server, any(), second_response).await; + let second_mock = responses::mount_sse_once_match(&server, any(), second_response).await; let session_model = session_configured.model.clone(); @@ -197,22 +172,13 @@ async fn update_plan_tool_emits_plan_update_event() -> anyhow::Result<()> { assert!(saw_plan_update, "expected PlanUpdate event"); - let requests = server.received_requests().await.expect("recorded requests"); - assert!(!requests.is_empty(), "expected at least one POST request"); - - let request_bodies = requests - .iter() - .map(|req| req.body_json::().expect("request json")) - .collect::>(); - - let body_with_tool_output = find_request_with_function_call_output(&request_bodies) - .expect("function_call_output item not found in requests"); - let output_item = function_call_output(body_with_tool_output).expect("tool output item"); + let req = second_mock.single_request(); + let output_item = req.function_call_output(call_id); assert_eq!( output_item.get("call_id").and_then(Value::as_str), Some(call_id) ); - let output_text = extract_output_text(output_item).expect("output text present"); + let output_text = extract_output_text(&output_item).expect("output text present"); assert_eq!(output_text, "Plan updated"); Ok(()) @@ -251,7 +217,7 @@ async fn update_plan_tool_rejects_malformed_payload() -> anyhow::Result<()> { ev_assistant_message("msg-1", "malformed plan payload"), ev_completed("resp-2"), ]); - responses::mount_sse_once_match(&server, any(), second_response).await; + let second_mock = responses::mount_sse_once_match(&server, any(), second_response).await; let session_model = session_configured.model.clone(); @@ -286,22 +252,13 @@ async fn update_plan_tool_rejects_malformed_payload() -> anyhow::Result<()> { "did not expect PlanUpdate event for malformed payload" ); - let requests = server.received_requests().await.expect("recorded requests"); - assert!(!requests.is_empty(), "expected at least one POST request"); - - let request_bodies = requests - .iter() - .map(|req| req.body_json::().expect("request json")) - .collect::>(); - - let body_with_tool_output = find_request_with_function_call_output(&request_bodies) - .expect("function_call_output item not found in requests"); - let output_item = function_call_output(body_with_tool_output).expect("tool output item"); + let req = second_mock.single_request(); + let output_item = req.function_call_output(call_id); assert_eq!( output_item.get("call_id").and_then(Value::as_str), Some(call_id) ); - let output_text = extract_output_text(output_item).expect("output text present"); + let output_text = extract_output_text(&output_item).expect("output text present"); assert!( output_text.contains("failed to parse function arguments"), "expected parse error message in output text, got {output_text:?}" @@ -354,7 +311,7 @@ async fn apply_patch_tool_executes_and_emits_patch_events() -> anyhow::Result<() ev_assistant_message("msg-1", "patch complete"), ev_completed("resp-2"), ]); - responses::mount_sse_once_match(&server, any(), second_response).await; + let second_mock = responses::mount_sse_once_match(&server, any(), second_response).await; let session_model = session_configured.model.clone(); @@ -395,22 +352,13 @@ async fn apply_patch_tool_executes_and_emits_patch_events() -> anyhow::Result<() let patch_end_success = patch_end_success.expect("expected PatchApplyEnd event to capture success flag"); - let requests = server.received_requests().await.expect("recorded requests"); - assert!(!requests.is_empty(), "expected at least one POST request"); - - let request_bodies = requests - .iter() - .map(|req| req.body_json::().expect("request json")) - .collect::>(); - - let body_with_tool_output = find_request_with_function_call_output(&request_bodies) - .expect("function_call_output item not found in requests"); - let output_item = function_call_output(body_with_tool_output).expect("tool output item"); + let req = second_mock.single_request(); + let output_item = req.function_call_output(call_id); assert_eq!( output_item.get("call_id").and_then(Value::as_str), Some(call_id) ); - let output_text = extract_output_text(output_item).expect("output text present"); + let output_text = extract_output_text(&output_item).expect("output text present"); if let Ok(exec_output) = serde_json::from_str::(output_text) { let exit_code = exec_output["metadata"]["exit_code"] @@ -480,7 +428,7 @@ async fn apply_patch_reports_parse_diagnostics() -> anyhow::Result<()> { ev_assistant_message("msg-1", "failed"), ev_completed("resp-2"), ]); - responses::mount_sse_once_match(&server, any(), second_response).await; + let second_mock = responses::mount_sse_once_match(&server, any(), second_response).await; let session_model = session_configured.model.clone(); @@ -501,22 +449,13 @@ async fn apply_patch_reports_parse_diagnostics() -> anyhow::Result<()> { wait_for_event(&codex, |event| matches!(event, EventMsg::TaskComplete(_))).await; - let requests = server.received_requests().await.expect("recorded requests"); - assert!(!requests.is_empty(), "expected at least one POST request"); - - let request_bodies = requests - .iter() - .map(|req| req.body_json::().expect("request json")) - .collect::>(); - - let body_with_tool_output = find_request_with_function_call_output(&request_bodies) - .expect("function_call_output item not found in requests"); - let output_item = function_call_output(body_with_tool_output).expect("tool output item"); + let req = second_mock.single_request(); + let output_item = req.function_call_output(call_id); assert_eq!( output_item.get("call_id").and_then(Value::as_str), Some(call_id) ); - let output_text = extract_output_text(output_item).expect("output text present"); + let output_text = extract_output_text(&output_item).expect("output text present"); assert!( output_text.contains("apply_patch verification failed"), diff --git a/codex-rs/core/tests/suite/tool_parallelism.rs b/codex-rs/core/tests/suite/tool_parallelism.rs index e667df4369..b4e3d1c9ad 100644 --- a/codex-rs/core/tests/suite/tool_parallelism.rs +++ b/codex-rs/core/tests/suite/tool_parallelism.rs @@ -63,8 +63,9 @@ async fn build_codex_with_test_tool(server: &wiremock::MockServer) -> anyhow::Re } fn assert_parallel_duration(actual: Duration) { + // Allow headroom for runtime overhead while still differentiating from serial execution. assert!( - actual < Duration::from_millis(500), + actual < Duration::from_millis(750), "expected parallel execution to finish quickly, got {actual:?}" ); } @@ -83,6 +84,16 @@ async fn read_file_tools_run_in_parallel() -> anyhow::Result<()> { let server = start_mock_server().await; let test = build_codex_with_test_tool(&server).await?; + let warmup_args = json!({ + "sleep_after_ms": 10, + "barrier": { + "id": "parallel-test-sync-warmup", + "participants": 2, + "timeout_ms": 1_000, + } + }) + .to_string(); + let parallel_args = json!({ "sleep_after_ms": 300, "barrier": { @@ -93,6 +104,17 @@ async fn read_file_tools_run_in_parallel() -> anyhow::Result<()> { }) .to_string(); + let warmup_first = sse(vec![ + json!({"type": "response.created", "response": {"id": "resp-warm-1"}}), + ev_function_call("warm-call-1", "test_sync_tool", &warmup_args), + ev_function_call("warm-call-2", "test_sync_tool", &warmup_args), + ev_completed("resp-warm-1"), + ]); + let warmup_second = sse(vec![ + ev_assistant_message("warm-msg-1", "warmup complete"), + ev_completed("resp-warm-2"), + ]); + let first_response = sse(vec![ json!({"type": "response.created", "response": {"id": "resp-1"}}), ev_function_call("call-1", "test_sync_tool", ¶llel_args), @@ -103,7 +125,13 @@ async fn read_file_tools_run_in_parallel() -> anyhow::Result<()> { ev_assistant_message("msg-1", "done"), ev_completed("resp-2"), ]); - mount_sse_sequence(&server, vec![first_response, second_response]).await; + mount_sse_sequence( + &server, + vec![warmup_first, warmup_second, first_response, second_response], + ) + .await; + + run_turn(&test, "warm up parallel tool").await?; let duration = run_turn_and_measure(&test, "exercise sync tool").await?; assert_parallel_duration(duration); diff --git a/codex-rs/core/tests/suite/tools.rs b/codex-rs/core/tests/suite/tools.rs index 08826a1a72..99c70dbd61 100644 --- a/codex-rs/core/tests/suite/tools.rs +++ b/codex-rs/core/tests/suite/tools.rs @@ -15,6 +15,7 @@ use core_test_support::responses::ev_completed; use core_test_support::responses::ev_custom_tool_call; use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_once; use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; @@ -25,7 +26,6 @@ use core_test_support::wait_for_event; use regex_lite::Regex; use serde_json::Value; use serde_json::json; -use wiremock::Request; async fn submit_turn( test: &TestCodex, @@ -58,27 +58,6 @@ async fn submit_turn( Ok(()) } -fn request_bodies(requests: &[Request]) -> Result> { - requests - .iter() - .map(|req| Ok(serde_json::from_slice::(&req.body)?)) - .collect() -} - -fn collect_output_items<'a>(bodies: &'a [Value], ty: &str) -> Vec<&'a Value> { - let mut out = Vec::new(); - for body in bodies { - if let Some(items) = body.get("input").and_then(Value::as_array) { - for item in items { - if item.get("type").and_then(Value::as_str) == Some(ty) { - out.push(item); - } - } - } - } - out -} - fn tool_names(body: &Value) -> Vec { body.get("tools") .and_then(Value::as_array) @@ -107,18 +86,23 @@ async fn custom_tool_unknown_returns_custom_output_error() -> Result<()> { let call_id = "custom-unsupported"; let tool_name = "unsupported_tool"; - let responses = vec![ + mount_sse_once( + &server, sse(vec![ ev_response_created("resp-1"), ev_custom_tool_call(call_id, tool_name, "\"payload\""), ev_completed("resp-1"), ]), + ) + .await; + let mock = mount_sse_once( + &server, sse(vec![ ev_assistant_message("msg-1", "done"), ev_completed("resp-2"), ]), - ]; - mount_sse_sequence(&server, responses).await; + ) + .await; submit_turn( &test, @@ -128,13 +112,7 @@ async fn custom_tool_unknown_returns_custom_output_error() -> Result<()> { ) .await?; - let requests = server.received_requests().await.expect("recorded requests"); - let bodies = request_bodies(&requests)?; - let custom_items = collect_output_items(&bodies, "custom_tool_call_output"); - assert_eq!(custom_items.len(), 1, "expected single custom tool output"); - let item = custom_items[0]; - assert_eq!(item.get("call_id").and_then(Value::as_str), Some(call_id)); - + let item = mock.single_request().custom_tool_call_output(call_id); let output = item .get("output") .and_then(Value::as_str) @@ -170,7 +148,8 @@ async fn shell_escalated_permissions_rejected_then_ok() -> Result<()> { "timeout_ms": 1_000, }); - let responses = vec![ + mount_sse_once( + &server, sse(vec![ ev_response_created("resp-1"), ev_function_call( @@ -180,6 +159,10 @@ async fn shell_escalated_permissions_rejected_then_ok() -> Result<()> { ), ev_completed("resp-1"), ]), + ) + .await; + let second_mock = mount_sse_once( + &server, sse(vec![ ev_response_created("resp-2"), ev_function_call( @@ -189,12 +172,16 @@ async fn shell_escalated_permissions_rejected_then_ok() -> Result<()> { ), ev_completed("resp-2"), ]), + ) + .await; + let third_mock = mount_sse_once( + &server, sse(vec![ ev_assistant_message("msg-1", "done"), ev_completed("resp-3"), ]), - ]; - mount_sse_sequence(&server, responses).await; + ) + .await; submit_turn( &test, @@ -204,46 +191,23 @@ async fn shell_escalated_permissions_rejected_then_ok() -> Result<()> { ) .await?; - let requests = server.received_requests().await.expect("recorded requests"); - let bodies = request_bodies(&requests)?; - let function_outputs = collect_output_items(&bodies, "function_call_output"); - for item in &function_outputs { - let call_id = item - .get("call_id") - .and_then(Value::as_str) - .unwrap_or_default(); - assert!( - call_id == call_id_blocked || call_id == call_id_success, - "unexpected call id {call_id}" - ); - } - let policy = AskForApproval::Never; let expected_message = format!( "approval policy is {policy:?}; reject command — you should not ask for escalated permissions if the approval policy is {policy:?}" ); - let blocked_outputs: Vec<&Value> = function_outputs - .iter() - .filter(|item| item.get("call_id").and_then(Value::as_str) == Some(call_id_blocked)) - .copied() - .collect(); - assert!( - !blocked_outputs.is_empty(), - "expected at least one rejection output for {call_id_blocked}" + let blocked_item = second_mock + .single_request() + .function_call_output(call_id_blocked); + assert_eq!( + blocked_item.get("output").and_then(Value::as_str), + Some(expected_message.as_str()), + "unexpected rejection message" ); - for item in blocked_outputs { - assert_eq!( - item.get("output").and_then(Value::as_str), - Some(expected_message.as_str()), - "unexpected rejection message" - ); - } - let success_item = function_outputs - .iter() - .find(|item| item.get("call_id").and_then(Value::as_str) == Some(call_id_success)) - .expect("success output present"); + let success_item = third_mock + .single_request() + .function_call_output(call_id_success); let output_json: Value = serde_json::from_str( success_item .get("output") @@ -282,18 +246,23 @@ async fn local_shell_missing_ids_maps_to_function_output_error() -> Result<()> { } }); - let responses = vec![ + mount_sse_once( + &server, sse(vec![ ev_response_created("resp-1"), local_shell_event, ev_completed("resp-1"), ]), + ) + .await; + let second_mock = mount_sse_once( + &server, sse(vec![ ev_assistant_message("msg-1", "done"), ev_completed("resp-2"), ]), - ]; - mount_sse_sequence(&server, responses).await; + ) + .await; submit_turn( &test, @@ -303,15 +272,7 @@ async fn local_shell_missing_ids_maps_to_function_output_error() -> Result<()> { ) .await?; - let requests = server.received_requests().await.expect("recorded requests"); - let bodies = request_bodies(&requests)?; - let function_outputs = collect_output_items(&bodies, "function_call_output"); - assert_eq!( - function_outputs.len(), - 1, - "expected a single function output" - ); - let item = function_outputs[0]; + let item = second_mock.single_request().function_call_output(""); assert_eq!(item.get("call_id").and_then(Value::as_str), Some("")); assert_eq!( item.get("output").and_then(Value::as_str), @@ -329,7 +290,7 @@ async fn collect_tools(use_unified_exec: bool) -> Result> { ev_assistant_message("msg-1", "done"), ev_completed("resp-1"), ])]; - mount_sse_sequence(&server, responses).await; + let mock = mount_sse_sequence(&server, responses).await; let mut builder = test_codex().with_config(move |config| { config.use_experimental_unified_exec_tool = use_unified_exec; @@ -344,15 +305,8 @@ async fn collect_tools(use_unified_exec: bool) -> Result> { ) .await?; - let requests = server.received_requests().await.expect("recorded requests"); - assert_eq!( - requests.len(), - 1, - "expected a single request for tools collection" - ); - let bodies = request_bodies(&requests)?; - let first_body = bodies.first().expect("request body present"); - Ok(tool_names(first_body)) + let first_body = mock.single_request().body_json(); + Ok(tool_names(&first_body)) } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -392,18 +346,23 @@ async fn shell_timeout_includes_timeout_prefix_and_metadata() -> Result<()> { "timeout_ms": timeout_ms, }); - let responses = vec![ + mount_sse_once( + &server, sse(vec![ ev_response_created("resp-1"), ev_function_call(call_id, "shell", &serde_json::to_string(&args)?), ev_completed("resp-1"), ]), + ) + .await; + let second_mock = mount_sse_once( + &server, sse(vec![ ev_assistant_message("msg-1", "done"), ev_completed("resp-2"), ]), - ]; - mount_sse_sequence(&server, responses).await; + ) + .await; submit_turn( &test, @@ -413,13 +372,7 @@ async fn shell_timeout_includes_timeout_prefix_and_metadata() -> Result<()> { ) .await?; - let requests = server.received_requests().await.expect("recorded requests"); - let bodies = request_bodies(&requests)?; - let function_outputs = collect_output_items(&bodies, "function_call_output"); - let timeout_item = function_outputs - .iter() - .find(|item| item.get("call_id").and_then(Value::as_str) == Some(call_id)) - .expect("timeout output present"); + let timeout_item = second_mock.single_request().function_call_output(call_id); let output_str = timeout_item .get("output") @@ -437,19 +390,9 @@ async fn shell_timeout_includes_timeout_prefix_and_metadata() -> Result<()> { ); let stdout = output_json["output"].as_str().unwrap_or_default(); - let timeout_pattern = r"(?s)^Total output lines: \d+ - -command timed out after (?P\d+) milliseconds -line -.*$"; - let captures = assert_regex_match(timeout_pattern, stdout); - let duration_ms = captures - .name("ms") - .and_then(|m| m.as_str().parse::().ok()) - .unwrap_or_default(); assert!( - duration_ms >= timeout_ms, - "expected duration >= configured timeout, got {duration_ms} (timeout {timeout_ms})" + stdout.contains("command timed out"), + "timeout output missing `command timed out`: {stdout}" ); } else { // Fallback: accept the signal classification path to deflake the test. @@ -478,18 +421,23 @@ async fn shell_sandbox_denied_truncates_error_output() -> Result<()> { "timeout_ms": 1_000, }); - let responses = vec![ + mount_sse_once( + &server, sse(vec![ ev_response_created("resp-1"), ev_function_call(call_id, "shell", &serde_json::to_string(&args)?), ev_completed("resp-1"), ]), + ) + .await; + let second_mock = mount_sse_once( + &server, sse(vec![ ev_assistant_message("msg-1", "done"), ev_completed("resp-2"), ]), - ]; - mount_sse_sequence(&server, responses).await; + ) + .await; submit_turn( &test, @@ -499,13 +447,7 @@ async fn shell_sandbox_denied_truncates_error_output() -> Result<()> { ) .await?; - let requests = server.received_requests().await.expect("recorded requests"); - let bodies = request_bodies(&requests)?; - let function_outputs = collect_output_items(&bodies, "function_call_output"); - let denied_item = function_outputs - .iter() - .find(|item| item.get("call_id").and_then(Value::as_str) == Some(call_id)) - .expect("denied output present"); + let denied_item = second_mock.single_request().function_call_output(call_id); let output = denied_item .get("output") @@ -558,18 +500,23 @@ async fn shell_spawn_failure_truncates_exec_error() -> Result<()> { "timeout_ms": 1_000, }); - let responses = vec![ + mount_sse_once( + &server, sse(vec![ ev_response_created("resp-1"), ev_function_call(call_id, "shell", &serde_json::to_string(&args)?), ev_completed("resp-1"), ]), + ) + .await; + let second_mock = mount_sse_once( + &server, sse(vec![ ev_assistant_message("msg-1", "done"), ev_completed("resp-2"), ]), - ]; - mount_sse_sequence(&server, responses).await; + ) + .await; submit_turn( &test, @@ -579,13 +526,7 @@ async fn shell_spawn_failure_truncates_exec_error() -> Result<()> { ) .await?; - let requests = server.received_requests().await.expect("recorded requests"); - let bodies = request_bodies(&requests)?; - let function_outputs = collect_output_items(&bodies, "function_call_output"); - let failure_item = function_outputs - .iter() - .find(|item| item.get("call_id").and_then(Value::as_str) == Some(call_id)) - .expect("spawn failure output present"); + let failure_item = second_mock.single_request().function_call_output(call_id); let output = failure_item .get("output") diff --git a/codex-rs/core/tests/suite/view_image.rs b/codex-rs/core/tests/suite/view_image.rs index 16913f4f93..bdb67ad631 100644 --- a/codex-rs/core/tests/suite/view_image.rs +++ b/codex-rs/core/tests/suite/view_image.rs @@ -22,16 +22,6 @@ use core_test_support::wait_for_event; use serde_json::Value; use wiremock::matchers::any; -fn function_call_output(body: &Value) -> Option<&Value> { - body.get("input") - .and_then(Value::as_array) - .and_then(|items| { - items.iter().find(|item| { - item.get("type").and_then(Value::as_str) == Some("function_call_output") - }) - }) -} - fn find_image_message(body: &Value) -> Option<&Value> { body.get("input") .and_then(Value::as_array) @@ -59,12 +49,6 @@ fn extract_output_text(item: &Value) -> Option<&str> { }) } -fn find_request_with_function_call_output(requests: &[Value]) -> Option<&Value> { - requests - .iter() - .find(|body| function_call_output(body).is_some()) -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn view_image_tool_attaches_local_image() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); @@ -100,7 +84,7 @@ async fn view_image_tool_attaches_local_image() -> anyhow::Result<()> { ev_assistant_message("msg-1", "done"), ev_completed("resp-2"), ]); - responses::mount_sse_once_match(&server, any(), second_response).await; + let mock = responses::mount_sse_once_match(&server, any(), second_response).await; let session_model = session_configured.model.clone(); @@ -137,25 +121,14 @@ async fn view_image_tool_attaches_local_image() -> anyhow::Result<()> { assert_eq!(tool_event.call_id, call_id); assert_eq!(tool_event.path, abs_path); - let requests = server.received_requests().await.expect("recorded requests"); - assert!( - requests.len() >= 2, - "expected at least two POST requests, got {}", - requests.len() - ); - let request_bodies = requests - .iter() - .map(|req| req.body_json::().expect("request json")) - .collect::>(); + let body = mock.single_request().body_json(); + let output_item = mock.single_request().function_call_output(call_id); - let body_with_tool_output = find_request_with_function_call_output(&request_bodies) - .expect("function_call_output item not found in requests"); - let output_item = function_call_output(body_with_tool_output).expect("tool output item"); - let output_text = extract_output_text(output_item).expect("output text present"); + let output_text = extract_output_text(&output_item).expect("output text present"); assert_eq!(output_text, "attached local image path"); - let image_message = find_image_message(body_with_tool_output) - .expect("pending input image message not included in request"); + let image_message = + find_image_message(&body).expect("pending input image message not included in request"); let image_url = image_message .get("content") .and_then(Value::as_array) @@ -210,7 +183,7 @@ async fn view_image_tool_errors_when_path_is_directory() -> anyhow::Result<()> { ev_assistant_message("msg-1", "done"), ev_completed("resp-2"), ]); - responses::mount_sse_once_match(&server, any(), second_response).await; + let mock = responses::mount_sse_once_match(&server, any(), second_response).await; let session_model = session_configured.model.clone(); @@ -231,26 +204,14 @@ async fn view_image_tool_errors_when_path_is_directory() -> anyhow::Result<()> { wait_for_event(&codex, |event| matches!(event, EventMsg::TaskComplete(_))).await; - let requests = server.received_requests().await.expect("recorded requests"); - assert!( - requests.len() >= 2, - "expected at least two POST requests, got {}", - requests.len() - ); - let request_bodies = requests - .iter() - .map(|req| req.body_json::().expect("request json")) - .collect::>(); - - let body_with_tool_output = find_request_with_function_call_output(&request_bodies) - .expect("function_call_output item not found in requests"); - let output_item = function_call_output(body_with_tool_output).expect("tool output item"); - let output_text = extract_output_text(output_item).expect("output text present"); + let body_with_tool_output = mock.single_request().body_json(); + let output_item = mock.single_request().function_call_output(call_id); + let output_text = extract_output_text(&output_item).expect("output text present"); let expected_message = format!("image path `{}` is not a file", abs_path.display()); assert_eq!(output_text, expected_message); assert!( - find_image_message(body_with_tool_output).is_none(), + find_image_message(&body_with_tool_output).is_none(), "directory path should not produce an input_image message" ); @@ -287,7 +248,7 @@ async fn view_image_tool_errors_when_file_missing() -> anyhow::Result<()> { ev_assistant_message("msg-1", "done"), ev_completed("resp-2"), ]); - responses::mount_sse_once_match(&server, any(), second_response).await; + let mock = responses::mount_sse_once_match(&server, any(), second_response).await; let session_model = session_configured.model.clone(); @@ -308,21 +269,9 @@ async fn view_image_tool_errors_when_file_missing() -> anyhow::Result<()> { wait_for_event(&codex, |event| matches!(event, EventMsg::TaskComplete(_))).await; - let requests = server.received_requests().await.expect("recorded requests"); - assert!( - requests.len() >= 2, - "expected at least two POST requests, got {}", - requests.len() - ); - let request_bodies = requests - .iter() - .map(|req| req.body_json::().expect("request json")) - .collect::>(); - - let body_with_tool_output = find_request_with_function_call_output(&request_bodies) - .expect("function_call_output item not found in requests"); - let output_item = function_call_output(body_with_tool_output).expect("tool output item"); - let output_text = extract_output_text(output_item).expect("output text present"); + let body_with_tool_output = mock.single_request().body_json(); + let output_item = mock.single_request().function_call_output(call_id); + let output_text = extract_output_text(&output_item).expect("output text present"); let expected_prefix = format!("unable to locate image at `{}`:", abs_path.display()); assert!( output_text.starts_with(&expected_prefix), @@ -330,7 +279,7 @@ async fn view_image_tool_errors_when_file_missing() -> anyhow::Result<()> { ); assert!( - find_image_message(body_with_tool_output).is_none(), + find_image_message(&body_with_tool_output).is_none(), "missing file should not produce an input_image message" ); diff --git a/codex-rs/docs/codex_mcp_interface.md b/codex-rs/docs/codex_mcp_interface.md index 1291d3ee9f..eb0442feb3 100644 --- a/codex-rs/docs/codex_mcp_interface.md +++ b/codex-rs/docs/codex_mcp_interface.md @@ -1,6 +1,6 @@ -# Codex MCP Interface [experimental] +# Codex MCP Server Interface [experimental] -This document describes Codex’s experimental MCP interface: a JSON‑RPC API that runs over the Model Context Protocol (MCP) transport to control a local Codex engine. +This document describes Codex’s experimental MCP server interface: a JSON‑RPC API that runs over the Model Context Protocol (MCP) transport to control a local Codex engine. - Status: experimental and subject to change without notice - Server binary: `codex mcp-server` (or `codex-mcp-server`) diff --git a/codex-rs/exec/src/cli.rs b/codex-rs/exec/src/cli.rs index 323cc879fe..d264eb61b0 100644 --- a/codex-rs/exec/src/cli.rs +++ b/codex-rs/exec/src/cli.rs @@ -77,7 +77,7 @@ pub struct Cli { /// Initial instructions for the agent. If not provided as an argument (or /// if `-` is used), instructions are read from stdin. - #[arg(value_name = "PROMPT")] + #[arg(value_name = "PROMPT", value_hint = clap::ValueHint::Other)] pub prompt: Option, } @@ -99,7 +99,7 @@ pub struct ResumeArgs { pub last: bool, /// Prompt to send after resuming the session. If `-` is used, read from stdin. - #[arg(value_name = "PROMPT")] + #[arg(value_name = "PROMPT", value_hint = clap::ValueHint::Other)] pub prompt: Option, } diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index 18e097782d..967da52b8a 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -48,7 +48,7 @@ use codex_core::default_client::set_default_originator; use codex_core::find_conversation_path_by_id_str; pub async fn run_main(cli: Cli, codex_linux_sandbox_exe: Option) -> anyhow::Result<()> { - if let Err(err) = set_default_originator("codex_exec") { + if let Err(err) = set_default_originator("codex_exec".to_string()) { tracing::warn!(?err, "Failed to set codex exec originator override {err:?}"); } diff --git a/codex-rs/exec/tests/suite/mod.rs b/codex-rs/exec/tests/suite/mod.rs index d04ecd2c93..052c43bf38 100644 --- a/codex-rs/exec/tests/suite/mod.rs +++ b/codex-rs/exec/tests/suite/mod.rs @@ -1,6 +1,7 @@ // Aggregates all former standalone integration tests as modules. mod apply_patch; mod auth_env; +mod originator; mod output_schema; mod resume; mod sandbox; diff --git a/codex-rs/exec/tests/suite/originator.rs b/codex-rs/exec/tests/suite/originator.rs new file mode 100644 index 0000000000..cf0954447a --- /dev/null +++ b/codex-rs/exec/tests/suite/originator.rs @@ -0,0 +1,52 @@ +#![cfg(not(target_os = "windows"))] +#![allow(clippy::expect_used, clippy::unwrap_used)] + +use core_test_support::responses; +use core_test_support::test_codex_exec::test_codex_exec; +use wiremock::matchers::header; + +/// Verify that when the server reports an error, `codex-exec` exits with a +/// non-zero status code so automation can detect failures. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn send_codex_exec_originator() -> anyhow::Result<()> { + let test = test_codex_exec(); + + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("response_1"), + responses::ev_assistant_message("response_1", "Hello, world!"), + responses::ev_completed("response_1"), + ]); + responses::mount_sse_once_match(&server, header("Originator", "codex_exec"), body).await; + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("tell me something") + .assert() + .code(0); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn supports_originator_override() -> anyhow::Result<()> { + let test = test_codex_exec(); + + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("response_1"), + responses::ev_assistant_message("response_1", "Hello, world!"), + responses::ev_completed("response_1"), + ]); + responses::mount_sse_once_match(&server, header("Originator", "codex_exec_override"), body) + .await; + + test.cmd_with_server(&server) + .env("CODEX_INTERNAL_ORIGINATOR_OVERRIDE", "codex_exec_override") + .arg("--skip-git-repo-check") + .arg("tell me something") + .assert() + .code(0); + + Ok(()) +} diff --git a/codex-rs/exec/tests/suite/output_schema.rs b/codex-rs/exec/tests/suite/output_schema.rs index b054484f7c..913270eff4 100644 --- a/codex-rs/exec/tests/suite/output_schema.rs +++ b/codex-rs/exec/tests/suite/output_schema.rs @@ -28,7 +28,7 @@ async fn exec_includes_output_schema_in_request() -> anyhow::Result<()> { responses::ev_assistant_message("m1", "fixture hello"), responses::ev_completed("resp1"), ]); - responses::mount_sse_once_match(&server, any(), body).await; + let response_mock = responses::mount_sse_once_match(&server, any(), body).await; test.cmd_with_server(&server) .arg("--skip-git-repo-check") @@ -43,12 +43,8 @@ async fn exec_includes_output_schema_in_request() -> anyhow::Result<()> { .assert() .success(); - let requests = server - .received_requests() - .await - .expect("failed to capture requests"); - assert_eq!(requests.len(), 1, "expected exactly one request"); - let payload: Value = serde_json::from_slice(&requests[0].body)?; + let request = response_mock.single_request(); + let payload: Value = request.body_json(); let text = payload.get("text").expect("request missing text field"); let format = text .get("format") diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index bdd1d3e241..3b6520be67 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -548,10 +548,15 @@ pub struct TaskStartedEvent { #[derive(Debug, Clone, Deserialize, Serialize, Default, TS)] pub struct TokenUsage { + #[ts(type = "number")] pub input_tokens: u64, + #[ts(type = "number")] pub cached_input_tokens: u64, + #[ts(type = "number")] pub output_tokens: u64, + #[ts(type = "number")] pub reasoning_output_tokens: u64, + #[ts(type = "number")] pub total_tokens: u64, } @@ -559,6 +564,7 @@ pub struct TokenUsage { pub struct TokenUsageInfo { pub total_token_usage: TokenUsage, pub last_token_usage: TokenUsage, + #[ts(type = "number | null")] pub model_context_window: Option, } @@ -634,8 +640,10 @@ pub struct RateLimitWindow { /// Percentage (0-100) of the window that has been consumed. pub used_percent: f64, /// Rolling window duration, in minutes. + #[ts(type = "number | null")] pub window_minutes: Option, /// Seconds until the window resets. + #[ts(type = "number | null")] pub resets_in_seconds: Option, } diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index ac69a100e9..0d15584f0f 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -5,6 +5,7 @@ mod perform_oauth_login; mod rmcp_client; mod utils; +pub use oauth::OAuthCredentialsStoreMode; pub use oauth::StoredOAuthTokens; pub use oauth::WrappedOAuthTokenResponse; pub use oauth::delete_oauth_tokens; diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index bb13b718fc..0534847651 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -58,6 +58,21 @@ pub struct StoredOAuthTokens { pub token_response: WrappedOAuthTokenResponse, } +/// Determine where Codex should store and read MCP credentials. +#[derive(Debug, Default, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum OAuthCredentialsStoreMode { + /// `Keyring` when available; otherwise, `File`. + /// Credentials stored in the keyring will only be readable by Codex unless the user explicitly grants access via OS-level keyring access. + #[default] + Auto, + /// CODEX_HOME/.credentials.json + /// This file will be readable to Codex and other applications running as the same user. + File, + /// Keyring when available, otherwise fail. + Keyring, +} + #[derive(Debug)] struct CredentialStoreError(anyhow::Error); @@ -83,15 +98,15 @@ impl fmt::Display for CredentialStoreError { impl std::error::Error for CredentialStoreError {} -trait CredentialStore { +trait KeyringStore { fn load(&self, service: &str, account: &str) -> Result, CredentialStoreError>; fn save(&self, service: &str, account: &str, value: &str) -> Result<(), CredentialStoreError>; fn delete(&self, service: &str, account: &str) -> Result; } -struct KeyringCredentialStore; +struct DefaultKeyringStore; -impl CredentialStore for KeyringCredentialStore { +impl KeyringStore for DefaultKeyringStore { fn load(&self, service: &str, account: &str) -> Result, CredentialStoreError> { let entry = Entry::new(service, account).map_err(CredentialStoreError::new)?; match entry.get_password() { @@ -129,47 +144,85 @@ impl PartialEq for WrappedOAuthTokenResponse { } } -pub(crate) fn load_oauth_tokens(server_name: &str, url: &str) -> Result> { - let store = KeyringCredentialStore; - load_oauth_tokens_with_store(&store, server_name, url) +pub(crate) fn load_oauth_tokens( + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, +) -> Result> { + let keyring_store = DefaultKeyringStore; + match store_mode { + OAuthCredentialsStoreMode::Auto => { + load_oauth_tokens_from_keyring_with_fallback_to_file(&keyring_store, server_name, url) + } + OAuthCredentialsStoreMode::File => load_oauth_tokens_from_file(server_name, url), + OAuthCredentialsStoreMode::Keyring => { + load_oauth_tokens_from_keyring(&keyring_store, server_name, url) + .with_context(|| "failed to read OAuth tokens from keyring".to_string()) + } + } } -fn load_oauth_tokens_with_store( - store: &C, +fn load_oauth_tokens_from_keyring_with_fallback_to_file( + keyring_store: &K, + server_name: &str, + url: &str, +) -> Result> { + match load_oauth_tokens_from_keyring(keyring_store, server_name, url) { + Ok(Some(tokens)) => Ok(Some(tokens)), + Ok(None) => load_oauth_tokens_from_file(server_name, url), + Err(error) => { + warn!("failed to read OAuth tokens from keyring: {error}"); + load_oauth_tokens_from_file(server_name, url) + .with_context(|| format!("failed to read OAuth tokens from keyring: {error}")) + } + } +} + +fn load_oauth_tokens_from_keyring( + keyring_store: &K, server_name: &str, url: &str, ) -> Result> { let key = compute_store_key(server_name, url)?; - match store.load(KEYRING_SERVICE, &key) { + match keyring_store.load(KEYRING_SERVICE, &key) { Ok(Some(serialized)) => { let tokens: StoredOAuthTokens = serde_json::from_str(&serialized) .context("failed to deserialize OAuth tokens from keyring")?; Ok(Some(tokens)) } - Ok(None) => load_oauth_tokens_from_file(server_name, url), - Err(error) => { - let message = error.message(); - warn!("failed to read OAuth tokens from keyring: {message}"); - load_oauth_tokens_from_file(server_name, url) - .with_context(|| format!("failed to read OAuth tokens from keyring: {message}")) + Ok(None) => Ok(None), + Err(error) => Err(error.into_error()), + } +} + +pub fn save_oauth_tokens( + server_name: &str, + tokens: &StoredOAuthTokens, + store_mode: OAuthCredentialsStoreMode, +) -> Result<()> { + let keyring_store = DefaultKeyringStore; + match store_mode { + OAuthCredentialsStoreMode::Auto => save_oauth_tokens_with_keyring_with_fallback_to_file( + &keyring_store, + server_name, + tokens, + ), + OAuthCredentialsStoreMode::File => save_oauth_tokens_to_file(tokens), + OAuthCredentialsStoreMode::Keyring => { + save_oauth_tokens_with_keyring(&keyring_store, server_name, tokens) } } } -pub fn save_oauth_tokens(server_name: &str, tokens: &StoredOAuthTokens) -> Result<()> { - let store = KeyringCredentialStore; - save_oauth_tokens_with_store(&store, server_name, tokens) -} - -fn save_oauth_tokens_with_store( - store: &C, +fn save_oauth_tokens_with_keyring( + keyring_store: &K, server_name: &str, tokens: &StoredOAuthTokens, ) -> Result<()> { let serialized = serde_json::to_string(tokens).context("failed to serialize OAuth tokens")?; let key = compute_store_key(server_name, &tokens.url)?; - match store.save(KEYRING_SERVICE, &key, &serialized) { + match keyring_store.save(KEYRING_SERVICE, &key, &serialized) { Ok(()) => { if let Err(error) = delete_oauth_tokens_from_file(&key) { warn!("failed to remove OAuth tokens from fallback storage: {error:?}"); @@ -177,31 +230,61 @@ fn save_oauth_tokens_with_store( Ok(()) } Err(error) => { - let message = error.message(); - warn!("failed to write OAuth tokens to keyring: {message}"); + let message = format!( + "failed to write OAuth tokens to keyring: {}", + error.message() + ); + warn!("{message}"); + Err(error.into_error().context(message)) + } + } +} + +fn save_oauth_tokens_with_keyring_with_fallback_to_file( + keyring_store: &K, + server_name: &str, + tokens: &StoredOAuthTokens, +) -> Result<()> { + match save_oauth_tokens_with_keyring(keyring_store, server_name, tokens) { + Ok(()) => Ok(()), + Err(error) => { + let message = error.to_string(); + warn!("falling back to file storage for OAuth tokens: {message}"); save_oauth_tokens_to_file(tokens) .with_context(|| format!("failed to write OAuth tokens to keyring: {message}")) } } } -pub fn delete_oauth_tokens(server_name: &str, url: &str) -> Result { - let store = KeyringCredentialStore; - delete_oauth_tokens_with_store(&store, server_name, url) +pub fn delete_oauth_tokens( + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, +) -> Result { + let keyring_store = DefaultKeyringStore; + delete_oauth_tokens_from_keyring_and_file(&keyring_store, store_mode, server_name, url) } -fn delete_oauth_tokens_with_store( - store: &C, +fn delete_oauth_tokens_from_keyring_and_file( + keyring_store: &K, + store_mode: OAuthCredentialsStoreMode, server_name: &str, url: &str, ) -> Result { let key = compute_store_key(server_name, url)?; - let keyring_removed = match store.delete(KEYRING_SERVICE, &key) { + let keyring_result = keyring_store.delete(KEYRING_SERVICE, &key); + let keyring_removed = match keyring_result { Ok(removed) => removed, Err(error) => { let message = error.message(); warn!("failed to delete OAuth tokens from keyring: {message}"); - return Err(error.into_error()).context("failed to delete OAuth tokens from keyring"); + match store_mode { + OAuthCredentialsStoreMode::Auto | OAuthCredentialsStoreMode::Keyring => { + return Err(error.into_error()) + .context("failed to delete OAuth tokens from keyring"); + } + OAuthCredentialsStoreMode::File => false, + } } }; @@ -218,6 +301,7 @@ struct OAuthPersistorInner { server_name: String, url: String, authorization_manager: Arc>, + store_mode: OAuthCredentialsStoreMode, last_credentials: Mutex>, } @@ -225,14 +309,16 @@ impl OAuthPersistor { pub(crate) fn new( server_name: String, url: String, - manager: Arc>, + authorization_manager: Arc>, + store_mode: OAuthCredentialsStoreMode, initial_credentials: Option, ) -> Self { Self { inner: Arc::new(OAuthPersistorInner { server_name, url, - authorization_manager: manager, + authorization_manager, + store_mode, last_credentials: Mutex::new(initial_credentials), }), } @@ -257,15 +343,18 @@ impl OAuthPersistor { }; let mut last_credentials = self.inner.last_credentials.lock().await; if last_credentials.as_ref() != Some(&stored) { - save_oauth_tokens(&self.inner.server_name, &stored)?; + save_oauth_tokens(&self.inner.server_name, &stored, self.inner.store_mode)?; *last_credentials = Some(stored); } } None => { let mut last_serialized = self.inner.last_credentials.lock().await; if last_serialized.take().is_some() - && let Err(error) = - delete_oauth_tokens(&self.inner.server_name, &self.inner.url) + && let Err(error) = delete_oauth_tokens( + &self.inner.server_name, + &self.inner.url, + self.inner.store_mode, + ) { warn!( "failed to remove OAuth tokens for server {}: {error}", @@ -542,7 +631,7 @@ mod tests { } } - impl CredentialStore for MockCredentialStore { + impl KeyringStore for MockCredentialStore { fn load( &self, _service: &str, @@ -643,7 +732,8 @@ mod tests { let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; store.save(KEYRING_SERVICE, &key, &serialized)?; - let loaded = super::load_oauth_tokens_with_store(&store, &tokens.server_name, &tokens.url)?; + let loaded = + super::load_oauth_tokens_from_keyring(&store, &tokens.server_name, &tokens.url)?; assert_eq!(loaded, Some(expected)); Ok(()) } @@ -657,8 +747,12 @@ mod tests { super::save_oauth_tokens_to_file(&tokens)?; - let loaded = super::load_oauth_tokens_with_store(&store, &tokens.server_name, &tokens.url)? - .expect("tokens should load from fallback"); + let loaded = super::load_oauth_tokens_from_keyring_with_fallback_to_file( + &store, + &tokens.server_name, + &tokens.url, + )? + .expect("tokens should load from fallback"); assert_tokens_match_without_expiry(&loaded, &expected); Ok(()) } @@ -674,8 +768,12 @@ mod tests { super::save_oauth_tokens_to_file(&tokens)?; - let loaded = super::load_oauth_tokens_with_store(&store, &tokens.server_name, &tokens.url)? - .expect("tokens should load from fallback"); + let loaded = super::load_oauth_tokens_from_keyring_with_fallback_to_file( + &store, + &tokens.server_name, + &tokens.url, + )? + .expect("tokens should load from fallback"); assert_tokens_match_without_expiry(&loaded, &expected); Ok(()) } @@ -689,7 +787,11 @@ mod tests { super::save_oauth_tokens_to_file(&tokens)?; - super::save_oauth_tokens_with_store(&store, &tokens.server_name, &tokens)?; + super::save_oauth_tokens_with_keyring_with_fallback_to_file( + &store, + &tokens.server_name, + &tokens, + )?; let fallback_path = super::fallback_file_path()?; assert!(!fallback_path.exists(), "fallback file should be removed"); @@ -706,7 +808,11 @@ mod tests { let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; store.set_error(&key, KeyringError::Invalid("error".into(), "save".into())); - super::save_oauth_tokens_with_store(&store, &tokens.server_name, &tokens)?; + super::save_oauth_tokens_with_keyring_with_fallback_to_file( + &store, + &tokens.server_name, + &tokens, + )?; let fallback_path = super::fallback_file_path()?; assert!(fallback_path.exists(), "fallback file should be created"); @@ -734,8 +840,34 @@ mod tests { store.save(KEYRING_SERVICE, &key, &serialized)?; super::save_oauth_tokens_to_file(&tokens)?; - let removed = - super::delete_oauth_tokens_with_store(&store, &tokens.server_name, &tokens.url)?; + let removed = super::delete_oauth_tokens_from_keyring_and_file( + &store, + OAuthCredentialsStoreMode::Auto, + &tokens.server_name, + &tokens.url, + )?; + assert!(removed); + assert!(!store.contains(&key)); + assert!(!super::fallback_file_path()?.exists()); + Ok(()) + } + + #[test] + fn delete_oauth_tokens_file_mode_removes_keyring_only_entry() -> Result<()> { + let _env = TempCodexHome::new(); + let store = MockCredentialStore::default(); + let tokens = sample_tokens(); + let serialized = serde_json::to_string(&tokens)?; + let key = super::compute_store_key(&tokens.server_name, &tokens.url)?; + store.save(KEYRING_SERVICE, &key, &serialized)?; + assert!(store.contains(&key)); + + let removed = super::delete_oauth_tokens_from_keyring_and_file( + &store, + OAuthCredentialsStoreMode::Auto, + &tokens.server_name, + &tokens.url, + )?; assert!(removed); assert!(!store.contains(&key)); assert!(!super::fallback_file_path()?.exists()); @@ -751,8 +883,12 @@ mod tests { store.set_error(&key, KeyringError::Invalid("error".into(), "delete".into())); super::save_oauth_tokens_to_file(&tokens).unwrap(); - let result = - super::delete_oauth_tokens_with_store(&store, &tokens.server_name, &tokens.url); + let result = super::delete_oauth_tokens_from_keyring_and_file( + &store, + OAuthCredentialsStoreMode::Auto, + &tokens.server_name, + &tokens.url, + ); assert!(result.is_err()); assert!(super::fallback_file_path().unwrap().exists()); Ok(()) diff --git a/codex-rs/rmcp-client/src/perform_oauth_login.rs b/codex-rs/rmcp-client/src/perform_oauth_login.rs index b5a893611a..c2d39a214f 100644 --- a/codex-rs/rmcp-client/src/perform_oauth_login.rs +++ b/codex-rs/rmcp-client/src/perform_oauth_login.rs @@ -12,6 +12,7 @@ use tokio::sync::oneshot; use tokio::time::timeout; use urlencoding::decode; +use crate::OAuthCredentialsStoreMode; use crate::StoredOAuthTokens; use crate::WrappedOAuthTokenResponse; use crate::save_oauth_tokens; @@ -26,7 +27,11 @@ impl Drop for CallbackServerGuard { } } -pub async fn perform_oauth_login(server_name: &str, server_url: &str) -> Result<()> { +pub async fn perform_oauth_login( + server_name: &str, + server_url: &str, + store_mode: OAuthCredentialsStoreMode, +) -> Result<()> { let server = Arc::new(Server::http("127.0.0.1:0").map_err(|err| anyhow!(err))?); let guard = CallbackServerGuard { server: Arc::clone(&server), @@ -81,7 +86,7 @@ pub async fn perform_oauth_login(server_name: &str, server_url: &str) -> Result< client_id, token_response: WrappedOAuthTokenResponse(credentials), }; - save_oauth_tokens(server_name, &stored)?; + save_oauth_tokens(server_name, &stored, store_mode)?; drop(guard); Ok(()) diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index d19bf8fe72..3d12e50868 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -35,6 +35,7 @@ use tracing::warn; use crate::load_oauth_tokens; use crate::logging_client_handler::LoggingClientHandler; +use crate::oauth::OAuthCredentialsStoreMode; use crate::oauth::OAuthPersistor; use crate::oauth::StoredOAuthTokens; use crate::utils::convert_call_tool_result; @@ -119,17 +120,22 @@ impl RmcpClient { server_name: &str, url: &str, bearer_token: Option, + store_mode: OAuthCredentialsStoreMode, ) -> Result { - let initial_tokens = match load_oauth_tokens(server_name, url) { - Ok(tokens) => tokens, - Err(err) => { - warn!("failed to read tokens for server `{server_name}`: {err}"); - None - } + let initial_oauth_tokens = match bearer_token { + Some(_) => None, + None => match load_oauth_tokens(server_name, url, store_mode) { + Ok(tokens) => tokens, + Err(err) => { + warn!("failed to read tokens for server `{server_name}`: {err}"); + None + } + }, }; - let transport = if let Some(initial_tokens) = initial_tokens.clone() { + let transport = if let Some(initial_tokens) = initial_oauth_tokens.clone() { let (transport, oauth_persistor) = - create_oauth_transport_and_runtime(server_name, url, initial_tokens).await?; + create_oauth_transport_and_runtime(server_name, url, initial_tokens, store_mode) + .await?; PendingTransport::StreamableHttpWithOAuth { transport, oauth_persistor, @@ -137,7 +143,7 @@ impl RmcpClient { } else { let mut http_config = StreamableHttpClientTransportConfig::with_uri(url.to_string()); if let Some(bearer_token) = bearer_token { - http_config = http_config.auth_header(format!("Bearer {bearer_token}")); + http_config = http_config.auth_header(bearer_token); } let transport = StreamableHttpClientTransport::from_config(http_config); @@ -283,6 +289,7 @@ async fn create_oauth_transport_and_runtime( server_name: &str, url: &str, initial_tokens: StoredOAuthTokens, + credentials_store: OAuthCredentialsStoreMode, ) -> Result<( StreamableHttpClientTransport>, OAuthPersistor, @@ -317,6 +324,7 @@ async fn create_oauth_transport_and_runtime( server_name.to_string(), url.to_string(), auth_manager, + credentials_store, Some(initial_tokens), ); diff --git a/codex-rs/tui/Cargo.toml b/codex-rs/tui/Cargo.toml index f42f555c9e..2c6f32a20a 100644 --- a/codex-rs/tui/Cargo.toml +++ b/codex-rs/tui/Cargo.toml @@ -68,6 +68,8 @@ strum_macros = { workspace = true } supports-color = { workspace = true } tempfile = { workspace = true } textwrap = { workspace = true } +tree-sitter-highlight = { workspace = true } +tree-sitter-bash = { workspace = true } tokio = { workspace = true, features = [ "io-std", "macros", diff --git a/codex-rs/tui/src/app_backtrack.rs b/codex-rs/tui/src/app_backtrack.rs index b5c1300b45..5528fa7bc4 100644 --- a/codex-rs/tui/src/app_backtrack.rs +++ b/codex-rs/tui/src/app_backtrack.rs @@ -134,8 +134,9 @@ impl App { /// Useful when switching sessions to ensure prior history remains visible. pub(crate) fn render_transcript_once(&mut self, tui: &mut tui::Tui) { if !self.transcript_cells.is_empty() { + let width = tui.terminal.last_known_screen_size.width; for cell in &self.transcript_cells { - tui.insert_history_lines(cell.transcript_lines()); + tui.insert_history_lines(cell.display_lines(width)); } } } diff --git a/codex-rs/tui/src/bottom_pane/approval_overlay.rs b/codex-rs/tui/src/bottom_pane/approval_overlay.rs index 7f52b11bfa..6585124c52 100644 --- a/codex-rs/tui/src/bottom_pane/approval_overlay.rs +++ b/codex-rs/tui/src/bottom_pane/approval_overlay.rs @@ -315,15 +315,16 @@ impl From for ApprovalRequestState { changes, } => { let mut header: Vec> = Vec::new(); - header.push(DiffSummary::new(changes, cwd).into()); if let Some(reason) = reason && !reason.is_empty() { - header.push(Box::new(Line::from(""))); header.push(Box::new( - Paragraph::new(reason.italic()).wrap(Wrap { trim: false }), + Paragraph::new(Line::from_iter(["Reason: ".into(), reason.italic()])) + .wrap(Wrap { trim: false }), )); + header.push(Box::new(Line::from(""))); } + header.push(DiffSummary::new(changes, cwd).into()); Self { variant: ApprovalVariant::ApplyPatch { id }, header: Box::new(ColumnRenderable::new(header)), diff --git a/codex-rs/tui/src/bottom_pane/chat_composer.rs b/codex-rs/tui/src/bottom_pane/chat_composer.rs index ad9770726a..2574ded461 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer.rs @@ -38,7 +38,6 @@ use crate::bottom_pane::prompt_args::prompt_has_numeric_placeholders; use crate::slash_command::SlashCommand; use crate::slash_command::built_in_slash_commands; use crate::style::user_message_style; -use crate::terminal_palette; use codex_protocol::custom_prompts::CustomPrompt; use codex_protocol::custom_prompts::PROMPTS_CMD_PREFIX; @@ -1533,7 +1532,7 @@ impl WidgetRef for ChatComposer { } } } - let style = user_message_style(terminal_palette::default_bg()); + let style = user_message_style(); let mut block_rect = composer_rect; block_rect.y = composer_rect.y.saturating_sub(1); block_rect.height = composer_rect.height.saturating_add(1); diff --git a/codex-rs/tui/src/bottom_pane/list_selection_view.rs b/codex-rs/tui/src/bottom_pane/list_selection_view.rs index 9f0ce3df24..238f74d6e4 100644 --- a/codex-rs/tui/src/bottom_pane/list_selection_view.rs +++ b/codex-rs/tui/src/bottom_pane/list_selection_view.rs @@ -20,7 +20,6 @@ use crate::render::RectExt as _; use crate::render::renderable::ColumnRenderable; use crate::render::renderable::Renderable; use crate::style::user_message_style; -use crate::terminal_palette; use super::CancellationEvent; use super::bottom_pane_view::BottomPaneView; @@ -350,7 +349,7 @@ impl Renderable for ListSelectionView { .areas(area); Block::default() - .style(user_message_style(terminal_palette::default_bg())) + .style(user_message_style()) .render(content_area, buf); let header_height = self diff --git a/codex-rs/tui/src/bottom_pane/mod.rs b/codex-rs/tui/src/bottom_pane/mod.rs index 6fe673a23b..db13a041cd 100644 --- a/codex-rs/tui/src/bottom_pane/mod.rs +++ b/codex-rs/tui/src/bottom_pane/mod.rs @@ -81,7 +81,7 @@ pub(crate) struct BottomPaneParams { } impl BottomPane { - const BOTTOM_PAD_LINES: u16 = 1; + const BOTTOM_PAD_LINES: u16 = 0; pub fn new(params: BottomPaneParams) -> Self { let enhanced_keys_supported = params.enhanced_keys_supported; Self { @@ -522,10 +522,29 @@ impl WidgetRef for &BottomPane { mod tests { use super::*; use crate::app_event::AppEvent; + use insta::assert_snapshot; use ratatui::buffer::Buffer; use ratatui::layout::Rect; use tokio::sync::mpsc::unbounded_channel; + fn snapshot_buffer(buf: &Buffer) -> String { + let mut lines = Vec::new(); + for y in 0..buf.area().height { + let mut row = String::new(); + for x in 0..buf.area().width { + row.push(buf[(x, y)].symbol().chars().next().unwrap_or(' ')); + } + lines.push(row); + } + lines.join("\n") + } + + fn render_snapshot(pane: &BottomPane, area: Rect) -> String { + let mut buf = Buffer::empty(area); + (&pane).render_ref(area, &mut buf); + snapshot_buffer(&buf) + } + fn exec_request() -> ApprovalRequest { ApprovalRequest::Exec { id: "1".to_string(), @@ -685,7 +704,7 @@ mod tests { } #[test] - fn bottom_padding_present_with_status_above_composer() { + fn status_and_composer_fill_height_without_bottom_padding() { let (tx_raw, _rx) = unbounded_channel::(); let tx = AppEventSender::new(tx_raw); let mut pane = BottomPane::new(BottomPaneParams { @@ -700,43 +719,21 @@ mod tests { // Activate spinner (status view replaces composer) with no live ring. pane.set_task_running(true); - // Use height == desired_height; expect 1 status row at top and 2 bottom padding rows. + // Use height == desired_height; expect spacer + status + composer rows without trailing padding. let height = pane.desired_height(30); assert!( height >= 3, - "expected at least 3 rows with bottom padding; got {height}" + "expected at least 3 rows to render spacer, status, and composer; got {height}" ); let area = Rect::new(0, 0, 30, height); - let mut buf = Buffer::empty(area); - (&pane).render_ref(area, &mut buf); - - // Row 1 contains the status header (row 0 is the spacer) - let mut top = String::new(); - for x in 0..area.width { - top.push(buf[(x, 1)].symbol().chars().next().unwrap_or(' ')); - } - assert!( - top.trim_start().starts_with("• Working"), - "expected top row to start with '• Working': {top:?}" - ); - assert!( - top.contains("Working"), - "expected Working header on top row: {top:?}" - ); - - // Last row should be blank padding; the row above should generally contain composer content. - let mut r_last = String::new(); - for x in 0..area.width { - r_last.push(buf[(x, height - 1)].symbol().chars().next().unwrap_or(' ')); - } - assert!( - r_last.trim().is_empty(), - "expected last row blank: {r_last:?}" + assert_snapshot!( + "status_and_composer_fill_height_without_bottom_padding", + render_snapshot(&pane, area) ); } #[test] - fn bottom_padding_shrinks_when_tiny() { + fn status_hidden_when_height_too_small() { let (tx_raw, _rx) = unbounded_channel::(); let tx = AppEventSender::new(tx_raw); let mut pane = BottomPane::new(BottomPaneParams { @@ -750,37 +747,18 @@ mod tests { pane.set_task_running(true); - // Height=2 → status on one row, composer on the other. + // Height=2 → composer takes the full space; status collapses when there is no room. let area2 = Rect::new(0, 0, 20, 2); - let mut buf2 = Buffer::empty(area2); - (&pane).render_ref(area2, &mut buf2); - let mut row0 = String::new(); - let mut row1 = String::new(); - for x in 0..area2.width { - row0.push(buf2[(x, 0)].symbol().chars().next().unwrap_or(' ')); - row1.push(buf2[(x, 1)].symbol().chars().next().unwrap_or(' ')); - } - let has_composer = row0.contains("Ask Codex") || row1.contains("Ask Codex"); - assert!( - has_composer, - "expected composer to be visible on one of the rows: row0={row0:?}, row1={row1:?}" - ); - assert!( - row0.contains("Working") || row1.contains("Working"), - "expected status header to be visible at height=2: row0={row0:?}, row1={row1:?}" + assert_snapshot!( + "status_hidden_when_height_too_small_height_2", + render_snapshot(&pane, area2) ); // Height=1 → no padding; single row is the composer (status hidden). let area1 = Rect::new(0, 0, 20, 1); - let mut buf1 = Buffer::empty(area1); - (&pane).render_ref(area1, &mut buf1); - let mut only = String::new(); - for x in 0..area1.width { - only.push(buf1[(x, 0)].symbol().chars().next().unwrap_or(' ')); - } - assert!( - only.contains("Ask Codex"), - "expected composer with no padding: {only:?}" + assert_snapshot!( + "status_hidden_when_height_too_small_height_1", + render_snapshot(&pane, area1) ); } } diff --git a/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__tests__status_and_composer_fill_height_without_bottom_padding.snap b/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__tests__status_and_composer_fill_height_without_bottom_padding.snap new file mode 100644 index 0000000000..ae96888f56 --- /dev/null +++ b/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__tests__status_and_composer_fill_height_without_bottom_padding.snap @@ -0,0 +1,11 @@ +--- +source: tui/src/bottom_pane/mod.rs +expression: "render_snapshot(&pane, area)" +--- + +• Working (0s • esc to interru + + +› Ask Codex to do anything + + ? for shortcuts diff --git a/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__tests__status_hidden_when_height_too_small_height_1.snap b/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__tests__status_hidden_when_height_too_small_height_1.snap new file mode 100644 index 0000000000..310c32b40c --- /dev/null +++ b/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__tests__status_hidden_when_height_too_small_height_1.snap @@ -0,0 +1,5 @@ +--- +source: tui/src/bottom_pane/mod.rs +expression: "render_snapshot(&pane, area1)" +--- +› Ask Codex to do an diff --git a/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__tests__status_hidden_when_height_too_small_height_2.snap b/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__tests__status_hidden_when_height_too_small_height_2.snap new file mode 100644 index 0000000000..ea0beeedf3 --- /dev/null +++ b/codex-rs/tui/src/bottom_pane/snapshots/codex_tui__bottom_pane__tests__status_hidden_when_height_too_small_height_2.snap @@ -0,0 +1,6 @@ +--- +source: tui/src/bottom_pane/mod.rs +expression: "render_snapshot(&pane, area2)" +--- + +› Ask Codex to do an diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__apply_patch_manual_flow_history_approved.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__apply_patch_manual_flow_history_approved.snap index 6e22bcebaa..e139b51088 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__apply_patch_manual_flow_history_approved.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__apply_patch_manual_flow_history_approved.snap @@ -3,4 +3,4 @@ source: tui/src/chatwidget/tests.rs expression: lines_to_single_string(&approved_lines) --- • Added foo.txt (+1 -0) - 1 +hello + 1 +hello diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__approval_modal_patch.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__approval_modal_patch.snap index ab88ffaf73..96dde8fb93 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__approval_modal_patch.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__approval_modal_patch.snap @@ -4,12 +4,12 @@ expression: terminal.backend().vt100().screen().contents() --- Would you like to make the following edits? + Reason: The model wants to apply changes + README.md (+2 -0) - 1 +hello - 2 +world - - The model wants to apply changes + 1 +hello + 2 +world › 1. Yes, proceed 2. No, and tell Codex what to do differently esc diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__chat_small_idle_h3.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__chat_small_idle_h3.snap index 60ef79386d..adb00d4bc6 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__chat_small_idle_h3.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__chat_small_idle_h3.snap @@ -1,7 +1,8 @@ --- source: tui/src/chatwidget/tests.rs +assertion_line: 1470 expression: terminal.backend() --- " " -"› Ask Codex to do anything " " " +"› Ask Codex to do anything " diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__chat_small_running_h2.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__chat_small_running_h2.snap index 630d1a3bbe..4e1b74a2c6 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__chat_small_running_h2.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__chat_small_running_h2.snap @@ -1,6 +1,7 @@ --- source: tui/src/chatwidget/tests.rs +assertion_line: 1500 expression: terminal.backend() --- -"• Thinking (0s • esc to interrupt) " +" " "› Ask Codex to do anything " diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__chat_small_running_h3.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__chat_small_running_h3.snap index 60ef79386d..45d69ab1bd 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__chat_small_running_h3.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__chat_small_running_h3.snap @@ -1,7 +1,8 @@ --- source: tui/src/chatwidget/tests.rs +assertion_line: 1500 expression: terminal.backend() --- " " +"• Thinking (0s • esc to interrupt) " "› Ask Codex to do anything " -" " diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__exec_approval_modal_exec.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__exec_approval_modal_exec.snap index f52a0f38a2..f0cf3a2d33 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__exec_approval_modal_exec.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__exec_approval_modal_exec.snap @@ -1,9 +1,10 @@ --- source: tui/src/chatwidget/tests.rs +assertion_line: 409 expression: "format!(\"{buf:?}\")" --- Buffer { - area: Rect { x: 0, y: 0, width: 80, height: 15 }, + area: Rect { x: 0, y: 0, width: 80, height: 14 }, content: [ " ", " ", @@ -19,7 +20,6 @@ Buffer { " 3. No, and tell Codex what to do differently esc ", " ", " Press enter to confirm or esc to cancel ", - " ", ], styles: [ x: 0, y: 0, fg: Reset, bg: Reset, underline: Reset, modifier: NONE, @@ -34,6 +34,5 @@ Buffer { x: 47, y: 11, fg: Reset, bg: Reset, underline: Reset, modifier: DIM, x: 50, y: 11, fg: Reset, bg: Reset, underline: Reset, modifier: NONE, x: 2, y: 13, fg: Reset, bg: Reset, underline: Reset, modifier: DIM, - x: 0, y: 14, fg: Reset, bg: Reset, underline: Reset, modifier: NONE, ] } diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__status_widget_active.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__status_widget_active.snap index 7c3253a2ea..e288807b00 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__status_widget_active.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__status_widget_active.snap @@ -1,5 +1,6 @@ --- source: tui/src/chatwidget/tests.rs +assertion_line: 1577 expression: terminal.backend() --- " " @@ -9,4 +10,3 @@ expression: terminal.backend() "› Ask Codex to do anything " " " " ? for shortcuts " -" " diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__status_widget_and_approval_modal.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__status_widget_and_approval_modal.snap index d1951cd082..086a3a8c7e 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__status_widget_and_approval_modal.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__status_widget_and_approval_modal.snap @@ -1,5 +1,6 @@ --- source: tui/src/chatwidget/tests.rs +assertion_line: 1548 expression: terminal.backend() --- " " @@ -16,4 +17,3 @@ expression: terminal.backend() " 3. No, and tell Codex what to do differently esc " " " " Press enter to confirm or esc to cancel " -" " diff --git a/codex-rs/tui/src/cli.rs b/codex-rs/tui/src/cli.rs index f0630a34c5..ee2a0355da 100644 --- a/codex-rs/tui/src/cli.rs +++ b/codex-rs/tui/src/cli.rs @@ -7,7 +7,7 @@ use std::path::PathBuf; #[command(version)] pub struct Cli { /// Optional user prompt to start the session. - #[arg(value_name = "PROMPT")] + #[arg(value_name = "PROMPT", value_hint = clap::ValueHint::Other)] pub prompt: Option, /// Optional image(s) to attach to the initial prompt. diff --git a/codex-rs/tui/src/custom_terminal.rs b/codex-rs/tui/src/custom_terminal.rs index aa79e980c1..bbd8900632 100644 --- a/codex-rs/tui/src/custom_terminal.rs +++ b/codex-rs/tui/src/custom_terminal.rs @@ -120,6 +120,8 @@ where /// Last known position of the cursor. Used to find the new area when the viewport is inlined /// and the terminal resized. pub last_known_cursor_pos: Position, + + use_custom_flush: bool, } impl Drop for Terminal @@ -158,6 +160,7 @@ where viewport_area: Rect::new(0, cursor_pos.y, 0, 0), last_known_screen_size: screen_size, last_known_cursor_pos: cursor_pos, + use_custom_flush: true, }) } @@ -190,15 +193,24 @@ where pub fn flush(&mut self) -> io::Result<()> { let previous_buffer = &self.buffers[1 - self.current]; let current_buffer = &self.buffers[self.current]; - let updates = diff_buffers(previous_buffer, current_buffer); - if let Some(DrawCommand::Put { x, y, .. }) = updates - .iter() - .rev() - .find(|cmd| matches!(cmd, DrawCommand::Put { .. })) - { - self.last_known_cursor_pos = Position { x: *x, y: *y }; + + if self.use_custom_flush { + let updates = diff_buffers(previous_buffer, current_buffer); + if let Some(DrawCommand::Put { x, y, .. }) = updates + .iter() + .rev() + .find(|cmd| matches!(cmd, DrawCommand::Put { .. })) + { + self.last_known_cursor_pos = Position { x: *x, y: *y }; + } + draw(&mut self.backend, updates.into_iter()) + } else { + let updates = previous_buffer.diff(current_buffer); + if let Some((x, y, _)) = updates.last() { + self.last_known_cursor_pos = Position { x: *x, y: *y }; + } + self.backend.draw(updates.into_iter()) } - draw(&mut self.backend, updates.into_iter()) } /// Updates the Terminal so that internal buffers match the requested area. @@ -408,11 +420,13 @@ fn diff_buffers<'a>(a: &'a Buffer, b: &'a Buffer) -> Vec> { let x = row .iter() - .rposition(|cell| cell.symbol() != " " || cell.bg != bg) + .rposition(|cell| { + cell.symbol() != " " || cell.bg != bg || cell.modifier != Modifier::empty() + }) .unwrap_or(0); last_nonblank_column[y as usize] = x as u16; - let (x_abs, y_abs) = a.pos_of(row_start + x + 1); if x < (a.area.width as usize).saturating_sub(1) { + let (x_abs, y_abs) = a.pos_of(row_start + x + 1); updates.push(DrawCommand::ClearToEnd { x: x_abs, y: y_abs, diff --git a/codex-rs/tui/src/diff_render.rs b/codex-rs/tui/src/diff_render.rs index 6f5f7e7d8f..a4b7c87f73 100644 --- a/codex-rs/tui/src/diff_render.rs +++ b/codex-rs/tui/src/diff_render.rs @@ -13,13 +13,14 @@ use std::path::Path; use std::path::PathBuf; use crate::exec_command::relativize_to_home; +use crate::render::Insets; +use crate::render::line_utils::prefix_lines; use crate::render::renderable::ColumnRenderable; +use crate::render::renderable::InsetRenderable; use crate::render::renderable::Renderable; use codex_core::git_info::get_git_repo_root; use codex_core::protocol::FileChange; -const SPACES_AFTER_LINE_NUMBER: usize = 6; - // Internal representation for diff line rendering enum DiffLineType { Insert, @@ -65,7 +66,10 @@ impl From for Box { path.extend(render_line_count_summary(row.added, row.removed)); rows.push(Box::new(path)); rows.push(Box::new(RtLine::from(""))); - rows.push(Box::new(row.change)); + rows.push(Box::new(InsetRenderable::new( + row.change, + Insets::tlbr(0, 2, 0, 0), + ))); } Box::new(ColumnRenderable::new(rows)) @@ -181,7 +185,9 @@ fn render_changes_block(rows: Vec, wrap_cols: usize, cwd: &Path) -> Vec, wrap_cols: usize, cwd: &Path) -> Vec>, width: usize) { match change { FileChange::Add { content } => { + let line_number_width = line_number_width(content.lines().count()); for (i, raw) in content.lines().enumerate() { out.extend(push_wrapped_diff_line( i + 1, DiffLineType::Insert, raw, width, + line_number_width, )); } } FileChange::Delete { content } => { + let line_number_width = line_number_width(content.lines().count()); for (i, raw) in content.lines().enumerate() { out.extend(push_wrapped_diff_line( i + 1, DiffLineType::Delete, raw, width, + line_number_width, )); } } FileChange::Update { unified_diff, .. } => { if let Ok(patch) = diffy::Patch::from_str(unified_diff) { + let mut max_line_number = 0; + for h in patch.hunks() { + let mut old_ln = h.old_range().start(); + let mut new_ln = h.new_range().start(); + for l in h.lines() { + match l { + diffy::Line::Insert(_) => { + max_line_number = max_line_number.max(new_ln); + new_ln += 1; + } + diffy::Line::Delete(_) => { + max_line_number = max_line_number.max(old_ln); + old_ln += 1; + } + diffy::Line::Context(_) => { + max_line_number = max_line_number.max(new_ln); + old_ln += 1; + new_ln += 1; + } + } + } + } + let line_number_width = line_number_width(max_line_number); let mut is_first_hunk = true; for h in patch.hunks() { if !is_first_hunk { - out.push(RtLine::from(vec![" ".into(), "⋮".dim()])); + let spacer = format!("{:width$} ", "", width = line_number_width.max(1)); + let spacer_span = RtSpan::styled(spacer, style_gutter()); + out.push(RtLine::from(vec![spacer_span, "⋮".dim()])); } is_first_hunk = false; @@ -229,6 +264,7 @@ fn render_change(change: &FileChange, out: &mut Vec>, width: usi DiffLineType::Insert, s, width, + line_number_width, )); new_ln += 1; } @@ -239,6 +275,7 @@ fn render_change(change: &FileChange, out: &mut Vec>, width: usi DiffLineType::Delete, s, width, + line_number_width, )); old_ln += 1; } @@ -249,6 +286,7 @@ fn render_change(change: &FileChange, out: &mut Vec>, width: usi DiffLineType::Context, s, width, + line_number_width, )); old_ln += 1; new_ln += 1; @@ -298,17 +336,15 @@ fn push_wrapped_diff_line( kind: DiffLineType, text: &str, width: usize, + line_number_width: usize, ) -> Vec> { - let indent = " "; let ln_str = line_number.to_string(); let mut remaining_text: &str = text; - // Reserve a fixed number of spaces after the line number so that content starts - // at a consistent column. Content includes a 1-character diff sign prefix - // ("+"/"-" for inserts/deletes, or a space for context lines) so alignment - // stays consistent across all diff lines. - let gap_after_ln = SPACES_AFTER_LINE_NUMBER.saturating_sub(ln_str.len()); - let prefix_cols = indent.len() + ln_str.len() + gap_after_ln; + // Reserve a fixed number of spaces (equal to the widest line number plus a + // trailing spacer) so the sign column stays aligned across the diff block. + let gutter_width = line_number_width.max(1); + let prefix_cols = gutter_width + 1; let mut first = true; let (sign_char, line_style) = match kind { @@ -332,8 +368,8 @@ fn push_wrapped_diff_line( remaining_text = rest; if first { - // Build gutter (indent + line number + spacing) as a dimmed span - let gutter = format!("{indent}{ln_str}{}", " ".repeat(gap_after_ln)); + // Build gutter (right-aligned line number plus spacer) as a dimmed span + let gutter = format!("{ln_str:>gutter_width$} "); // Content with a sign ('+'/'-'/' ') styled per diff kind let content = format!("{sign_char}{chunk}"); lines.push(RtLine::from(vec![ @@ -343,7 +379,7 @@ fn push_wrapped_diff_line( first = false; } else { // Continuation lines keep a space for the sign column so content aligns - let gutter = format!("{indent}{} ", " ".repeat(ln_str.len() + gap_after_ln)); + let gutter = format!("{:gutter_width$} ", ""); lines.push(RtLine::from(vec![ RtSpan::styled(gutter, style_gutter()), RtSpan::styled(chunk.to_string(), line_style), @@ -356,6 +392,14 @@ fn push_wrapped_diff_line( lines } +fn line_number_width(max_line_number: usize) -> usize { + if max_line_number == 0 { + 1 + } else { + max_line_number.to_string().len() + } +} + fn style_gutter() -> Style { Style::default().add_modifier(Modifier::DIM) } @@ -421,7 +465,8 @@ mod tests { let long_line = "this is a very long line that should wrap across multiple terminal columns and continue"; // Call the wrapping function directly so we can precisely control the width - let lines = push_wrapped_diff_line(1, DiffLineType::Insert, long_line, 80); + let lines = + push_wrapped_diff_line(1, DiffLineType::Insert, long_line, 80, line_number_width(1)); // Render into a small terminal to capture the visual layout snapshot_lines("wrap_behavior_insert", lines, 90, 8); @@ -442,11 +487,9 @@ mod tests { }, ); - for name in ["apply_update_block", "apply_update_block_manual"] { - let lines = diff_summary_for_tests(&changes); + let lines = diff_summary_for_tests(&changes); - snapshot_lines(name, lines, 80, 12); - } + snapshot_lines("apply_update_block", lines, 80, 12); } #[test] @@ -573,14 +616,37 @@ mod tests { }, ); - let mut lines = create_diff_summary(&changes, &PathBuf::from("/"), 28); - // Drop the combined header for this text-only snapshot - if !lines.is_empty() { - lines.remove(0); - } + let lines = create_diff_summary(&changes, &PathBuf::from("/"), 28); snapshot_lines_text("apply_update_block_wraps_long_lines_text", &lines); } + #[test] + fn ui_snapshot_apply_update_block_line_numbers_three_digits_text() { + let original = (1..=110).map(|i| format!("line {i}\n")).collect::(); + let modified = (1..=110) + .map(|i| { + if i == 100 { + format!("line {i} changed\n") + } else { + format!("line {i}\n") + } + }) + .collect::(); + let patch = diffy::create_patch(&original, &modified).to_string(); + + let mut changes: HashMap = HashMap::new(); + changes.insert( + PathBuf::from("hundreds.txt"), + FileChange::Update { + unified_diff: patch, + move_path: None, + }, + ); + + let lines = create_diff_summary(&changes, &PathBuf::from("/"), 80); + snapshot_lines_text("apply_update_block_line_numbers_three_digits_text", &lines); + } + #[test] fn ui_snapshot_apply_update_block_relativizes_path() { let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("/")); diff --git a/codex-rs/tui/src/exec_cell/render.rs b/codex-rs/tui/src/exec_cell/render.rs index 9c1231f0cf..a3cc8cacae 100644 --- a/codex-rs/tui/src/exec_cell/render.rs +++ b/codex-rs/tui/src/exec_cell/render.rs @@ -8,8 +8,10 @@ use crate::history_cell::HistoryCell; use crate::render::highlight::highlight_bash_to_lines; use crate::render::line_utils::prefix_lines; use crate::render::line_utils::push_owned_lines; +use crate::shimmer::shimmer_spans; use crate::wrapping::RtOptions; use crate::wrapping::word_wrap_line; +use crate::wrapping::word_wrap_lines; use codex_ansi_escape::ansi_escape_line; use codex_common::elapsed::format_duration; use codex_protocol::parse_command::ParsedCommand; @@ -116,10 +118,16 @@ pub(crate) fn output_lines( } pub(crate) fn spinner(start_time: Option) -> Span<'static> { - let blink_on = start_time - .map(|st| ((st.elapsed().as_millis() / 600) % 2) == 0) - .unwrap_or(false); - if blink_on { "•".into() } else { "◦".dim() } + let elapsed = start_time.map(|st| st.elapsed()).unwrap_or_default(); + if supports_color::on_cached(supports_color::Stream::Stdout) + .map(|level| level.has_16m) + .unwrap_or(false) + { + shimmer_spans("•")[0].clone() + } else { + let blink_on = (elapsed.as_millis() / 600).is_multiple_of(2); + if blink_on { "•".into() } else { "◦".dim() } + } } impl HistoryCell for ExecCell { @@ -131,17 +139,25 @@ impl HistoryCell for ExecCell { } } - fn transcript_lines(&self) -> Vec> { + fn desired_transcript_height(&self, width: u16) -> u16 { + self.transcript_lines(width).len() as u16 + } + + fn transcript_lines(&self, width: u16) -> Vec> { let mut lines: Vec> = vec![]; - for call in self.iter_calls() { - let cmd_display = strip_bash_lc_and_escape(&call.command); - for (i, part) in cmd_display.lines().enumerate() { - if i == 0 { - lines.push(vec!["$ ".magenta(), part.to_string().into()].into()); - } else { - lines.push(vec![" ".into(), part.to_string().into()].into()); - } + for (i, call) in self.iter_calls().enumerate() { + if i > 0 { + lines.push("".into()); } + let script = strip_bash_lc_and_escape(&call.command); + let highlighted_script = highlight_bash_to_lines(&script); + let cmd_display = word_wrap_lines( + &highlighted_script, + RtOptions::new(width as usize) + .initial_indent("$ ".magenta().into()) + .subsequent_indent(" ".into()), + ); + lines.extend(cmd_display); if let Some(output) = call.output.as_ref() { lines.extend(output.formatted_output.lines().map(ansi_escape_line)); @@ -160,7 +176,6 @@ impl HistoryCell for ExecCell { result.push_span(format!(" • {duration}").dim()); lines.push(result); } - lines.push("".into()); } lines } diff --git a/codex-rs/tui/src/history_cell.rs b/codex-rs/tui/src/history_cell.rs index e3a89da340..7c4d18e1a3 100644 --- a/codex-rs/tui/src/history_cell.rs +++ b/codex-rs/tui/src/history_cell.rs @@ -11,7 +11,6 @@ use crate::markdown::append_markdown; use crate::render::line_utils::line_to_static; use crate::render::line_utils::prefix_lines; use crate::style::user_message_style; -use crate::terminal_palette::default_bg; use crate::text_formatting::format_and_truncate_tool_result; use crate::ui_consts::LIVE_PREFIX_COLS; use crate::wrapping::RtOptions; @@ -56,10 +55,6 @@ use unicode_width::UnicodeWidthStr; pub(crate) trait HistoryCell: std::fmt::Debug + Send + Sync + Any { fn display_lines(&self, width: u16) -> Vec>; - fn transcript_lines(&self) -> Vec> { - self.display_lines(u16::MAX) - } - fn desired_height(&self, width: u16) -> u16 { Paragraph::new(Text::from(self.display_lines(width))) .wrap(Wrap { trim: false }) @@ -68,6 +63,29 @@ pub(crate) trait HistoryCell: std::fmt::Debug + Send + Sync + Any { .unwrap_or(0) } + fn transcript_lines(&self, width: u16) -> Vec> { + self.display_lines(width) + } + + fn desired_transcript_height(&self, width: u16) -> u16 { + let lines = self.transcript_lines(width); + // Workaround for ratatui bug: if there's only one line and it's whitespace-only, ratatui gives 2 lines. + if let [line] = &lines[..] + && line + .spans + .iter() + .all(|s| s.content.chars().all(char::is_whitespace)) + { + return 1; + } + + Paragraph::new(Text::from(lines)) + .wrap(Wrap { trim: false }) + .line_count(width) + .try_into() + .unwrap_or(0) + } + fn is_stream_continuation(&self) -> bool { false } @@ -92,19 +110,18 @@ impl HistoryCell for UserHistoryCell { fn display_lines(&self, width: u16) -> Vec> { let mut lines: Vec> = Vec::new(); - // Use ratatui-aware word wrapping and prefixing to avoid lifetime issues. - let wrap_width = width.saturating_sub(LIVE_PREFIX_COLS); // account for the ▌ prefix and trailing space + let wrap_width = width.saturating_sub(LIVE_PREFIX_COLS); - let style = user_message_style(default_bg()); + let style = user_message_style(); - // Use our ratatui wrapping helpers for correct styling and lifetimes. let wrapped = word_wrap_lines( &self .message .lines() .map(|l| Line::from(l).style(style)) .collect::>(), - RtOptions::new(wrap_width as usize), + // Wrap algorithm matches textarea.rs. + RtOptions::new(wrap_width as usize).wrap_algorithm(textwrap::WrapAlgorithm::FirstFit), ); lines.push(Line::from("").style(style)); @@ -112,13 +129,6 @@ impl HistoryCell for UserHistoryCell { lines.push(Line::from("").style(style)); lines } - - fn transcript_lines(&self) -> Vec> { - let mut lines: Vec> = Vec::new(); - lines.push("user".cyan().bold().into()); - lines.extend(self.message.lines().map(|l| l.to_string().into())); - lines - } } #[derive(Debug)] @@ -126,6 +136,7 @@ pub(crate) struct ReasoningSummaryCell { _header: String, content: String, citation_context: MarkdownCitationContext, + transcript_only: bool, } impl ReasoningSummaryCell { @@ -133,17 +144,17 @@ impl ReasoningSummaryCell { header: String, content: String, citation_context: MarkdownCitationContext, + transcript_only: bool, ) -> Self { Self { _header: header, content, citation_context, + transcript_only, } } -} -impl HistoryCell for ReasoningSummaryCell { - fn display_lines(&self, width: u16) -> Vec> { + fn lines(&self, width: u16) -> Vec> { let mut lines: Vec> = Vec::new(); append_markdown( &self.content, @@ -151,7 +162,7 @@ impl HistoryCell for ReasoningSummaryCell { &mut lines, self.citation_context.clone(), ); - let summary_style = Style::default().add_modifier(Modifier::DIM | Modifier::ITALIC); + let summary_style = Style::default().dim().italic(); let summary_lines = lines .into_iter() .map(|mut line| { @@ -171,19 +182,31 @@ impl HistoryCell for ReasoningSummaryCell { .subsequent_indent(" ".into()), ) } +} - fn transcript_lines(&self) -> Vec> { - let mut out: Vec> = Vec::new(); - out.push("thinking".magenta().bold().into()); - let mut lines = Vec::new(); - append_markdown( - &self.content, - None, - &mut lines, - self.citation_context.clone(), - ); - out.extend(lines); - out +impl HistoryCell for ReasoningSummaryCell { + fn display_lines(&self, width: u16) -> Vec> { + if self.transcript_only { + Vec::new() + } else { + self.lines(width) + } + } + + fn desired_height(&self, width: u16) -> u16 { + if self.transcript_only { + 0 + } else { + self.lines(width).len() as u16 + } + } + + fn transcript_lines(&self, width: u16) -> Vec> { + self.lines(width) + } + + fn desired_transcript_height(&self, width: u16) -> u16 { + self.lines(width).len() as u16 } } @@ -216,15 +239,6 @@ impl HistoryCell for AgentMessageCell { ) } - fn transcript_lines(&self) -> Vec> { - let mut out: Vec> = Vec::new(); - if self.is_first_line { - out.push("codex".magenta().bold().into()); - } - out.extend(self.lines.clone()); - out - } - fn is_stream_continuation(&self) -> bool { !self.is_first_line } @@ -247,21 +261,6 @@ impl HistoryCell for PlainHistoryCell { } } -#[derive(Debug)] -pub(crate) struct TranscriptOnlyHistoryCell { - lines: Vec>, -} - -impl HistoryCell for TranscriptOnlyHistoryCell { - fn display_lines(&self, _width: u16) -> Vec> { - Vec::new() - } - - fn transcript_lines(&self) -> Vec> { - self.lines.clone() - } -} - /// Cyan history cell line showing the current review status. pub(crate) fn new_review_status_line(message: String) -> PlainHistoryCell { PlainHistoryCell { @@ -1049,16 +1048,6 @@ pub(crate) fn new_view_image_tool_call(path: PathBuf, cwd: &Path) -> PlainHistor PlainHistoryCell { lines } } -pub(crate) fn new_reasoning_block( - full_reasoning_buffer: String, - config: &Config, -) -> TranscriptOnlyHistoryCell { - let mut lines: Vec> = Vec::new(); - lines.push(Line::from("thinking".magenta().italic())); - append_markdown(&full_reasoning_buffer, None, &mut lines, config); - TranscriptOnlyHistoryCell { lines } -} - pub(crate) fn new_reasoning_summary_block( full_reasoning_buffer: String, config: &Config, @@ -1084,12 +1073,18 @@ pub(crate) fn new_reasoning_summary_block( header_buffer, summary_buffer, config.into(), + false, )); } } } } - Box::new(new_reasoning_block(full_reasoning_buffer, config)) + Box::new(ReasoningSummaryCell::new( + "".to_string(), + full_reasoning_buffer, + config.into(), + true, + )) } #[derive(Debug)] @@ -1120,10 +1115,6 @@ impl HistoryCell for FinalMessageSeparator { vec![Line::from_iter(["─".repeat(width as usize).dim()])] } } - - fn transcript_lines(&self) -> Vec> { - vec![] - } } fn format_mcp_invocation<'a>(invocation: McpInvocation) -> Line<'a> { @@ -1187,7 +1178,14 @@ mod tests { } fn render_transcript(cell: &dyn HistoryCell) -> Vec { - render_lines(&cell.transcript_lines()) + render_lines(&cell.transcript_lines(u16::MAX)) + } + + #[test] + fn empty_agent_message_cell_transcript() { + let cell = AgentMessageCell::new(vec![Line::default()], false); + assert_eq!(cell.transcript_lines(80), vec![Line::from(" ")]); + assert_eq!(cell.desired_transcript_height(80), 1); } #[test] @@ -1882,10 +1880,7 @@ mod tests { assert_eq!(rendered_display, vec!["• Detailed reasoning goes here."]); let rendered_transcript = render_transcript(cell.as_ref()); - assert_eq!( - rendered_transcript, - vec!["thinking", "Detailed reasoning goes here."] - ); + assert_eq!(rendered_transcript, vec!["• Detailed reasoning goes here."]); } #[test] @@ -1897,7 +1892,7 @@ mod tests { new_reasoning_summary_block("Detailed reasoning goes here.".to_string(), &config); let rendered = render_transcript(cell.as_ref()); - assert_eq!(rendered, vec!["thinking", "Detailed reasoning goes here."]); + assert_eq!(rendered, vec!["• Detailed reasoning goes here."]); } #[test] @@ -1911,10 +1906,7 @@ mod tests { ); let rendered = render_transcript(cell.as_ref()); - assert_eq!( - rendered, - vec!["thinking", "**High level reasoning without closing"] - ); + assert_eq!(rendered, vec!["• **High level reasoning without closing"]); } #[test] @@ -1928,10 +1920,7 @@ mod tests { ); let rendered = render_transcript(cell.as_ref()); - assert_eq!( - rendered, - vec!["thinking", "High level reasoning without closing"] - ); + assert_eq!(rendered, vec!["• High level reasoning without closing"]); let cell = new_reasoning_summary_block( "**High level reasoning without closing**\n\n ".to_string(), @@ -1939,10 +1928,7 @@ mod tests { ); let rendered = render_transcript(cell.as_ref()); - assert_eq!( - rendered, - vec!["thinking", "High level reasoning without closing"] - ); + assert_eq!(rendered, vec!["• High level reasoning without closing"]); } #[test] @@ -1959,9 +1945,6 @@ mod tests { assert_eq!(rendered_display, vec!["• We should fix the bug next."]); let rendered_transcript = render_transcript(cell.as_ref()); - assert_eq!( - rendered_transcript, - vec!["thinking", "We should fix the bug next."] - ); + assert_eq!(rendered_transcript, vec!["• We should fix the bug next."]); } } diff --git a/codex-rs/tui/src/pager_overlay.rs b/codex-rs/tui/src/pager_overlay.rs index 7997625aae..46c6c19a18 100644 --- a/codex-rs/tui/src/pager_overlay.rs +++ b/codex-rs/tui/src/pager_overlay.rs @@ -3,9 +3,13 @@ use std::sync::Arc; use std::time::Duration; use crate::history_cell::HistoryCell; +use crate::history_cell::UserHistoryCell; use crate::key_hint; use crate::key_hint::KeyBinding; +use crate::render::Insets; +use crate::render::renderable::InsetRenderable; use crate::render::renderable::Renderable; +use crate::style::user_message_style; use crate::tui; use crate::tui::TuiEvent; use crossterm::event::KeyCode; @@ -13,6 +17,7 @@ use crossterm::event::KeyEvent; use ratatui::buffer::Buffer; use ratatui::buffer::Cell; use ratatui::layout::Rect; +use ratatui::style::Style; use ratatui::style::Stylize; use ratatui::text::Line; use ratatui::text::Span; @@ -21,7 +26,6 @@ use ratatui::widgets::Clear; use ratatui::widgets::Paragraph; use ratatui::widgets::Widget; use ratatui::widgets::WidgetRef; -use ratatui::widgets::Wrap; pub(crate) enum Overlay { Transcript(TranscriptOverlay), @@ -317,29 +321,30 @@ impl PagerView { } } -struct CachedParagraph { - paragraph: Paragraph<'static>, +/// A renderable that caches its desired height. +struct CachedRenderable { + renderable: Box, height: std::cell::Cell>, last_width: std::cell::Cell>, } -impl CachedParagraph { - fn new(paragraph: Paragraph<'static>) -> Self { +impl CachedRenderable { + fn new(renderable: Box) -> Self { Self { - paragraph, + renderable, height: std::cell::Cell::new(None), last_width: std::cell::Cell::new(None), } } } -impl Renderable for CachedParagraph { +impl Renderable for CachedRenderable { fn render(&self, area: Rect, buf: &mut Buffer) { - self.paragraph.render_ref(area, buf); + self.renderable.render(area, buf); } fn desired_height(&self, width: u16) -> u16 { if self.last_width.get() != Some(width) { - let height = self.paragraph.line_count(width) as u16; + let height = self.renderable.desired_height(width); self.height.set(Some(height)); self.last_width.set(Some(width)); } @@ -347,6 +352,23 @@ impl Renderable for CachedParagraph { } } +struct CellRenderable { + cell: Arc, + style: Style, +} + +impl Renderable for CellRenderable { + fn render(&self, area: Rect, buf: &mut Buffer) { + let p = + Paragraph::new(Text::from(self.cell.transcript_lines(area.width))).style(self.style); + p.render(area, buf); + } + + fn desired_height(&self, width: u16) -> u16 { + self.cell.desired_transcript_height(width) + } +} + pub(crate) struct TranscriptOverlay { view: PagerView, cells: Vec>, @@ -358,7 +380,7 @@ impl TranscriptOverlay { pub(crate) fn new(transcript_cells: Vec>) -> Self { Self { view: PagerView::new( - Self::render_cells_to_texts(&transcript_cells, None), + Self::render_cells(&transcript_cells, None), "T R A N S C R I P T".to_string(), usize::MAX, ), @@ -368,46 +390,46 @@ impl TranscriptOverlay { } } - fn render_cells_to_texts( + fn render_cells( cells: &[Arc], highlight_cell: Option, ) -> Vec> { - let mut texts: Vec> = Vec::new(); - let mut first = true; - for (idx, cell) in cells.iter().enumerate() { - let mut lines: Vec> = Vec::new(); - if !cell.is_stream_continuation() && !first { - lines.push(Line::from("")); - } - let cell_lines = if Some(idx) == highlight_cell { - cell.transcript_lines() - .into_iter() - .map(Stylize::reversed) - .collect() - } else { - cell.transcript_lines() - }; - lines.extend(cell_lines); - texts.push(Box::new(CachedParagraph::new( - Paragraph::new(Text::from(lines)).wrap(Wrap { trim: false }), - ))); - first = false; - } - texts + cells + .iter() + .enumerate() + .flat_map(|(i, c)| { + let mut v: Vec> = Vec::new(); + let mut cell_renderable = if c.as_any().is::() { + Box::new(CachedRenderable::new(Box::new(CellRenderable { + cell: c.clone(), + style: if highlight_cell == Some(i) { + user_message_style().reversed() + } else { + user_message_style() + }, + }))) as Box + } else { + Box::new(CachedRenderable::new(Box::new(CellRenderable { + cell: c.clone(), + style: Style::default(), + }))) as Box + }; + if !c.is_stream_continuation() && i > 0 { + cell_renderable = Box::new(InsetRenderable::new( + cell_renderable, + Insets::tlbr(1, 0, 0, 0), + )); + } + v.push(cell_renderable); + v + }) + .collect() } pub(crate) fn insert_cell(&mut self, cell: Arc) { let follow_bottom = self.view.is_scrolled_to_bottom(); - // Append as a new Text chunk (with a separating blank if needed) - let mut lines: Vec> = Vec::new(); - if !cell.is_stream_continuation() && !self.cells.is_empty() { - lines.push(Line::from("")); - } - lines.extend(cell.transcript_lines()); - self.view.renderables.push(Box::new(CachedParagraph::new( - Paragraph::new(Text::from(lines)).wrap(Wrap { trim: false }), - ))); self.cells.push(cell); + self.view.renderables = Self::render_cells(&self.cells, self.highlight_cell); if follow_bottom { self.view.scroll_offset = usize::MAX; } @@ -415,7 +437,7 @@ impl TranscriptOverlay { pub(crate) fn set_highlight_cell(&mut self, cell: Option) { self.highlight_cell = cell; - self.view.renderables = Self::render_cells_to_texts(&self.cells, self.highlight_cell); + self.view.renderables = Self::render_cells(&self.cells, self.highlight_cell); if let Some(idx) = self.highlight_cell { self.view.scroll_chunk_into_view(idx); } @@ -475,8 +497,8 @@ pub(crate) struct StaticOverlay { impl StaticOverlay { pub(crate) fn with_title(lines: Vec>, title: String) -> Self { Self::with_renderables( - vec![Box::new(CachedParagraph::new(Paragraph::new(Text::from( - lines, + vec![Box::new(CachedRenderable::new(Box::new(Paragraph::new( + Text::from(lines), ))))], title, ) @@ -585,7 +607,7 @@ mod tests { self.lines.clone() } - fn transcript_lines(&self) -> Vec> { + fn transcript_lines(&self, _width: u16) -> Vec> { self.lines.clone() } } diff --git a/codex-rs/tui/src/render/highlight.rs b/codex-rs/tui/src/render/highlight.rs index 393aa33705..e6d200cc3d 100644 --- a/codex-rs/tui/src/render/highlight.rs +++ b/codex-rs/tui/src/render/highlight.rs @@ -1,81 +1,146 @@ -use codex_core::bash::try_parse_bash; +use ratatui::style::Style; use ratatui::style::Stylize; use ratatui::text::Line; use ratatui::text::Span; +use std::sync::OnceLock; +use tree_sitter_highlight::Highlight; +use tree_sitter_highlight::HighlightConfiguration; +use tree_sitter_highlight::HighlightEvent; +use tree_sitter_highlight::Highlighter; -/// Convert the full bash script into per-line styled content by first -/// computing operator-dimmed spans across the entire script, then splitting -/// by newlines and dimming heredoc body lines. Performs a single parse and -/// reuses it for both highlighting and heredoc detection. -pub(crate) fn highlight_bash_to_lines(script: &str) -> Vec> { - // Parse once; use the tree for both highlighting and heredoc body detection. - let spans: Vec> = if let Some(tree) = try_parse_bash(script) { - // Single walk: collect operator ranges and heredoc rows. - let root = tree.root_node(); - let mut cursor = root.walk(); - let mut stack = vec![root]; - let mut ranges: Vec<(usize, usize)> = Vec::new(); - while let Some(node) = stack.pop() { - if !node.is_named() && !node.is_extra() { - let kind = node.kind(); - let is_quote = matches!(kind, "\"" | "'" | "`"); - let is_whitespace = kind.trim().is_empty(); - if !is_quote && !is_whitespace { - ranges.push((node.start_byte(), node.end_byte())); - } - } else if node.kind() == "heredoc_body" { - ranges.push((node.start_byte(), node.end_byte())); - } - for child in node.children(&mut cursor) { - stack.push(child); - } - } - if ranges.is_empty() { - ranges.push((script.len(), script.len())); - } - ranges.sort_by_key(|(st, _)| *st); - let mut spans: Vec> = Vec::new(); - let mut i = 0usize; - for (start, end) in ranges.into_iter() { - let dim_start = start.max(i); - let dim_end = end; - if dim_start < dim_end { - if dim_start > i { - spans.push(script[i..dim_start].to_string().into()); - } - spans.push(script[dim_start..dim_end].to_string().dim()); - i = dim_end; - } - } - if i < script.len() { - spans.push(script[i..].to_string().into()); - } - spans - } else { - vec![script.to_string().into()] - }; - // Split spans into lines preserving style boundaries and highlights across newlines. - let mut lines: Vec> = vec![Line::from("")]; - for sp in spans { - let style = sp.style; - let text = sp.content.into_owned(); - for (i, part) in text.split('\n').enumerate() { - if i > 0 { - lines.push(Line::from("")); - } - if part.is_empty() { - continue; - } - let span = Span { - style, - content: std::borrow::Cow::Owned(part.to_string()), - }; - if let Some(last) = lines.last_mut() { - last.spans.push(span); - } +// Ref: https://github.com/tree-sitter/tree-sitter-bash/blob/master/queries/highlights.scm +#[derive(Copy, Clone)] +enum BashHighlight { + Comment, + Constant, + Embedded, + Function, + Keyword, + Number, + Operator, + Property, + String, +} + +impl BashHighlight { + const ALL: [Self; 9] = [ + Self::Comment, + Self::Constant, + Self::Embedded, + Self::Function, + Self::Keyword, + Self::Number, + Self::Operator, + Self::Property, + Self::String, + ]; + + const fn as_str(self) -> &'static str { + match self { + Self::Comment => "comment", + Self::Constant => "constant", + Self::Embedded => "embedded", + Self::Function => "function", + Self::Keyword => "keyword", + Self::Number => "number", + Self::Operator => "operator", + Self::Property => "property", + Self::String => "string", } } - lines + + fn style(self) -> Style { + match self { + Self::Comment | Self::Operator | Self::String => Style::default().dim(), + _ => Style::default(), + } + } +} + +static HIGHLIGHT_CONFIG: OnceLock = OnceLock::new(); + +fn highlight_names() -> &'static [&'static str] { + static NAMES: OnceLock<[&'static str; BashHighlight::ALL.len()]> = OnceLock::new(); + NAMES + .get_or_init(|| BashHighlight::ALL.map(BashHighlight::as_str)) + .as_slice() +} + +fn highlight_config() -> &'static HighlightConfiguration { + HIGHLIGHT_CONFIG.get_or_init(|| { + let language = tree_sitter_bash::LANGUAGE.into(); + #[expect(clippy::expect_used)] + let mut config = HighlightConfiguration::new( + language, + "bash", + tree_sitter_bash::HIGHLIGHT_QUERY, + "", + "", + ) + .expect("load bash highlight query"); + config.configure(highlight_names()); + config + }) +} + +fn highlight_for(highlight: Highlight) -> BashHighlight { + BashHighlight::ALL[highlight.0] +} + +fn push_segment(lines: &mut Vec>, segment: &str, style: Option