From 296996d74e345b1b05d8c3451a06ace21c5ada96 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Wed, 25 Jun 2025 13:29:03 -0700 Subject: [PATCH 1/4] feat: standalone file search CLI (#1386) Standalone fuzzy filename search library that should be helpful in addressing https://github.com/openai/codex/issues/1261. --- codex-rs/Cargo.lock | 52 ++++++ codex-rs/Cargo.toml | 1 + codex-rs/file-search/Cargo.toml | 20 +++ codex-rs/file-search/README.md | 5 + codex-rs/file-search/src/cli.rs | 38 +++++ codex-rs/file-search/src/lib.rs | 284 +++++++++++++++++++++++++++++++ codex-rs/file-search/src/main.rs | 50 ++++++ codex-rs/justfile | 4 + 8 files changed, 454 insertions(+) create mode 100644 codex-rs/file-search/Cargo.toml create mode 100644 codex-rs/file-search/README.md create mode 100644 codex-rs/file-search/src/cli.rs create mode 100644 codex-rs/file-search/src/lib.rs create mode 100644 codex-rs/file-search/src/main.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index bb533be143..e034a99357 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -691,6 +691,18 @@ dependencies = [ "tempfile", ] +[[package]] +name = "codex-file-search" +version = "0.0.0" +dependencies = [ + "anyhow", + "clap", + "ignore", + "nucleo-matcher", + "serde_json", + "tokio", +] + [[package]] name = "codex-linux-sandbox" version = "0.0.0" @@ -1601,6 +1613,19 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a8d1add55171497b4705a648c6b583acafb01d58050a51727785f0b2c8e0a2b2" +[[package]] +name = "globset" +version = "0.4.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "54a1028dfc5f5df5da8a56a73e6c153c9a9708ec57232470703592a3f18e49f5" +dependencies = [ + "aho-corasick", + "bstr", + "log", + "regex-automata 0.4.9", + "regex-syntax 0.8.5", +] + [[package]] name = "h2" version = "0.4.9" @@ -1985,6 +2010,22 @@ dependencies = [ "icu_properties", ] +[[package]] +name = "ignore" +version = "0.4.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d89fd380afde86567dfba715db065673989d6253f42b88179abd3eae47bda4b" +dependencies = [ + "crossbeam-deque", + "globset", + "log", + "memchr", + "regex-automata 0.4.9", + "same-file", + "walkdir", + "winapi-util", +] + [[package]] name = "image" version = "0.25.6" @@ -2577,6 +2618,16 @@ dependencies = [ "winapi", ] +[[package]] +name = "nucleo-matcher" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf33f538733d1a5a3494b836ba913207f14d9d4a1d3cd67030c5061bdd2cac85" +dependencies = [ + "memchr", + "unicode-segmentation", +] + [[package]] name = "num-bigint" version = "0.4.6" @@ -4362,6 +4413,7 @@ dependencies = [ "bytes", "libc", "mio", + "parking_lot", "pin-project-lite", "signal-hook-registry", "socket2", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 6991a6223a..f93cbbaa37 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -8,6 +8,7 @@ members = [ "core", "exec", "execpolicy", + "file-search", "linux-sandbox", "login", "mcp-client", diff --git a/codex-rs/file-search/Cargo.toml b/codex-rs/file-search/Cargo.toml new file mode 100644 index 0000000000..1850d5ac13 --- /dev/null +++ b/codex-rs/file-search/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "codex-file-search" +version = { workspace = true } +edition = "2024" + +[[bin]] +name = "codex-file-search" +path = "src/main.rs" + +[lib] +name = "codex_file_search" +path = "src/lib.rs" + +[dependencies] +anyhow = "1" +clap = { version = "4", features = ["derive"] } +ignore = "0.4.23" +nucleo-matcher = "0.3.1" +serde_json = "1.0.110" +tokio = { version = "1", features = ["full"] } diff --git a/codex-rs/file-search/README.md b/codex-rs/file-search/README.md new file mode 100644 index 0000000000..c47d494a18 --- /dev/null +++ b/codex-rs/file-search/README.md @@ -0,0 +1,5 @@ +# codex_file_search + +Fast fuzzy file search tool for Codex. + +Uses under the hood (which is what `ripgrep` uses) to traverse a directory (while honoring `.gitignore`, etc.) to produce the list of files to search and then uses to fuzzy-match the user supplied `PATTERN` against the corpus. diff --git a/codex-rs/file-search/src/cli.rs b/codex-rs/file-search/src/cli.rs new file mode 100644 index 0000000000..27afcbc140 --- /dev/null +++ b/codex-rs/file-search/src/cli.rs @@ -0,0 +1,38 @@ +use std::num::NonZero; +use std::path::PathBuf; + +use clap::ArgAction; +use clap::Parser; + +/// Fuzzy matches filenames under a directory. +#[derive(Parser)] +#[command(version)] +pub struct Cli { + /// Whether to output results in JSON format. + #[clap(long, default_value = "false")] + pub json: bool, + + /// Maximum number of results to return. + #[clap(long, short = 'l', default_value = "64")] + pub limit: NonZero, + + /// Directory to search. + #[clap(long, short = 'C')] + pub cwd: Option, + + // While it is common to default to the number of logical CPUs when creating + // a thread pool, empirically, the I/O of the filetree traversal offers + // limited parallelism and is the bottleneck, so using a smaller number of + // threads is more efficient. (Empirically, using more than 2 threads doesn't seem to provide much benefit.) + // + /// Number of worker threads to use. + #[clap(long, default_value = "2")] + pub threads: NonZero, + + /// Exclude patterns + #[arg(short, long, action = ArgAction::Append)] + pub exclude: Vec, + + /// Search pattern. + pub pattern: Option, +} diff --git a/codex-rs/file-search/src/lib.rs b/codex-rs/file-search/src/lib.rs new file mode 100644 index 0000000000..8754181670 --- /dev/null +++ b/codex-rs/file-search/src/lib.rs @@ -0,0 +1,284 @@ +use ignore::WalkBuilder; +use ignore::overrides::OverrideBuilder; +use nucleo_matcher::Matcher; +use nucleo_matcher::Utf32Str; +use nucleo_matcher::pattern::AtomKind; +use nucleo_matcher::pattern::CaseMatching; +use nucleo_matcher::pattern::Normalization; +use nucleo_matcher::pattern::Pattern; +use std::cell::UnsafeCell; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::num::NonZero; +use std::path::Path; +use std::path::PathBuf; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use tokio::process::Command; + +mod cli; + +pub use cli::Cli; + +pub struct FileSearchResults { + pub matches: Vec<(u32, String)>, + pub total_match_count: usize, +} + +pub trait Reporter { + fn report_match(&self, file: &str, score: u32); + fn warn_matches_truncated(&self, total_match_count: usize, shown_match_count: usize); + fn warn_no_search_pattern(&self, search_directory: &Path); +} + +pub async fn run_main( + Cli { + pattern, + limit, + cwd, + json: _, + exclude, + threads, + }: Cli, + reporter: T, +) -> anyhow::Result<()> { + let search_directory = match cwd { + Some(dir) => dir, + None => std::env::current_dir()?, + }; + let pattern_text = match pattern { + Some(pattern) => pattern, + None => { + reporter.warn_no_search_pattern(&search_directory); + #[cfg(unix)] + Command::new("ls") + .arg("-al") + .current_dir(search_directory) + .stdout(std::process::Stdio::inherit()) + .stderr(std::process::Stdio::inherit()) + .status() + .await?; + #[cfg(windows)] + { + Command::new("cmd") + .arg("/c") + .arg(search_directory) + .stdout(std::process::Stdio::inherit()) + .stderr(std::process::Stdio::inherit()) + .status() + .await?; + } + return Ok(()); + } + }; + + let FileSearchResults { + total_match_count, + matches, + } = run(&pattern_text, limit, search_directory, exclude, threads).await?; + let match_count = matches.len(); + let matches_truncated = total_match_count > match_count; + + for (score, file) in matches { + reporter.report_match(&file, score); + } + if matches_truncated { + reporter.warn_matches_truncated(total_match_count, match_count); + } + + Ok(()) +} + +pub async fn run( + pattern_text: &str, + limit: NonZero, + search_directory: PathBuf, + exclude: Vec, + threads: NonZero, +) -> anyhow::Result { + let pattern = create_pattern(pattern_text); + // Create one BestMatchesList per worker thread so that each worker can + // operate independently. The results across threads will be merged when + // the traversal is complete. + let WorkerCount { + num_walk_builder_threads, + num_best_matches_lists, + } = create_worker_count(threads); + let best_matchers_per_worker: Vec> = (0..num_best_matches_lists) + .map(|_| { + UnsafeCell::new(BestMatchesList::new( + limit.get(), + pattern.clone(), + Matcher::new(nucleo_matcher::Config::DEFAULT), + )) + }) + .collect(); + + // Use the same tree-walker library that ripgrep uses. We use it directly so + // that we can leverage the parallelism it provides. + let mut walk_builder = WalkBuilder::new(&search_directory); + walk_builder.threads(num_walk_builder_threads); + if !exclude.is_empty() { + let mut override_builder = OverrideBuilder::new(&search_directory); + for exclude in exclude { + // The `!` prefix is used to indicate an exclude pattern. + let exclude_pattern = format!("!{}", exclude); + override_builder.add(&exclude_pattern)?; + } + let override_matcher = override_builder.build()?; + walk_builder.overrides(override_matcher); + } + let walker = walk_builder.build_parallel(); + + // Each worker created by `WalkParallel::run()` will have its own + // `BestMatchesList` to update. + let index_counter = AtomicUsize::new(0); + walker.run(|| { + let search_directory = search_directory.clone(); + let index = index_counter.fetch_add(1, Ordering::Relaxed); + let best_list_ptr = best_matchers_per_worker[index].get(); + let best_list = unsafe { &mut *best_list_ptr }; + Box::new(move |entry| { + if let Some(path) = get_file_path(&entry, &search_directory) { + best_list.insert(path); + } + ignore::WalkState::Continue + }) + }); + + fn get_file_path<'a>( + entry_result: &'a Result, + search_directory: &std::path::Path, + ) -> Option<&'a str> { + let entry = match entry_result { + Ok(e) => e, + Err(_) => return None, + }; + if entry.file_type().is_some_and(|ft| ft.is_dir()) { + return None; + } + let path = entry.path(); + match path.strip_prefix(search_directory) { + Ok(rel_path) => rel_path.to_str(), + Err(_) => None, + } + } + + // Merge results across best_matchers_per_worker. + let mut global_heap: BinaryHeap> = BinaryHeap::new(); + let mut total_match_count = 0; + for best_list_cell in best_matchers_per_worker.iter() { + let best_list = unsafe { &*best_list_cell.get() }; + total_match_count += best_list.num_matches; + for &Reverse((score, ref line)) in best_list.binary_heap.iter() { + if global_heap.len() < limit.get() { + global_heap.push(Reverse((score, line.clone()))); + } else if let Some(min_element) = global_heap.peek() { + if score > min_element.0.0 { + global_heap.pop(); + global_heap.push(Reverse((score, line.clone()))); + } + } + } + } + + let mut matches: Vec<(u32, String)> = global_heap.into_iter().map(|r| r.0).collect(); + matches.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)); + + Ok(FileSearchResults { + matches, + total_match_count, + }) +} + +/// Maintains the `max_count` best matches for a given pattern. +struct BestMatchesList { + max_count: usize, + num_matches: usize, + pattern: Pattern, + matcher: Matcher, + binary_heap: BinaryHeap>, + + /// Internal buffer for converting strings to UTF-32. + utf32buf: Vec, +} + +impl BestMatchesList { + fn new(max_count: usize, pattern: Pattern, matcher: Matcher) -> Self { + Self { + max_count, + num_matches: 0, + pattern, + matcher, + binary_heap: BinaryHeap::new(), + utf32buf: Vec::::new(), + } + } + + fn insert(&mut self, line: &str) { + let haystack: Utf32Str<'_> = Utf32Str::new(line, &mut self.utf32buf); + if let Some(score) = self.pattern.score(haystack, &mut self.matcher) { + // In the tests below, we verify that score() returns None for a + // non-match, so we can categorically increment the count here. + self.num_matches += 1; + + if self.binary_heap.len() < self.max_count { + self.binary_heap.push(Reverse((score, line.to_string()))); + } else if let Some(min_element) = self.binary_heap.peek() { + if score > min_element.0.0 { + self.binary_heap.pop(); + self.binary_heap.push(Reverse((score, line.to_string()))); + } + } + } + } +} + +struct WorkerCount { + num_walk_builder_threads: usize, + num_best_matches_lists: usize, +} + +fn create_worker_count(num_workers: NonZero) -> WorkerCount { + // It appears that the number of times the function passed to + // `WalkParallel::run()` is called is: the number of threads specified to + // the builder PLUS ONE. + // + // In `WalkParallel::visit()`, the builder function gets called once here: + // https://github.com/BurntSushi/ripgrep/blob/79cbe89deb1151e703f4d91b19af9cdcc128b765/crates/ignore/src/walk.rs#L1233 + // + // And then once for every worker here: + // https://github.com/BurntSushi/ripgrep/blob/79cbe89deb1151e703f4d91b19af9cdcc128b765/crates/ignore/src/walk.rs#L1288 + let num_walk_builder_threads = num_workers.get(); + let num_best_matches_lists = num_walk_builder_threads + 1; + + WorkerCount { + num_walk_builder_threads, + num_best_matches_lists, + } +} + +fn create_pattern(pattern: &str) -> Pattern { + Pattern::new( + pattern, + CaseMatching::Smart, + Normalization::Smart, + AtomKind::Fuzzy, + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn verify_score_is_none_for_non_match() { + let mut utf32buf = Vec::::new(); + let line = "hello"; + let mut matcher = Matcher::new(nucleo_matcher::Config::DEFAULT); + let haystack: Utf32Str<'_> = Utf32Str::new(line, &mut utf32buf); + let pattern = create_pattern("zzz"); + let score = pattern.score(haystack, &mut matcher); + assert_eq!(score, None); + } +} diff --git a/codex-rs/file-search/src/main.rs b/codex-rs/file-search/src/main.rs new file mode 100644 index 0000000000..c25122c141 --- /dev/null +++ b/codex-rs/file-search/src/main.rs @@ -0,0 +1,50 @@ +use std::path::Path; + +use clap::Parser; +use codex_file_search::Cli; +use codex_file_search::Reporter; +use codex_file_search::run_main; +use serde_json::json; + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let cli = Cli::parse(); + let reporter = StdioReporter { + write_output_as_json: cli.json, + }; + run_main(cli, reporter).await?; + Ok(()) +} + +struct StdioReporter { + write_output_as_json: bool, +} + +impl Reporter for StdioReporter { + fn report_match(&self, file: &str, score: u32) { + if self.write_output_as_json { + let value = json!({ "file": file, "score": score }); + println!("{}", serde_json::to_string(&value).unwrap()); + } else { + println!("{file}"); + } + } + + fn warn_matches_truncated(&self, total_match_count: usize, shown_match_count: usize) { + if self.write_output_as_json { + let value = json!({"matches_truncated": true}); + println!("{}", serde_json::to_string(&value).unwrap()); + } else { + eprintln!( + "Warning: showing {shown_match_count} out of {total_match_count} results. Provide a more specific pattern or increase the --limit.", + ); + } + } + + fn warn_no_search_pattern(&self, search_directory: &Path) { + eprintln!( + "No search pattern specified. Showing the contents of the current directory ({}):", + search_directory.to_string_lossy() + ); + } +} diff --git a/codex-rs/justfile b/codex-rs/justfile index c09465a482..83a390ec56 100644 --- a/codex-rs/justfile +++ b/codex-rs/justfile @@ -16,6 +16,10 @@ exec *args: tui *args: cargo run --bin codex -- tui "$@" +# Run the CLI version of the file-search crate. +file-search *args: + cargo run --bin codex-file-search -- "$@" + # format code fmt: cargo fmt -- --config imports_granularity=Item From fcfe43c7df46836a1c60cec4dfd1591d3036a0c8 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Wed, 25 Jun 2025 23:31:11 -0700 Subject: [PATCH 2/4] feat: show number of tokens remaining in UI (#1388) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When using the OpenAI Responses API, we now record the `usage` field for a `"response.completed"` event, which includes metrics about the number of tokens consumed. We also introduce `openai_model_info.rs`, which includes current data about the most common OpenAI models available via the API (specifically `context_window` and `max_output_tokens`). If Codex does not recognize the model, you can set `model_context_window` and `model_max_output_tokens` explicitly in `config.toml`. When then introduce a new event type to `protocol.rs`, `TokenCount`, which includes the `TokenUsage` for the most recent turn. Finally, we update the TUI to record the running sum of tokens used so the percentage of available context window remaining can be reported via the placeholder text for the composer: ![Screenshot 2025-06-25 at 11 20 55 PM](https://github.com/user-attachments/assets/6fd6982f-7247-4f14-84b2-2e600cb1fd49) We could certainly get much fancier with this (such as reporting the estimated cost of the conversation), but for now, we are just trying to achieve feature parity with the TypeScript CLI. Though arguably this improves upon the TypeScript CLI, as the TypeScript CLI uses heuristics to estimate the number of tokens used rather than using the `usage` information directly: https://github.com/openai/codex/blob/296996d74e345b1b05d8c3451a06ace21c5ada96/codex-cli/src/utils/approximate-tokens-used.ts#L3-L16 Fixes https://github.com/openai/codex/issues/1242 --- codex-rs/config.md | 10 +++ codex-rs/core/src/chat_completions.rs | 18 ++++- codex-rs/core/src/client.rs | 49 +++++++++++-- codex-rs/core/src/client_common.rs | 6 +- codex-rs/core/src/codex.rs | 15 +++- codex-rs/core/src/config.rs | 39 ++++++++-- codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/openai_model_info.rs | 71 +++++++++++++++++++ codex-rs/core/src/protocol.rs | 13 ++++ codex-rs/exec/src/event_processor.rs | 4 ++ codex-rs/mcp-server/src/codex_tool_runner.rs | 1 + codex-rs/tui/src/bottom_pane/chat_composer.rs | 40 ++++++++++- codex-rs/tui/src/bottom_pane/mod.rs | 13 ++++ codex-rs/tui/src/chatwidget.rs | 36 ++++++++++ 14 files changed, 301 insertions(+), 15 deletions(-) create mode 100644 codex-rs/core/src/openai_model_info.rs diff --git a/codex-rs/config.md b/codex-rs/config.md index 14d5fd2252..bb8b67162c 100644 --- a/codex-rs/config.md +++ b/codex-rs/config.md @@ -407,6 +407,16 @@ Setting `hide_agent_reasoning` to `true` suppresses these events in **both** the hide_agent_reasoning = true # defaults to false ``` +## model_context_window + +The size of the context window for the model, in tokens. + +In general, Codex knows the context window for the most common OpenAI models, but if you are using a new model with an old version of the Codex CLI, then you can use `model_context_window` to tell Codex what value to use to determine how much context is left during a conversation. + +## model_max_output_tokens + +This is analogous to `model_context_window`, but for the maximum number of output tokens for the model. + ## project_doc_max_bytes Maximum number of bytes to read from an `AGENTS.md` file to include in the instructions sent with the first turn of a session. Defaults to 32 KiB. diff --git a/codex-rs/core/src/chat_completions.rs b/codex-rs/core/src/chat_completions.rs index f381c72e51..12c5b7afca 100644 --- a/codex-rs/core/src/chat_completions.rs +++ b/codex-rs/core/src/chat_completions.rs @@ -215,6 +215,7 @@ where let _ = tx_event .send(Ok(ResponseEvent::Completed { response_id: String::new(), + token_usage: None, })) .await; return; @@ -232,6 +233,7 @@ where let _ = tx_event .send(Ok(ResponseEvent::Completed { response_id: String::new(), + token_usage: None, })) .await; return; @@ -317,6 +319,7 @@ where let _ = tx_event .send(Ok(ResponseEvent::Completed { response_id: String::new(), + token_usage: None, })) .await; @@ -394,7 +397,10 @@ where // Not an assistant message – forward immediately. return Poll::Ready(Some(Ok(ResponseEvent::OutputItemDone(item)))); } - Poll::Ready(Some(Ok(ResponseEvent::Completed { response_id }))) => { + Poll::Ready(Some(Ok(ResponseEvent::Completed { + response_id, + token_usage, + }))) => { if !this.cumulative.is_empty() { let aggregated_item = crate::models::ResponseItem::Message { role: "assistant".to_string(), @@ -404,7 +410,10 @@ where }; // Buffer Completed so it is returned *after* the aggregated message. - this.pending_completed = Some(ResponseEvent::Completed { response_id }); + this.pending_completed = Some(ResponseEvent::Completed { + response_id, + token_usage, + }); return Poll::Ready(Some(Ok(ResponseEvent::OutputItemDone( aggregated_item, @@ -412,7 +421,10 @@ where } // Nothing aggregated – forward Completed directly. - return Poll::Ready(Some(Ok(ResponseEvent::Completed { response_id }))); + return Poll::Ready(Some(Ok(ResponseEvent::Completed { + response_id, + token_usage, + }))); } // No other `Ok` variants exist at the moment, continue polling. } } diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index aff838887a..4770796dbb 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -35,6 +35,7 @@ use crate::model_provider_info::ModelProviderInfo; use crate::model_provider_info::WireApi; use crate::models::ResponseItem; use crate::openai_tools::create_tools_json_for_responses_api; +use crate::protocol::TokenUsage; use crate::util::backoff; #[derive(Clone)] @@ -210,6 +211,38 @@ struct SseEvent { #[derive(Debug, Deserialize)] struct ResponseCompleted { id: String, + usage: Option, +} + +#[derive(Debug, Deserialize)] +struct ResponseCompletedUsage { + input_tokens: u64, + input_tokens_details: Option, + output_tokens: u64, + output_tokens_details: Option, + total_tokens: u64, +} + +impl From for TokenUsage { + fn from(val: ResponseCompletedUsage) -> Self { + TokenUsage { + input_tokens: val.input_tokens, + cached_input_tokens: val.input_tokens_details.map(|d| d.cached_tokens), + output_tokens: val.output_tokens, + reasoning_output_tokens: val.output_tokens_details.map(|d| d.reasoning_tokens), + total_tokens: val.total_tokens, + } + } +} + +#[derive(Debug, Deserialize)] +struct ResponseCompletedInputTokensDetails { + cached_tokens: u64, +} + +#[derive(Debug, Deserialize)] +struct ResponseCompletedOutputTokensDetails { + reasoning_tokens: u64, } async fn process_sse(stream: S, tx_event: mpsc::Sender>) @@ -221,7 +254,7 @@ where // If the stream stays completely silent for an extended period treat it as disconnected. let idle_timeout = *OPENAI_STREAM_IDLE_TIMEOUT_MS; // The response id returned from the "complete" message. - let mut response_id = None; + let mut response_completed: Option = None; loop { let sse = match timeout(idle_timeout, stream.next()).await { @@ -233,9 +266,15 @@ where return; } Ok(None) => { - match response_id { - Some(response_id) => { - let event = ResponseEvent::Completed { response_id }; + match response_completed { + Some(ResponseCompleted { + id: response_id, + usage, + }) => { + let event = ResponseEvent::Completed { + response_id, + token_usage: usage.map(Into::into), + }; let _ = tx_event.send(Ok(event)).await; } None => { @@ -301,7 +340,7 @@ where if let Some(resp_val) = event.response { match serde_json::from_value::(resp_val) { Ok(r) => { - response_id = Some(r.id); + response_completed = Some(r); } Err(e) => { debug!("failed to parse ResponseCompleted: {e}"); diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index a2633475df..e17cf22c59 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -2,6 +2,7 @@ use crate::config_types::ReasoningEffort as ReasoningEffortConfig; use crate::config_types::ReasoningSummary as ReasoningSummaryConfig; use crate::error::Result; use crate::models::ResponseItem; +use crate::protocol::TokenUsage; use codex_apply_patch::APPLY_PATCH_TOOL_INSTRUCTIONS; use futures::Stream; use serde::Serialize; @@ -51,7 +52,10 @@ impl Prompt { #[derive(Debug)] pub enum ResponseEvent { OutputItemDone(ResponseItem), - Completed { response_id: String }, + Completed { + response_id: String, + token_usage: Option, + }, } #[derive(Debug, Serialize)] diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index e12a3a600b..a43f75a731 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1078,7 +1078,20 @@ async fn try_run_turn( let response = handle_response_item(sess, sub_id, item.clone()).await?; output.push(ProcessedResponseItem { item, response }); } - ResponseEvent::Completed { response_id } => { + ResponseEvent::Completed { + response_id, + token_usage, + } => { + if let Some(token_usage) = token_usage { + sess.tx_event + .send(Event { + id: sub_id.to_string(), + msg: EventMsg::TokenCount(token_usage), + }) + .await + .ok(); + } + let mut state = sess.state.lock().unwrap(); state.previous_response_id = Some(response_id); break; diff --git a/codex-rs/core/src/config.rs b/codex-rs/core/src/config.rs index e01bb3f423..6652d7c78d 100644 --- a/codex-rs/core/src/config.rs +++ b/codex-rs/core/src/config.rs @@ -10,6 +10,7 @@ use crate::config_types::UriBasedFileOpener; use crate::flags::OPENAI_DEFAULT_MODEL; use crate::model_provider_info::ModelProviderInfo; use crate::model_provider_info::built_in_model_providers; +use crate::openai_model_info::get_model_info; use crate::protocol::AskForApproval; use crate::protocol::SandboxPolicy; use dirs::home_dir; @@ -30,6 +31,12 @@ pub struct Config { /// Optional override of model selection. pub model: String, + /// Size of the context window for the model, in tokens. + pub model_context_window: Option, + + /// Maximum number of output tokens. + pub model_max_output_tokens: Option, + /// Key into the model_providers map that specifies which provider to use. pub model_provider_id: String, @@ -234,6 +241,12 @@ pub struct ConfigToml { /// Provider to use from the model_providers map. pub model_provider: Option, + /// Size of the context window for the model, in tokens. + pub model_context_window: Option, + + /// Maximum number of output tokens. + pub model_max_output_tokens: Option, + /// Default approval policy for executing commands. pub approval_policy: Option, @@ -387,11 +400,23 @@ impl Config { let history = cfg.history.unwrap_or_default(); + let model = model + .or(config_profile.model) + .or(cfg.model) + .unwrap_or_else(default_model); + let openai_model_info = get_model_info(&model); + let model_context_window = cfg + .model_context_window + .or_else(|| openai_model_info.as_ref().map(|info| info.context_window)); + let model_max_output_tokens = cfg.model_max_output_tokens.or_else(|| { + openai_model_info + .as_ref() + .map(|info| info.max_output_tokens) + }); let config = Self { - model: model - .or(config_profile.model) - .or(cfg.model) - .unwrap_or_else(default_model), + model, + model_context_window, + model_max_output_tokens, model_provider_id, model_provider, cwd: resolved_cwd, @@ -687,6 +712,8 @@ disable_response_storage = true assert_eq!( Config { model: "o3".to_string(), + model_context_window: Some(200_000), + model_max_output_tokens: Some(100_000), model_provider_id: "openai".to_string(), model_provider: fixture.openai_provider.clone(), approval_policy: AskForApproval::Never, @@ -729,6 +756,8 @@ disable_response_storage = true )?; let expected_gpt3_profile_config = Config { model: "gpt-3.5-turbo".to_string(), + model_context_window: Some(16_385), + model_max_output_tokens: Some(4_096), model_provider_id: "openai-chat-completions".to_string(), model_provider: fixture.openai_chat_completions_provider.clone(), approval_policy: AskForApproval::UnlessTrusted, @@ -786,6 +815,8 @@ disable_response_storage = true )?; let expected_zdr_profile_config = Config { model: "o3".to_string(), + model_context_window: Some(200_000), + model_max_output_tokens: Some(100_000), model_provider_id: "openai".to_string(), model_provider: fixture.openai_provider.clone(), approval_policy: AskForApproval::OnFailure, diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 16cf190588..6812260c97 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -28,6 +28,7 @@ pub use model_provider_info::ModelProviderInfo; pub use model_provider_info::WireApi; mod models; pub mod openai_api_key; +mod openai_model_info; mod openai_tools; mod project_doc; pub mod protocol; diff --git a/codex-rs/core/src/openai_model_info.rs b/codex-rs/core/src/openai_model_info.rs new file mode 100644 index 0000000000..9ffd831a91 --- /dev/null +++ b/codex-rs/core/src/openai_model_info.rs @@ -0,0 +1,71 @@ +/// Metadata about a model, particularly OpenAI models. +/// We may want to consider including details like the pricing for +/// input tokens, output tokens, etc., though users will need to be able to +/// override this in config.toml, as this information can get out of date. +/// Though this would help present more accurate pricing information in the UI. +#[derive(Debug)] +pub(crate) struct ModelInfo { + /// Size of the context window in tokens. + pub(crate) context_window: u64, + + /// Maximum number of output tokens that can be generated for the model. + pub(crate) max_output_tokens: u64, +} + +/// Note details such as what a model like gpt-4o is aliased to may be out of +/// date. +pub(crate) fn get_model_info(name: &str) -> Option { + match name { + // https://platform.openai.com/docs/models/o3 + "o3" => Some(ModelInfo { + context_window: 200_000, + max_output_tokens: 100_000, + }), + + // https://platform.openai.com/docs/models/o4-mini + "o4-mini" => Some(ModelInfo { + context_window: 200_000, + max_output_tokens: 100_000, + }), + + // https://platform.openai.com/docs/models/codex-mini-latest + "codex-mini-latest" => Some(ModelInfo { + context_window: 200_000, + max_output_tokens: 100_000, + }), + + // As of Jun 25, 2025, gpt-4.1 defaults to gpt-4.1-2025-04-14. + // https://platform.openai.com/docs/models/gpt-4.1 + "gpt-4.1" | "gpt-4.1-2025-04-14" => Some(ModelInfo { + context_window: 1_047_576, + max_output_tokens: 32_768, + }), + + // As of Jun 25, 2025, gpt-4o defaults to gpt-4o-2024-08-06. + // https://platform.openai.com/docs/models/gpt-4o + "gpt-4o" | "gpt-4o-2024-08-06" => Some(ModelInfo { + context_window: 128_000, + max_output_tokens: 16_384, + }), + + // https://platform.openai.com/docs/models/gpt-4o?snapshot=gpt-4o-2024-05-13 + "gpt-4o-2024-05-13" => Some(ModelInfo { + context_window: 128_000, + max_output_tokens: 4_096, + }), + + // https://platform.openai.com/docs/models/gpt-4o?snapshot=gpt-4o-2024-11-20 + "gpt-4o-2024-11-20" => Some(ModelInfo { + context_window: 128_000, + max_output_tokens: 16_384, + }), + + // https://platform.openai.com/docs/models/gpt-3.5-turbo + "gpt-3.5-turbo" => Some(ModelInfo { + context_window: 16_385, + max_output_tokens: 4_096, + }), + + _ => None, + } +} diff --git a/codex-rs/core/src/protocol.rs b/codex-rs/core/src/protocol.rs index d4aa769852..fa25a2fe38 100644 --- a/codex-rs/core/src/protocol.rs +++ b/codex-rs/core/src/protocol.rs @@ -275,6 +275,10 @@ pub enum EventMsg { /// Agent has completed all actions TaskComplete(TaskCompleteEvent), + /// Token count event, sent periodically to report the number of tokens + /// used in the current session. + TokenCount(TokenUsage), + /// Agent text output message AgentMessage(AgentMessageEvent), @@ -322,6 +326,15 @@ pub struct TaskCompleteEvent { pub last_agent_message: Option, } +#[derive(Debug, Clone, Deserialize, Serialize, Default)] +pub struct TokenUsage { + pub input_tokens: u64, + pub cached_input_tokens: Option, + pub output_tokens: u64, + pub reasoning_output_tokens: Option, + pub total_tokens: u64, +} + #[derive(Debug, Clone, Deserialize, Serialize)] pub struct AgentMessageEvent { pub message: String, diff --git a/codex-rs/exec/src/event_processor.rs b/codex-rs/exec/src/event_processor.rs index e2a8bbb20a..5320c572b9 100644 --- a/codex-rs/exec/src/event_processor.rs +++ b/codex-rs/exec/src/event_processor.rs @@ -16,6 +16,7 @@ use codex_core::protocol::McpToolCallEndEvent; use codex_core::protocol::PatchApplyBeginEvent; use codex_core::protocol::PatchApplyEndEvent; use codex_core::protocol::SessionConfiguredEvent; +use codex_core::protocol::TokenUsage; use owo_colors::OwoColorize; use owo_colors::Style; use shlex::try_join; @@ -180,6 +181,9 @@ impl EventProcessor { EventMsg::TaskStarted | EventMsg::TaskComplete(_) => { // Ignore. } + EventMsg::TokenCount(TokenUsage { total_tokens, .. }) => { + ts_println!(self, "tokens used: {total_tokens}"); + } EventMsg::AgentMessage(AgentMessageEvent { message }) => { ts_println!( self, diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs index 67c990b00c..796a119e5c 100644 --- a/codex-rs/mcp-server/src/codex_tool_runner.rs +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -162,6 +162,7 @@ pub async fn run_codex_tool_session( } EventMsg::Error(_) | EventMsg::TaskStarted + | EventMsg::TokenCount(_) | EventMsg::AgentReasoning(_) | EventMsg::McpToolCallBegin(_) | EventMsg::McpToolCallEnd(_) diff --git a/codex-rs/tui/src/bottom_pane/chat_composer.rs b/codex-rs/tui/src/bottom_pane/chat_composer.rs index 1218f76ec7..4ec8299081 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer.rs @@ -1,3 +1,4 @@ +use codex_core::protocol::TokenUsage; use crossterm::event::KeyEvent; use ratatui::buffer::Buffer; use ratatui::layout::Alignment; @@ -24,6 +25,8 @@ const MIN_TEXTAREA_ROWS: usize = 1; /// Rows consumed by the border. const BORDER_LINES: u16 = 2; +const BASE_PLACEHOLDER_TEXT: &str = "send a message"; + /// Result returned when the user interacts with the text area. pub enum InputResult { Submitted(String), @@ -40,7 +43,7 @@ pub(crate) struct ChatComposer<'a> { impl ChatComposer<'_> { pub fn new(has_input_focus: bool, app_event_tx: AppEventSender) -> Self { let mut textarea = TextArea::default(); - textarea.set_placeholder_text("send a message"); + textarea.set_placeholder_text(BASE_PLACEHOLDER_TEXT); textarea.set_cursor_line_style(ratatui::style::Style::default()); let mut this = Self { @@ -53,6 +56,41 @@ impl ChatComposer<'_> { this } + /// Update the cached *context-left* percentage and refresh the placeholder + /// text. The UI relies on the placeholder to convey the remaining + /// context when the composer is empty. + pub(crate) fn set_token_usage( + &mut self, + token_usage: TokenUsage, + model_context_window: Option, + ) { + let placeholder = match (token_usage.total_tokens, model_context_window) { + (total_tokens, Some(context_window)) => { + let percent_remaining: u8 = if context_window > 0 { + // Calculate the percentage of context left. + let percent = 100.0 - (total_tokens as f32 / context_window as f32 * 100.0); + percent.clamp(0.0, 100.0) as u8 + } else { + // If we don't have a context window, we cannot compute the + // percentage. + 100 + }; + if percent_remaining > 25 { + format!("{BASE_PLACEHOLDER_TEXT} — {percent_remaining}% context left") + } else { + format!( + "{BASE_PLACEHOLDER_TEXT} — {percent_remaining}% context left (consider /compact)" + ) + } + } + (total_tokens, None) => { + format!("{BASE_PLACEHOLDER_TEXT} — {total_tokens} tokens used") + } + }; + + self.textarea.set_placeholder_text(placeholder); + } + /// Record the history metadata advertised by `SessionConfiguredEvent` so /// that the composer can navigate cross-session history. pub(crate) fn set_history_metadata(&mut self, log_id: u64, entry_count: usize) { diff --git a/codex-rs/tui/src/bottom_pane/mod.rs b/codex-rs/tui/src/bottom_pane/mod.rs index c654581ccd..e3234e99a6 100644 --- a/codex-rs/tui/src/bottom_pane/mod.rs +++ b/codex-rs/tui/src/bottom_pane/mod.rs @@ -2,6 +2,7 @@ use bottom_pane_view::BottomPaneView; use bottom_pane_view::ConditionalUpdate; +use codex_core::protocol::TokenUsage; use crossterm::event::KeyEvent; use ratatui::buffer::Buffer; use ratatui::layout::Rect; @@ -129,6 +130,18 @@ impl BottomPane<'_> { } } + /// Update the *context-window remaining* indicator in the composer. This + /// is forwarded directly to the underlying `ChatComposer`. + pub(crate) fn set_token_usage( + &mut self, + token_usage: TokenUsage, + model_context_window: Option, + ) { + self.composer + .set_token_usage(token_usage, model_context_window); + self.request_redraw(); + } + /// Called when the agent requests user approval. pub fn push_approval_request(&mut self, request: ApprovalRequest) { let request = if let Some(view) = self.active_view.as_mut() { diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index bd5197c73b..fad72e3ab9 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -18,6 +18,7 @@ use codex_core::protocol::McpToolCallEndEvent; use codex_core::protocol::Op; use codex_core::protocol::PatchApplyBeginEvent; use codex_core::protocol::TaskCompleteEvent; +use codex_core::protocol::TokenUsage; use crossterm::event::KeyEvent; use ratatui::buffer::Buffer; use ratatui::layout::Constraint; @@ -46,6 +47,7 @@ pub(crate) struct ChatWidget<'a> { input_focus: InputFocus, config: Config, initial_user_message: Option, + token_usage: TokenUsage, } #[derive(Clone, Copy, Eq, PartialEq)] @@ -131,6 +133,7 @@ impl ChatWidget<'_> { initial_prompt.unwrap_or_default(), initial_images, ), + token_usage: TokenUsage::default(), } } @@ -250,6 +253,11 @@ impl ChatWidget<'_> { self.bottom_pane.set_task_running(false); self.request_redraw(); } + EventMsg::TokenCount(token_usage) => { + self.token_usage = add_token_usage(&self.token_usage, &token_usage); + self.bottom_pane + .set_token_usage(self.token_usage.clone(), self.config.model_context_window); + } EventMsg::Error(ErrorEvent { message }) => { self.conversation_history.add_error(message); self.bottom_pane.set_task_running(false); @@ -410,3 +418,31 @@ impl WidgetRef for &ChatWidget<'_> { (&self.bottom_pane).render(chunks[1], buf); } } + +fn add_token_usage(current_usage: &TokenUsage, new_usage: &TokenUsage) -> TokenUsage { + let cached_input_tokens = match ( + current_usage.cached_input_tokens, + new_usage.cached_input_tokens, + ) { + (Some(current), Some(new)) => Some(current + new), + (Some(current), None) => Some(current), + (None, Some(new)) => Some(new), + (None, None) => None, + }; + let reasoning_output_tokens = match ( + current_usage.reasoning_output_tokens, + new_usage.reasoning_output_tokens, + ) { + (Some(current), Some(new)) => Some(current + new), + (Some(current), None) => Some(current), + (None, Some(new)) => Some(new), + (None, None) => None, + }; + TokenUsage { + input_tokens: current_usage.input_tokens + new_usage.input_tokens, + cached_input_tokens, + output_tokens: current_usage.output_tokens + new_usage.output_tokens, + reasoning_output_tokens, + total_tokens: current_usage.total_tokens + new_usage.total_tokens, + } +} From a339a7bcce153974d7f590b38f6e53dd01c2cc66 Mon Sep 17 00:00:00 2001 From: Gabriel Peal Date: Thu, 26 Jun 2025 14:40:42 -0400 Subject: [PATCH 3/4] [Rust] Allow resuming a session that was killed with ctrl + c (#1387) Previously, if you ctrl+c'd a conversation, all subsequent turns would 400 because the Responses API never got a response for one of its call ids. This ensures that if we aren't sending a call id by hand, we generate a synthetic aborted call. Fixes #1244 https://github.com/user-attachments/assets/5126354f-b970-45f5-8c65-f811bca8294a --- codex-rs/core/src/chat_completions.rs | 9 ++- codex-rs/core/src/client.rs | 11 ++- codex-rs/core/src/client_common.rs | 1 + codex-rs/core/src/codex.rs | 106 ++++++++++++++++++++++---- 4 files changed, 108 insertions(+), 19 deletions(-) diff --git a/codex-rs/core/src/chat_completions.rs b/codex-rs/core/src/chat_completions.rs index 12c5b7afca..dfe06d1fec 100644 --- a/codex-rs/core/src/chat_completions.rs +++ b/codex-rs/core/src/chat_completions.rs @@ -425,7 +425,12 @@ where response_id, token_usage, }))); - } // No other `Ok` variants exist at the moment, continue polling. + } + Poll::Ready(Some(Ok(ResponseEvent::Created))) => { + // These events are exclusive to the Responses API and + // will never appear in a Chat Completions stream. + continue; + } } } } @@ -439,7 +444,7 @@ pub(crate) trait AggregateStreamExt: Stream> + Size /// /// ```ignore /// OutputItemDone() - /// Completed { .. } + /// Completed /// ``` /// /// No other `OutputItemDone` events will be seen by the caller. diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 4770796dbb..6daa3a8969 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -168,7 +168,7 @@ impl ModelClient { // negligible. if !(status == StatusCode::TOO_MANY_REQUESTS || status.is_server_error()) { // Surface the error body to callers. Use `unwrap_or_default` per Clippy. - let body = (res.text().await).unwrap_or_default(); + let body = res.text().await.unwrap_or_default(); return Err(CodexErr::UnexpectedStatus(status, body)); } @@ -208,6 +208,9 @@ struct SseEvent { item: Option, } +#[derive(Debug, Deserialize)] +struct ResponseCreated {} + #[derive(Debug, Deserialize)] struct ResponseCompleted { id: String, @@ -335,6 +338,11 @@ where return; } } + "response.created" => { + if event.response.is_some() { + let _ = tx_event.send(Ok(ResponseEvent::Created {})).await; + } + } // Final response completed – includes array of output items & id "response.completed" => { if let Some(resp_val) = event.response { @@ -350,7 +358,6 @@ where }; } "response.content_part.done" - | "response.created" | "response.function_call_arguments.delta" | "response.in_progress" | "response.output_item.added" diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index e17cf22c59..b08880a0df 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -51,6 +51,7 @@ impl Prompt { #[derive(Debug)] pub enum ResponseEvent { + Created, OutputItemDone(ResponseItem), Completed { response_id: String, diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index a43f75a731..ec6e0bd185 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1,6 +1,7 @@ // Poisoned mutex should fail the program #![allow(clippy::unwrap_used)] +use std::borrow::Cow; use std::collections::HashMap; use std::collections::HashSet; use std::path::Path; @@ -188,7 +189,7 @@ pub(crate) struct Session { /// Optional rollout recorder for persisting the conversation transcript so /// sessions can be replayed or inspected later. - rollout: Mutex>, + rollout: Mutex>, state: Mutex, codex_linux_sandbox_exe: Option, } @@ -206,6 +207,9 @@ impl Session { struct State { approved_commands: HashSet>, current_task: Option, + /// Call IDs that have been sent from the Responses API but have not been sent back yet. + /// You CANNOT send a Responses API follow-up message unless you have sent back the output for all pending calls or else it will 400. + pending_call_ids: HashSet, previous_response_id: Option, pending_approvals: HashMap>, pending_input: Vec, @@ -312,7 +316,7 @@ impl Session { /// Append the given items to the session's rollout transcript (if enabled) /// and persist them to disk. async fn record_rollout_items(&self, items: &[ResponseItem]) { - // Clone the recorder outside of the mutex so we don’t hold the lock + // Clone the recorder outside of the mutex so we don't hold the lock // across an await point (MutexGuard is not Send). let recorder = { let guard = self.rollout.lock().unwrap(); @@ -411,6 +415,8 @@ impl Session { pub fn abort(&self) { info!("Aborting existing session"); let mut state = self.state.lock().unwrap(); + // Don't clear pending_call_ids because we need to keep track of them to ensure we don't 400 on the next turn. + // We will generate a synthetic aborted response for each pending call id. state.pending_approvals.clear(); state.pending_input.clear(); if let Some(task) = state.current_task.take() { @@ -431,7 +437,7 @@ impl Session { } let Ok(json) = serde_json::to_string(¬ification) else { - tracing::error!("failed to serialise notification payload"); + error!("failed to serialise notification payload"); return; }; @@ -443,7 +449,7 @@ impl Session { // Fire-and-forget – we do not wait for completion. if let Err(e) = command.spawn() { - tracing::warn!("failed to spawn notifier '{}': {e}", notify_command[0]); + warn!("failed to spawn notifier '{}': {e}", notify_command[0]); } } } @@ -647,7 +653,7 @@ async fn submission_loop( match RolloutRecorder::new(&config, session_id, instructions.clone()).await { Ok(r) => Some(r), Err(e) => { - tracing::warn!("failed to initialise rollout recorder: {e}"); + warn!("failed to initialise rollout recorder: {e}"); None } }; @@ -742,7 +748,7 @@ async fn submission_loop( tokio::spawn(async move { if let Err(e) = crate::message_history::append_entry(&text, &id, &config).await { - tracing::warn!("failed to append to message history: {e}"); + warn!("failed to append to message history: {e}"); } }); } @@ -772,7 +778,7 @@ async fn submission_loop( }; if let Err(e) = tx_event.send(event).await { - tracing::warn!("failed to send GetHistoryEntryResponse event: {e}"); + warn!("failed to send GetHistoryEntryResponse event: {e}"); } }); } @@ -1052,6 +1058,7 @@ async fn run_turn( /// events map to a `ResponseItem`. A `ResponseItem` may need to be /// "handled" such that it produces a `ResponseInputItem` that needs to be /// sent back to the model on the next turn. +#[derive(Debug)] struct ProcessedResponseItem { item: ResponseItem, response: Option, @@ -1062,7 +1069,57 @@ async fn try_run_turn( sub_id: &str, prompt: &Prompt, ) -> CodexResult> { - let mut stream = sess.client.clone().stream(prompt).await?; + // call_ids that are part of this response. + let completed_call_ids = prompt + .input + .iter() + .filter_map(|ri| match ri { + ResponseItem::FunctionCallOutput { call_id, .. } => Some(call_id), + ResponseItem::LocalShellCall { + call_id: Some(call_id), + .. + } => Some(call_id), + _ => None, + }) + .collect::>(); + + // call_ids that were pending but are not part of this response. + // This usually happens because the user interrupted the model before we responded to one of its tool calls + // and then the user sent a follow-up message. + let missing_calls = { + sess.state + .lock() + .unwrap() + .pending_call_ids + .iter() + .filter_map(|call_id| { + if completed_call_ids.contains(&call_id) { + None + } else { + Some(call_id.clone()) + } + }) + .map(|call_id| ResponseItem::FunctionCallOutput { + call_id: call_id.clone(), + output: FunctionCallOutputPayload { + content: "aborted".to_string(), + success: Some(false), + }, + }) + .collect::>() + }; + let prompt: Cow = if missing_calls.is_empty() { + Cow::Borrowed(prompt) + } else { + // Add the synthetic aborted missing calls to the beginning of the input to ensure all call ids have responses. + let input = [missing_calls, prompt.input.clone()].concat(); + Cow::Owned(Prompt { + input, + ..prompt.clone() + }) + }; + + let mut stream = sess.client.clone().stream(&prompt).await?; // Buffer all the incoming messages from the stream first, then execute them. // If we execute a function call in the middle of handling the stream, it can time out. @@ -1074,8 +1131,27 @@ async fn try_run_turn( let mut output = Vec::new(); for event in input { match event { + ResponseEvent::Created => { + let mut state = sess.state.lock().unwrap(); + // We successfully created a new response and ensured that all pending calls were included so we can clear the pending call ids. + state.pending_call_ids.clear(); + } ResponseEvent::OutputItemDone(item) => { + let call_id = match &item { + ResponseItem::LocalShellCall { + call_id: Some(call_id), + .. + } => Some(call_id), + ResponseItem::FunctionCall { call_id, .. } => Some(call_id), + _ => None, + }; + if let Some(call_id) = call_id { + // We just got a new call id so we need to make sure to respond to it in the next turn. + let mut state = sess.state.lock().unwrap(); + state.pending_call_ids.insert(call_id.clone()); + } let response = handle_response_item(sess, sub_id, item.clone()).await?; + output.push(ProcessedResponseItem { item, response }); } ResponseEvent::Completed { @@ -1138,7 +1214,7 @@ async fn handle_response_item( arguments, call_id, } => { - tracing::info!("FunctionCall: {arguments}"); + info!("FunctionCall: {arguments}"); Some(handle_function_call(sess, sub_id.to_string(), name, arguments, call_id).await) } ResponseItem::LocalShellCall { @@ -1220,7 +1296,7 @@ async fn handle_function_call( // Unknown function: reply with structured failure so the model can adapt. ResponseInputItem::FunctionCallOutput { call_id, - output: crate::models::FunctionCallOutputPayload { + output: FunctionCallOutputPayload { content: format!("unsupported call: {}", name), success: None, }, @@ -1252,7 +1328,7 @@ fn parse_container_exec_arguments( // allow model to re-sample let output = ResponseInputItem::FunctionCallOutput { call_id: call_id.to_string(), - output: crate::models::FunctionCallOutputPayload { + output: FunctionCallOutputPayload { content: format!("failed to parse function arguments: {e}"), success: None, }, @@ -1320,7 +1396,7 @@ async fn handle_container_exec_with_params( ReviewDecision::Denied | ReviewDecision::Abort => { return ResponseInputItem::FunctionCallOutput { call_id, - output: crate::models::FunctionCallOutputPayload { + output: FunctionCallOutputPayload { content: "exec command rejected by user".to_string(), success: None, }, @@ -1336,7 +1412,7 @@ async fn handle_container_exec_with_params( SafetyCheck::Reject { reason } => { return ResponseInputItem::FunctionCallOutput { call_id, - output: crate::models::FunctionCallOutputPayload { + output: FunctionCallOutputPayload { content: format!("exec command rejected: {reason}"), success: None, }, @@ -1870,7 +1946,7 @@ fn apply_changes_from_apply_patch(action: &ApplyPatchAction) -> anyhow::Result Vec { +fn get_writable_roots(cwd: &Path) -> Vec { let mut writable_roots = Vec::new(); if cfg!(target_os = "macos") { // On macOS, $TMPDIR is private to the user. @@ -1898,7 +1974,7 @@ fn get_writable_roots(cwd: &Path) -> Vec { } /// Exec output is a pre-serialized JSON payload -fn format_exec_output(output: &str, exit_code: i32, duration: std::time::Duration) -> String { +fn format_exec_output(output: &str, exit_code: i32, duration: Duration) -> String { #[derive(Serialize)] struct ExecMetadata { exit_code: i32, From 5c3ff9fc109721e6bcb709315a3cd60bcba50baf Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 26 Jun 2025 12:10:49 -0700 Subject: [PATCH 4/4] feat: add support for /diff command --- codex-rs/tui/src/app.rs | 23 ++++ codex-rs/tui/src/bottom_pane/command_popup.rs | 29 ++--- codex-rs/tui/src/chatwidget.rs | 8 ++ codex-rs/tui/src/get_git_diff.rs | 114 ++++++++++++++++++ codex-rs/tui/src/history_cell.rs | 9 +- codex-rs/tui/src/lib.rs | 1 + codex-rs/tui/src/slash_command.rs | 16 ++- 7 files changed, 176 insertions(+), 24 deletions(-) create mode 100644 codex-rs/tui/src/get_git_diff.rs diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index ff61b5c941..ecfa513b2b 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -1,6 +1,7 @@ use crate::app_event::AppEvent; use crate::app_event_sender::AppEventSender; use crate::chatwidget::ChatWidget; +use crate::get_git_diff::get_git_diff; use crate::git_warning_screen::GitWarningOutcome; use crate::git_warning_screen::GitWarningScreen; use crate::login_screen::LoginScreen; @@ -250,6 +251,28 @@ impl<'a> App<'a> { SlashCommand::Quit => { break; } + SlashCommand::Diff => { + let (is_repo, diff_text) = match get_git_diff() { + Ok(v) => v, + Err(e) => { + let msg = format!("Failed to compute diff: {e}"); + if let AppState::Chat { widget } = &mut self.app_state { + widget.add_background_event(msg); + } + continue; + } + }; + + let text = if is_repo { + diff_text + } else { + "`/diff` — _not inside a git repository_".to_string() + }; + + if let AppState::Chat { widget } = &mut self.app_state { + widget.add_background_event(text); + } + } }, } } diff --git a/codex-rs/tui/src/bottom_pane/command_popup.rs b/codex-rs/tui/src/bottom_pane/command_popup.rs index 0dcb98865c..fd865047ef 100644 --- a/codex-rs/tui/src/bottom_pane/command_popup.rs +++ b/codex-rs/tui/src/bottom_pane/command_popup.rs @@ -1,5 +1,3 @@ -use std::collections::HashMap; - use ratatui::buffer::Buffer; use ratatui::layout::Rect; use ratatui::style::Color; @@ -25,7 +23,7 @@ use ratatui::style::Modifier; pub(crate) struct CommandPopup { command_filter: String, - all_commands: HashMap<&'static str, SlashCommand>, + all_commands: Vec<(&'static str, SlashCommand)>, selected_idx: Option, } @@ -84,23 +82,20 @@ impl CommandPopup { /// Return the list of commands that match the current filter. Matching is /// performed using a *prefix* comparison on the command name. fn filtered_commands(&self) -> Vec<&SlashCommand> { - let mut cmds: Vec<&SlashCommand> = self - .all_commands - .values() - .filter(|cmd| { - if self.command_filter.is_empty() { - true - } else { - cmd.command() + self.all_commands + .iter() + .filter_map(|(_name, cmd)| { + if self.command_filter.is_empty() + || cmd + .command() .starts_with(&self.command_filter.to_ascii_lowercase()) + { + Some(cmd) + } else { + None } }) - .collect(); - - // Sort the commands alphabetically so the order is stable and - // predictable. - cmds.sort_by(|a, b| a.command().cmp(b.command())); - cmds + .collect::>() } /// Move the selection cursor one step up. diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index fad72e3ab9..6850ee761a 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -384,6 +384,14 @@ impl ChatWidget<'_> { self.app_event_tx.send(AppEvent::Redraw); } + /// Inject a background event into the conversation history. This is used + /// for displaying informational messages that originate from the UI + /// itself (e.g. the `/diff` command) rather than from the backend agent. + pub(crate) fn add_background_event(&mut self, message: String) { + self.conversation_history.add_background_event(message); + self.request_redraw(); + } + pub(crate) fn handle_scroll_delta(&mut self, scroll_delta: i32) { // If the user is trying to scroll exactly one line, we let them, but // otherwise we assume they are trying to scroll in larger increments. diff --git a/codex-rs/tui/src/get_git_diff.rs b/codex-rs/tui/src/get_git_diff.rs new file mode 100644 index 0000000000..ff89fdcf1e --- /dev/null +++ b/codex-rs/tui/src/get_git_diff.rs @@ -0,0 +1,114 @@ +//! Utility to compute the current Git diff for the working directory. +//! +//! The implementation mirrors the behaviour of the TypeScript version in +//! `codex-cli`: it returns the diff for tracked changes as well as any +//! untracked files. When the current directory is not inside a Git +//! repository, the function returns `Ok((false, String::new()))`. + +use std::io; +use std::path::Path; +use std::process::Command; +use std::process::Stdio; + +/// Return value of [`get_git_diff`]. +/// +/// * `bool` – Whether the current working directory is inside a Git repo. +/// * `String` – The concatenated diff (may be empty). +pub(crate) fn get_git_diff() -> io::Result<(bool, String)> { + // First check if we are inside a Git repository. + if !inside_git_repo()? { + return Ok((false, String::new())); + } + + // 1. Diff for tracked files. + let tracked_diff = run_git_capture_diff(&["diff", "--color"])?; + + // 2. Determine untracked files. + let untracked_output = run_git_capture_stdout(&["ls-files", "--others", "--exclude-standard"])?; + + let mut untracked_diff = String::new(); + let null_device: &Path = if cfg!(windows) { + Path::new("NUL") + } else { + Path::new("/dev/null") + }; + + for file in untracked_output + .split('\n') + .map(str::trim) + .filter(|s| !s.is_empty()) + { + // Use `git diff --no-index` to generate a diff against the null device. + let args = [ + "diff", + "--color", + "--no-index", + "--", + null_device.to_str().unwrap_or("/dev/null"), + file, + ]; + + match run_git_capture_diff(&args) { + Ok(diff) => untracked_diff.push_str(&diff), + // If the file disappeared between ls-files and diff we ignore the error. + Err(err) if err.kind() == io::ErrorKind::NotFound => {} + Err(err) => return Err(err), + } + } + + Ok((true, format!("{}{}", tracked_diff, untracked_diff))) +} + +/// Helper that executes `git` with the given `args` and returns `stdout` as a +/// UTF-8 string. Any non-zero exit status is considered an *error*. +fn run_git_capture_stdout(args: &[&str]) -> io::Result { + let output = Command::new("git") + .args(args) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .output()?; + + if output.status.success() { + Ok(String::from_utf8_lossy(&output.stdout).into_owned()) + } else { + Err(io::Error::other(format!( + "git {:?} failed with status {}", + args, output.status + ))) + } +} + +/// Like [`run_git_capture_stdout`] but treats exit status 1 as success and +/// returns stdout. Git returns 1 for diffs when differences are present. +fn run_git_capture_diff(args: &[&str]) -> io::Result { + let output = Command::new("git") + .args(args) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .output()?; + + if output.status.success() || output.status.code() == Some(1) { + Ok(String::from_utf8_lossy(&output.stdout).into_owned()) + } else { + Err(io::Error::other(format!( + "git {:?} failed with status {}", + args, output.status + ))) + } +} + +/// Determine if the current directory is inside a Git repository. +fn inside_git_repo() -> io::Result { + let status = Command::new("git") + .args(["rev-parse", "--is-inside-work-tree"]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); + + match status { + Ok(s) if s.success() => Ok(true), + Ok(_) => Ok(false), + Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(false), // git not installed + Err(e) => Err(e), + } +} diff --git a/codex-rs/tui/src/history_cell.rs b/codex-rs/tui/src/history_cell.rs index e2a54283c1..e3707f3e61 100644 --- a/codex-rs/tui/src/history_cell.rs +++ b/codex-rs/tui/src/history_cell.rs @@ -453,7 +453,14 @@ impl HistoryCell { pub(crate) fn new_background_event(message: String) -> Self { let mut lines: Vec> = Vec::new(); lines.push(Line::from("event".dim())); - lines.extend(message.lines().map(|l| Line::from(l.to_string()).dim())); + + for raw in message.lines() { + // Parse ANSI color sequences so they render correctly in Ratatui. + // We preserve any colors encoded in the input; additionally mark + // the text as dim to distinguish background events from regular + // conversation. + lines.push(ansi_escape_line(raw).dim()); + } lines.push(Line::from("")); HistoryCell::BackgroundEvent { view: TextBlock::new(lines), diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 156951fff4..b17bb0421b 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -29,6 +29,7 @@ mod citation_regex; mod cli; mod conversation_history_widget; mod exec_command; +mod get_git_diff; mod git_warning_screen; mod history_cell; mod log_layer; diff --git a/codex-rs/tui/src/slash_command.rs b/codex-rs/tui/src/slash_command.rs index bfc02ceb13..0ae68fe7f1 100644 --- a/codex-rs/tui/src/slash_command.rs +++ b/codex-rs/tui/src/slash_command.rs @@ -1,7 +1,5 @@ -use std::collections::HashMap; - use strum::IntoEnumIterator; -use strum_macros::AsRefStr; // derive macro +use strum_macros::AsRefStr; use strum_macros::EnumIter; use strum_macros::EnumString; use strum_macros::IntoStaticStr; @@ -12,9 +10,12 @@ use strum_macros::IntoStaticStr; )] #[strum(serialize_all = "kebab-case")] pub enum SlashCommand { + // DO NOT ALPHA-SORT! Enum order is presentation order in the popup, so + // more frequently used commands should be listed first. New, - ToggleMouseMode, + Diff, Quit, + ToggleMouseMode, } impl SlashCommand { @@ -26,6 +27,9 @@ impl SlashCommand { "Toggle mouse mode (enable for scrolling, disable for text selection)" } SlashCommand::Quit => "Exit the application.", + SlashCommand::Diff => { + "Show git diff of the working directory (including untracked files)" + } } } @@ -36,7 +40,7 @@ impl SlashCommand { } } -/// Return all built-in commands in a HashMap keyed by their command string. -pub fn built_in_slash_commands() -> HashMap<&'static str, SlashCommand> { +/// Return all built-in commands in a BTreeMap keyed by their command string. +pub fn built_in_slash_commands() -> Vec<(&'static str, SlashCommand)> { SlashCommand::iter().map(|c| (c.command(), c)).collect() }