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/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..dfe06d1fec 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,8 +421,16 @@ where } // Nothing aggregated – forward Completed directly. - return Poll::Ready(Some(Ok(ResponseEvent::Completed { response_id }))); - } // No other `Ok` variants exist at the moment, continue polling. + return Poll::Ready(Some(Ok(ResponseEvent::Completed { + response_id, + token_usage, + }))); + } + Poll::Ready(Some(Ok(ResponseEvent::Created))) => { + // These events are exclusive to the Responses API and + // will never appear in a Chat Completions stream. + continue; + } } } } @@ -427,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 aff838887a..6daa3a8969 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)] @@ -167,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)); } @@ -207,9 +208,44 @@ struct SseEvent { item: Option, } +#[derive(Debug, Deserialize)] +struct ResponseCreated {} + #[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 +257,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 +269,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 => { @@ -296,12 +338,17 @@ 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 { 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}"); @@ -311,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 a2633475df..b08880a0df 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; @@ -50,8 +51,12 @@ impl Prompt { #[derive(Debug)] pub enum ResponseEvent { + Created, 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..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,11 +1131,43 @@ 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 { 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; @@ -1125,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 { @@ -1207,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, }, @@ -1239,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, }, @@ -1307,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, }, @@ -1323,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, }, @@ -1857,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. @@ -1885,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, 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/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 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 2c705c7c55..6850ee761a 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); @@ -418,3 +426,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, + } +}