diff --git a/codex-rs/agent-graph-store/src/error.rs b/codex-rs/agent-graph-store/src/error.rs new file mode 100644 index 0000000000000000000000000000000000000000..ddd8eeef380d4446baf8dd430ec1f984cc3e1ddd --- /dev/null +++ b/codex-rs/agent-graph-store/src/error.rs @@ -0,0 +1,20 @@ +/// Result type returned by agent graph store operations. +pub type AgentGraphStoreResult = Result; + +/// Error type shared by agent graph store implementations. +#[derive(Debug, thiserror::Error)] +pub enum AgentGraphStoreError { + /// The caller supplied invalid request data. + #[error("invalid agent graph store request: {message}")] + InvalidRequest { + /// User-facing explanation of the invalid request. + message: String, + }, + + /// Catch-all for implementation failures that do not fit a more specific category. + #[error("agent graph store internal error: {message}")] + Internal { + /// User-facing explanation of the implementation failure. + message: String, + }, +} diff --git a/codex-rs/agent-graph-store/src/lib.rs b/codex-rs/agent-graph-store/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..d5f40331b255f509349a13e8e70e8ccfa1ebf5bf --- /dev/null +++ b/codex-rs/agent-graph-store/src/lib.rs @@ -0,0 +1,13 @@ +//! Storage-neutral parent/child topology for thread-spawned agents. + +mod error; +mod local; +mod store; +mod types; + +pub use error::AgentGraphStoreError; +pub use error::AgentGraphStoreResult; +pub use local::LocalAgentGraphStore; +pub use store::AgentGraphStore; +pub use store::AgentGraphStoreFuture; +pub use types::ThreadSpawnEdgeStatus; diff --git a/codex-rs/agent-graph-store/src/local.rs b/codex-rs/agent-graph-store/src/local.rs new file mode 100644 index 0000000000000000000000000000000000000000..a7c1fd4a3395e9070f65cf7cbeff45f402c0f255 --- /dev/null +++ b/codex-rs/agent-graph-store/src/local.rs @@ -0,0 +1,344 @@ +use codex_protocol::ThreadId; +use codex_state::StateRuntime; +use std::sync::Arc; + +use crate::AgentGraphStore; +use crate::AgentGraphStoreError; +use crate::AgentGraphStoreFuture; +use crate::ThreadSpawnEdgeStatus; + +/// SQLite-backed implementation of [`AgentGraphStore`] using an existing state runtime. +#[derive(Clone)] +pub struct LocalAgentGraphStore { + state_db: Arc, +} + +impl std::fmt::Debug for LocalAgentGraphStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LocalAgentGraphStore") + .field("sqlite", self.state_db.sqlite()) + .finish_non_exhaustive() + } +} + +impl LocalAgentGraphStore { + /// Create a local graph store from an already-initialized state runtime. + pub fn new(state_db: Arc) -> Self { + Self { state_db } + } +} + +impl AgentGraphStore for LocalAgentGraphStore { + fn upsert_thread_spawn_edge( + &self, + parent_thread_id: ThreadId, + child_thread_id: ThreadId, + status: ThreadSpawnEdgeStatus, + ) -> AgentGraphStoreFuture<'_, ()> { + Box::pin(async move { + self.state_db + .upsert_thread_spawn_edge( + parent_thread_id, + child_thread_id, + to_state_status(status), + ) + .await + .map_err(internal_error) + }) + } + + fn set_thread_spawn_edge_status( + &self, + child_thread_id: ThreadId, + status: ThreadSpawnEdgeStatus, + ) -> AgentGraphStoreFuture<'_, ()> { + Box::pin(async move { + self.state_db + .set_thread_spawn_edge_status(child_thread_id, to_state_status(status)) + .await + .map_err(internal_error) + }) + } + + fn list_thread_spawn_children( + &self, + parent_thread_id: ThreadId, + status_filter: Option, + ) -> AgentGraphStoreFuture<'_, Vec> { + Box::pin(async move { + if let Some(status) = status_filter { + return self + .state_db + .list_thread_spawn_children_with_status( + parent_thread_id, + to_state_status(status), + ) + .await + .map_err(internal_error); + } + + self.state_db + .list_thread_spawn_children(parent_thread_id) + .await + .map_err(internal_error) + }) + } + + fn list_thread_spawn_descendants( + &self, + root_thread_id: ThreadId, + status_filter: Option, + ) -> AgentGraphStoreFuture<'_, Vec> { + Box::pin(async move { + match status_filter { + Some(status) => self + .state_db + .list_thread_spawn_descendants_with_status( + root_thread_id, + to_state_status(status), + ) + .await + .map_err(internal_error), + None => self + .state_db + .list_thread_spawn_descendants(root_thread_id) + .await + .map_err(internal_error), + } + }) + } +} + +fn to_state_status(status: ThreadSpawnEdgeStatus) -> codex_state::DirectionalThreadSpawnEdgeStatus { + match status { + ThreadSpawnEdgeStatus::Open => codex_state::DirectionalThreadSpawnEdgeStatus::Open, + ThreadSpawnEdgeStatus::Closed => codex_state::DirectionalThreadSpawnEdgeStatus::Closed, + } +} + +fn internal_error(err: impl std::fmt::Display) -> AgentGraphStoreError { + AgentGraphStoreError::Internal { + message: err.to_string(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use codex_state::DirectionalThreadSpawnEdgeStatus; + use codex_utils_absolute_path::test_support::PathExt; + use pretty_assertions::assert_eq; + use tempfile::TempDir; + + struct TestRuntime { + state_db: Arc, + _codex_home: TempDir, + } + + fn thread_id(suffix: u128) -> ThreadId { + ThreadId::from_string(&format!("00000000-0000-0000-0000-{suffix:012}")) + .expect("valid thread id") + } + + async fn state_runtime() -> TestRuntime { + let codex_home = TempDir::new().expect("tempdir should be created"); + let state_db = StateRuntime::init( + codex_state::SqliteConfig::new_for_testing(codex_home.path().abs()), + "test-provider".to_string(), + ) + .await + .expect("state db should initialize"); + TestRuntime { + state_db, + _codex_home: codex_home, + } + } + + #[tokio::test] + async fn local_store_upserts_and_lists_direct_children_with_status_filters() { + let fixture = state_runtime().await; + let state_db = fixture.state_db; + let store = LocalAgentGraphStore::new(state_db.clone()); + let parent_thread_id = thread_id(/*suffix*/ 1); + let first_child_thread_id = thread_id(/*suffix*/ 2); + let second_child_thread_id = thread_id(/*suffix*/ 3); + + store + .upsert_thread_spawn_edge( + parent_thread_id, + second_child_thread_id, + ThreadSpawnEdgeStatus::Closed, + ) + .await + .expect("closed child edge should insert"); + store + .upsert_thread_spawn_edge( + parent_thread_id, + first_child_thread_id, + ThreadSpawnEdgeStatus::Open, + ) + .await + .expect("open child edge should insert"); + + let all_children = store + .list_thread_spawn_children(parent_thread_id, /*status_filter*/ None) + .await + .expect("all children should load"); + assert_eq!( + all_children, + vec![first_child_thread_id, second_child_thread_id] + ); + + let open_children = store + .list_thread_spawn_children(parent_thread_id, Some(ThreadSpawnEdgeStatus::Open)) + .await + .expect("open children should load"); + let state_open_children = state_db + .list_thread_spawn_children_with_status( + parent_thread_id, + DirectionalThreadSpawnEdgeStatus::Open, + ) + .await + .expect("state open children should load"); + assert_eq!(open_children, state_open_children); + assert_eq!(open_children, vec![first_child_thread_id]); + + let closed_children = store + .list_thread_spawn_children(parent_thread_id, Some(ThreadSpawnEdgeStatus::Closed)) + .await + .expect("closed children should load"); + assert_eq!(closed_children, vec![second_child_thread_id]); + } + + #[tokio::test] + async fn local_store_updates_edge_status() { + let fixture = state_runtime().await; + let state_db = fixture.state_db; + let store = LocalAgentGraphStore::new(state_db); + let parent_thread_id = thread_id(/*suffix*/ 10); + let child_thread_id = thread_id(/*suffix*/ 11); + + store + .upsert_thread_spawn_edge( + parent_thread_id, + child_thread_id, + ThreadSpawnEdgeStatus::Open, + ) + .await + .expect("child edge should insert"); + store + .set_thread_spawn_edge_status(child_thread_id, ThreadSpawnEdgeStatus::Closed) + .await + .expect("child edge should close"); + + let open_children = store + .list_thread_spawn_children(parent_thread_id, Some(ThreadSpawnEdgeStatus::Open)) + .await + .expect("open children should load"); + assert_eq!(open_children, Vec::::new()); + + let closed_children = store + .list_thread_spawn_children(parent_thread_id, Some(ThreadSpawnEdgeStatus::Closed)) + .await + .expect("closed children should load"); + assert_eq!(closed_children, vec![child_thread_id]); + } + + #[tokio::test] + async fn local_store_lists_descendants_breadth_first_with_status_filters() { + let fixture = state_runtime().await; + let state_db = fixture.state_db; + let store = LocalAgentGraphStore::new(state_db.clone()); + let root_thread_id = thread_id(/*suffix*/ 20); + let later_child_thread_id = thread_id(/*suffix*/ 22); + let earlier_child_thread_id = thread_id(/*suffix*/ 21); + let closed_grandchild_thread_id = thread_id(/*suffix*/ 23); + let open_grandchild_thread_id = thread_id(/*suffix*/ 24); + let closed_child_thread_id = thread_id(/*suffix*/ 25); + let closed_great_grandchild_thread_id = thread_id(/*suffix*/ 26); + + for (parent_thread_id, child_thread_id, status) in [ + ( + root_thread_id, + later_child_thread_id, + ThreadSpawnEdgeStatus::Open, + ), + ( + root_thread_id, + earlier_child_thread_id, + ThreadSpawnEdgeStatus::Open, + ), + ( + earlier_child_thread_id, + open_grandchild_thread_id, + ThreadSpawnEdgeStatus::Open, + ), + ( + later_child_thread_id, + closed_grandchild_thread_id, + ThreadSpawnEdgeStatus::Closed, + ), + ( + root_thread_id, + closed_child_thread_id, + ThreadSpawnEdgeStatus::Closed, + ), + ( + closed_child_thread_id, + closed_great_grandchild_thread_id, + ThreadSpawnEdgeStatus::Closed, + ), + ] { + store + .upsert_thread_spawn_edge(parent_thread_id, child_thread_id, status) + .await + .expect("edge should insert"); + } + + let all_descendants = store + .list_thread_spawn_descendants(root_thread_id, /*status_filter*/ None) + .await + .expect("all descendants should load"); + assert_eq!( + all_descendants, + vec![ + earlier_child_thread_id, + later_child_thread_id, + closed_child_thread_id, + closed_grandchild_thread_id, + open_grandchild_thread_id, + closed_great_grandchild_thread_id, + ] + ); + + let open_descendants = store + .list_thread_spawn_descendants(root_thread_id, Some(ThreadSpawnEdgeStatus::Open)) + .await + .expect("open descendants should load"); + let state_open_descendants = state_db + .list_thread_spawn_descendants_with_status( + root_thread_id, + DirectionalThreadSpawnEdgeStatus::Open, + ) + .await + .expect("state open descendants should load"); + assert_eq!(open_descendants, state_open_descendants); + assert_eq!( + open_descendants, + vec![ + earlier_child_thread_id, + later_child_thread_id, + open_grandchild_thread_id, + ] + ); + + let closed_descendants = store + .list_thread_spawn_descendants(root_thread_id, Some(ThreadSpawnEdgeStatus::Closed)) + .await + .expect("closed descendants should load"); + assert_eq!( + closed_descendants, + vec![closed_child_thread_id, closed_great_grandchild_thread_id] + ); + } +} diff --git a/codex-rs/agent-graph-store/src/store.rs b/codex-rs/agent-graph-store/src/store.rs new file mode 100644 index 0000000000000000000000000000000000000000..0760cb15d4ac94cddc58c765d70961e7c46d931c --- /dev/null +++ b/codex-rs/agent-graph-store/src/store.rs @@ -0,0 +1,60 @@ +use std::future::Future; +use std::pin::Pin; + +use codex_protocol::ThreadId; + +use crate::AgentGraphStoreResult; +use crate::ThreadSpawnEdgeStatus; + +/// Future returned by [`AgentGraphStore`] operations. +pub type AgentGraphStoreFuture<'a, T> = + Pin> + Send + 'a>>; + +/// Storage-neutral boundary for persisted thread-spawn parent/child topology. +/// +/// Implementations are expected to return stable ordering for list methods so callers can merge +/// persisted graph state with live in-memory state without introducing nondeterministic output. +pub trait AgentGraphStore: Send + Sync { + /// Insert or replace the directional parent/child edge for a spawned thread. + /// + /// `child_thread_id` has at most one persisted parent. Re-inserting the same child should + /// update both the parent and status to match the supplied values. + fn upsert_thread_spawn_edge( + &self, + parent_thread_id: ThreadId, + child_thread_id: ThreadId, + status: ThreadSpawnEdgeStatus, + ) -> AgentGraphStoreFuture<'_, ()>; + + /// Update the persisted lifecycle status of a spawned thread's incoming edge. + /// + /// Implementations should treat missing children as a successful no-op. + fn set_thread_spawn_edge_status( + &self, + child_thread_id: ThreadId, + status: ThreadSpawnEdgeStatus, + ) -> AgentGraphStoreFuture<'_, ()>; + + /// List direct spawned children of a parent thread. + /// + /// When `status_filter` is `Some`, only child edges with that exact status are returned. When + /// it is `None`, all direct child edges are returned regardless of status, including statuses + /// that may be added by a future store implementation. + fn list_thread_spawn_children( + &self, + parent_thread_id: ThreadId, + status_filter: Option, + ) -> AgentGraphStoreFuture<'_, Vec>; + + /// List spawned descendants breadth-first by depth, then by thread id. + /// + /// `status_filter` is applied to every traversed edge, not just to the returned descendants. + /// For example, `Some(Open)` walks only open edges, so descendants under a closed edge are not + /// included even if their own incoming edge is open. `None` walks and returns every persisted + /// edge regardless of status. + fn list_thread_spawn_descendants( + &self, + root_thread_id: ThreadId, + status_filter: Option, + ) -> AgentGraphStoreFuture<'_, Vec>; +} diff --git a/codex-rs/agent-graph-store/src/types.rs b/codex-rs/agent-graph-store/src/types.rs new file mode 100644 index 0000000000000000000000000000000000000000..2a9f6caedb6c4e3ba3d07f5bce50c14085db3164 --- /dev/null +++ b/codex-rs/agent-graph-store/src/types.rs @@ -0,0 +1,42 @@ +use serde::Deserialize; +use serde::Serialize; + +/// Lifecycle status attached to a directional thread-spawn edge. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ThreadSpawnEdgeStatus { + /// The child thread is still live or resumable as an open spawned agent. + Open, + /// The child thread has been closed from the parent/child graph's perspective. + Closed, +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + + #[test] + fn thread_spawn_edge_status_serializes_as_snake_case() { + assert_eq!( + serde_json::to_string(&ThreadSpawnEdgeStatus::Open) + .expect("open status should serialize"), + "\"open\"" + ); + assert_eq!( + serde_json::to_string(&ThreadSpawnEdgeStatus::Closed) + .expect("closed status should serialize"), + "\"closed\"" + ); + assert_eq!( + serde_json::from_str::("\"open\"") + .expect("open status should deserialize"), + ThreadSpawnEdgeStatus::Open + ); + assert_eq!( + serde_json::from_str::("\"closed\"") + .expect("closed status should deserialize"), + ThreadSpawnEdgeStatus::Closed + ); + } +} diff --git a/codex-rs/apply-patch/src/file_update.rs b/codex-rs/apply-patch/src/file_update.rs new file mode 100644 index 0000000000000000000000000000000000000000..08bb0ebefe2197388562153db4888a89cefe2a53 --- /dev/null +++ b/codex-rs/apply-patch/src/file_update.rs @@ -0,0 +1,335 @@ +use codex_exec_server::ExecutorFileSystem; +use codex_exec_server::FileSystemSandboxContext; +use codex_exec_server::ReadFileOptions; +use codex_utils_path_uri::PathUri; +use similar::TextDiff; + +use crate::ApplyPatchError; +use crate::ApplyPatchFileUpdateMode; +use crate::IoError; +use crate::UpdateFileChunk; +use crate::seek_sequence; +use crate::text_file::Replacement; +use crate::text_file::SourceFile; + +#[cfg(test)] +#[path = "file_update_tests.rs"] +mod tests; + +pub(crate) struct AppliedPatch { + pub(crate) original_contents: String, + pub(crate) new_contents: String, +} + +/// Return *only* the new file contents (joined into a single `String`) after +/// applying the chunks to the file at `path`. +pub(crate) async fn derive_new_contents_from_chunks( + path: &PathUri, + chunks: &[UpdateFileChunk], + update_file_mode: ApplyPatchFileUpdateMode, + fs: &dyn ExecutorFileSystem, + follow_symlinks: bool, + sandbox: Option<&FileSystemSandboxContext>, +) -> std::result::Result { + let original_contents = fs + .read_file_text(path, ReadFileOptions { follow_symlinks }, sandbox) + .await + .map_err(|err| { + ApplyPatchError::IoError(IoError { + context: format!( + "Failed to read file to update {}", + path.inferred_native_path_string() + ), + source: err, + }) + })?; + + let path_text = path.inferred_native_path_string(); + let new_contents = match update_file_mode { + ApplyPatchFileUpdateMode::NormalizeToLf => { + let mut original_lines = original_contents + .split('\n') + .map(String::from) + .collect::>(); + + // Drop the trailing empty element that results from the final newline so + // that line counts match the behaviour of standard `diff`. + if original_lines.last().is_some_and(String::is_empty) { + original_lines.pop(); + } + + let replacements = + compute_replacements(&original_lines, &path_text, chunks, update_file_mode)?; + let mut new_lines = apply_replacements(original_lines, &replacements); + if !new_lines.last().is_some_and(String::is_empty) { + new_lines.push(String::new()); + } + new_lines.join("\n") + } + ApplyPatchFileUpdateMode::PreserveLineEndings => { + let mut source_file = SourceFile::parse(&original_contents); + let original_lines = source_file.line_texts(); + let replacements = + compute_replacements(&original_lines, &path_text, chunks, update_file_mode)?; + source_file.apply_replacements(&replacements); + source_file.into_contents() + } + }; + Ok(AppliedPatch { + original_contents, + new_contents, + }) +} + +/// Compute a list of replacements needed to transform `original_lines` into the +/// new lines, given the patch `chunks`. Each replacement is returned as +/// `(start_index, old_len, new_lines)`. +fn compute_replacements( + original_lines: &[String], + path: &str, + chunks: &[UpdateFileChunk], + update_file_mode: ApplyPatchFileUpdateMode, +) -> std::result::Result, ApplyPatchError> { + let mut replacements: Vec = Vec::new(); + let mut line_index: usize = 0; + + for chunk in chunks { + // If a chunk has a `change_context`, we use seek_sequence to find it, then + // adjust our `line_index` to continue from there. + if let Some(ctx_line) = &chunk.change_context { + if let Some(idx) = seek_sequence::seek_sequence( + original_lines, + std::slice::from_ref(ctx_line), + line_index, + /*eof*/ false, + update_file_mode, + ) { + line_index = idx + 1; + } else { + return Err(ApplyPatchError::ComputeReplacements(format!( + "Failed to find context '{ctx_line}' in {path}" + ))); + } + } + + if chunk.old_lines.is_empty() { + // Preserve the legacy split representation's handling of a final + // empty line. `SourceFile` only exposes real source lines, so its + // insertion point is always after the final line. + let insertion_idx = match update_file_mode { + ApplyPatchFileUpdateMode::NormalizeToLf => { + if original_lines.last().is_some_and(String::is_empty) { + original_lines.len() - 1 + } else { + original_lines.len() + } + } + ApplyPatchFileUpdateMode::PreserveLineEndings => original_lines.len(), + }; + replacements.push((insertion_idx, 0, chunk.new_lines.clone())); + continue; + } + + // Otherwise, try to match the existing lines in the file with the old lines + // from the chunk. If found, schedule that region for replacement. + // Attempt to locate the `old_lines` verbatim within the file. In many + // real‑world diffs the last element of `old_lines` is an *empty* string + // representing the terminating newline of the region being replaced. + // This sentinel is not present in `original_lines` because `SourceFile` + // stores the terminator on the preceding line rather than as an extra + // trailing element. If a direct search fails and the pattern ends with + // an empty string, retry without that final element so modifications + // touching the end‑of‑file can be located reliably. + + let mut pattern: &[String] = &chunk.old_lines; + let mut found = seek_sequence::seek_sequence( + original_lines, + pattern, + line_index, + chunk.is_end_of_file, + update_file_mode, + ); + + let mut new_slice: &[String] = &chunk.new_lines; + + if found.is_none() && pattern.last().is_some_and(String::is_empty) { + // Retry without the trailing empty line which represents the final + // newline in the file. + pattern = &pattern[..pattern.len() - 1]; + if new_slice.last().is_some_and(String::is_empty) { + new_slice = &new_slice[..new_slice.len() - 1]; + } + + found = seek_sequence::seek_sequence( + original_lines, + pattern, + line_index, + chunk.is_end_of_file, + update_file_mode, + ); + } + + if let Some(start_idx) = found { + match update_file_mode { + ApplyPatchFileUpdateMode::NormalizeToLf => { + replacements.push((start_idx, pattern.len(), new_slice.to_vec())); + } + ApplyPatchFileUpdateMode::PreserveLineEndings => { + // Context lines occur in both sides of a patch chunk. Keep those + // original lines in place so their exact contents and terminators + // survive, especially when the file has mixed line endings. + let mut old_start = 0; + let mut new_start = 0; + for &(old_context, new_context) in &chunk.context_line_indices { + // A trailing empty context line can be removed from `pattern` + // and `new_slice` above when it represents the final newline. + if old_context >= pattern.len() || new_context >= new_slice.len() { + break; + } + if old_start != old_context || new_start != new_context { + replacements.push(( + start_idx + old_start, + old_context - old_start, + new_slice[new_start..new_context].to_vec(), + )); + } + old_start = old_context + 1; + new_start = new_context + 1; + } + if old_start != pattern.len() || new_start != new_slice.len() { + replacements.push(( + start_idx + old_start, + pattern.len() - old_start, + new_slice[new_start..].to_vec(), + )); + } + } + } + line_index = start_idx + pattern.len(); + } else { + return Err(ApplyPatchError::ComputeReplacements(format!( + "Failed to find expected lines in {}:\n{}", + path, + chunk.old_lines.join("\n"), + ))); + } + } + + replacements.sort_by_key(|(index, _, _)| *index); + + Ok(replacements) +} + +/// Apply the `(start_index, old_len, new_lines)` replacements to `original_lines`, +/// returning the modified file contents as a vector of lines. +fn apply_replacements(mut lines: Vec, replacements: &[Replacement]) -> Vec { + // We must apply replacements in descending order so that earlier replacements + // don't shift the positions of later ones. + for (start_idx, old_len, new_segment) in replacements.iter().rev() { + let start_idx = *start_idx; + let old_len = *old_len; + + // Remove old lines. + for _ in 0..old_len { + if start_idx < lines.len() { + lines.remove(start_idx); + } + } + + // Insert new lines. + for (offset, new_line) in new_segment.iter().enumerate() { + lines.insert(start_idx + offset, new_line.clone()); + } + } + + lines +} + +/// Intended result of a file update for apply_patch. +#[derive(Debug, Eq, PartialEq)] +pub struct ApplyPatchFileUpdate { + pub(crate) unified_diff: String, + pub(crate) original_content: String, + pub(crate) content: String, +} + +pub async fn unified_diff_from_chunks( + path: &PathUri, + chunks: &[UpdateFileChunk], + fs: &dyn ExecutorFileSystem, + sandbox: Option<&FileSystemSandboxContext>, +) -> std::result::Result { + unified_diff_from_chunks_with_mode( + path, + chunks, + ApplyPatchFileUpdateMode::default(), + fs, + sandbox, + ) + .await +} + +pub(crate) async fn unified_diff_from_chunks_with_mode( + path: &PathUri, + chunks: &[UpdateFileChunk], + update_file_mode: ApplyPatchFileUpdateMode, + fs: &dyn ExecutorFileSystem, + sandbox: Option<&FileSystemSandboxContext>, +) -> std::result::Result { + unified_diff_from_chunks_with_context_and_mode( + path, + chunks, + /*context*/ 1, + update_file_mode, + fs, + sandbox, + ) + .await +} + +pub async fn unified_diff_from_chunks_with_context( + path: &PathUri, + chunks: &[UpdateFileChunk], + context: usize, + fs: &dyn ExecutorFileSystem, + sandbox: Option<&FileSystemSandboxContext>, +) -> std::result::Result { + unified_diff_from_chunks_with_context_and_mode( + path, + chunks, + context, + ApplyPatchFileUpdateMode::default(), + fs, + sandbox, + ) + .await +} + +async fn unified_diff_from_chunks_with_context_and_mode( + path: &PathUri, + chunks: &[UpdateFileChunk], + context: usize, + update_file_mode: ApplyPatchFileUpdateMode, + fs: &dyn ExecutorFileSystem, + sandbox: Option<&FileSystemSandboxContext>, +) -> std::result::Result { + let AppliedPatch { + original_contents, + new_contents, + } = derive_new_contents_from_chunks( + path, + chunks, + update_file_mode, + fs, + /*follow_symlinks*/ true, + sandbox, + ) + .await?; + let text_diff = TextDiff::from_lines(&original_contents, &new_contents); + let unified_diff = text_diff.unified_diff().context_radius(context).to_string(); + Ok(ApplyPatchFileUpdate { + unified_diff, + original_content: original_contents, + content: new_contents, + }) +} diff --git a/codex-rs/apply-patch/src/file_update_tests.rs b/codex-rs/apply-patch/src/file_update_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..968082e26d64acb9ee4a9feea468763971cb2d4f --- /dev/null +++ b/codex-rs/apply-patch/src/file_update_tests.rs @@ -0,0 +1,278 @@ +use super::*; +use crate::Hunk; +use crate::apply_patch; +use crate::parse_patch; +use codex_exec_server::LOCAL_FS; +use codex_utils_path_uri::PathUri; +use pretty_assertions::assert_eq; +use std::fs; +use tempfile::tempdir; + +fn wrap_patch(body: &str) -> String { + format!("*** Begin Patch\n{body}\n*** End Patch") +} + +#[tokio::test] +async fn test_unified_diff() { + // Start with a file containing four lines. + let dir = tempdir().unwrap(); + let path = dir.path().join("multi.txt"); + fs::write(&path, "foo\nbar\nbaz\nqux\n").unwrap(); + let patch = wrap_patch(&format!( + r#"*** Update File: {} +@@ + foo +-bar ++BAR +@@ + baz +-qux ++QUX"#, + path.display() + )); + let patch = parse_patch(&patch).unwrap(); + + let update_file_chunks = match patch.hunks.as_slice() { + [Hunk::UpdateFile { chunks, .. }] => chunks, + _ => panic!("Expected a single UpdateFile hunk"), + }; + let path_uri = PathUri::from_host_native_path(&path).expect("absolute test path"); + let diff = unified_diff_from_chunks( + &path_uri, + update_file_chunks, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .unwrap(); + let expected_diff = r#"@@ -1,4 +1,4 @@ + foo +-bar ++BAR + baz +-qux ++QUX +"#; + let expected = ApplyPatchFileUpdate { + unified_diff: expected_diff.to_string(), + original_content: "foo\nbar\nbaz\nqux\n".to_string(), + content: "foo\nBAR\nbaz\nQUX\n".to_string(), + }; + assert_eq!(expected, diff); +} + +#[tokio::test] +async fn test_unified_diff_first_line_replacement() { + // Replace the very first line of the file. + let dir = tempdir().unwrap(); + let path = dir.path().join("first.txt"); + fs::write(&path, "foo\nbar\nbaz\n").unwrap(); + + let patch = wrap_patch(&format!( + r#"*** Update File: {} +@@ +-foo ++FOO + bar +"#, + path.display() + )); + + let patch = parse_patch(&patch).unwrap(); + let chunks = match patch.hunks.as_slice() { + [Hunk::UpdateFile { chunks, .. }] => chunks, + _ => panic!("Expected a single UpdateFile hunk"), + }; + + let resolved_path = PathUri::from_host_native_path(&path).expect("absolute test path"); + let diff = unified_diff_from_chunks( + &resolved_path, + chunks, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .unwrap(); + let expected_diff = r#"@@ -1,2 +1,2 @@ +-foo ++FOO + bar +"#; + let expected = ApplyPatchFileUpdate { + unified_diff: expected_diff.to_string(), + original_content: "foo\nbar\nbaz\n".to_string(), + content: "FOO\nbar\nbaz\n".to_string(), + }; + assert_eq!(expected, diff); +} + +#[tokio::test] +async fn test_unified_diff_last_line_replacement() { + // Replace the very last line of the file. + let dir = tempdir().unwrap(); + let path = dir.path().join("last.txt"); + fs::write(&path, "foo\nbar\nbaz\n").unwrap(); + + let patch = wrap_patch(&format!( + r#"*** Update File: {} +@@ + foo + bar +-baz ++BAZ +"#, + path.display() + )); + + let patch = parse_patch(&patch).unwrap(); + let chunks = match patch.hunks.as_slice() { + [Hunk::UpdateFile { chunks, .. }] => chunks, + _ => panic!("Expected a single UpdateFile hunk"), + }; + + let resolved_path = PathUri::from_host_native_path(&path).expect("absolute test path"); + let diff = unified_diff_from_chunks( + &resolved_path, + chunks, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .unwrap(); + let expected_diff = r#"@@ -2,2 +2,2 @@ + bar +-baz ++BAZ +"#; + let expected = ApplyPatchFileUpdate { + unified_diff: expected_diff.to_string(), + original_content: "foo\nbar\nbaz\n".to_string(), + content: "foo\nbar\nBAZ\n".to_string(), + }; + assert_eq!(expected, diff); +} + +#[tokio::test] +async fn test_unified_diff_insert_at_eof() { + // Insert a new line at end-of-file. + let dir = tempdir().unwrap(); + let path = dir.path().join("insert.txt"); + fs::write(&path, "foo\nbar\nbaz\n").unwrap(); + + let patch = wrap_patch(&format!( + r#"*** Update File: {} +@@ ++quux +*** End of File +"#, + path.display() + )); + + let patch = parse_patch(&patch).unwrap(); + let chunks = match patch.hunks.as_slice() { + [Hunk::UpdateFile { chunks, .. }] => chunks, + _ => panic!("Expected a single UpdateFile hunk"), + }; + + let path_uri = PathUri::from_host_native_path(&path).expect("absolute test path"); + let diff = + unified_diff_from_chunks(&path_uri, chunks, LOCAL_FS.as_ref(), /*sandbox*/ None) + .await + .unwrap(); + let expected_diff = r#"@@ -3 +3,2 @@ + baz ++quux +"#; + let expected = ApplyPatchFileUpdate { + unified_diff: expected_diff.to_string(), + original_content: "foo\nbar\nbaz\n".to_string(), + content: "foo\nbar\nbaz\nquux\n".to_string(), + }; + assert_eq!(expected, diff); +} + +#[tokio::test] +async fn test_unified_diff_interleaved_changes() { + // Original file with six lines. + let dir = tempdir().unwrap(); + let path = dir.path().join("interleaved.txt"); + fs::write(&path, "a\nb\nc\nd\ne\nf\n").unwrap(); + + // Patch replaces two separate lines and appends a new one at EOF using + // three distinct chunks. + let patch_body = format!( + r#"*** Update File: {} +@@ + a +-b ++B +@@ + d +-e ++E +@@ + f ++g +*** End of File"#, + path.display() + ); + let patch = wrap_patch(&patch_body); + + // Extract chunks then build the unified diff. + let parsed = parse_patch(&patch).unwrap(); + let chunks = match parsed.hunks.as_slice() { + [Hunk::UpdateFile { chunks, .. }] => chunks, + _ => panic!("Expected a single UpdateFile hunk"), + }; + + let path_uri = PathUri::from_host_native_path(&path).expect("absolute test path"); + let diff = + unified_diff_from_chunks(&path_uri, chunks, LOCAL_FS.as_ref(), /*sandbox*/ None) + .await + .unwrap(); + + let expected_diff = r#"@@ -1,6 +1,7 @@ + a +-b ++B + c + d +-e ++E + f ++g +"#; + + let expected = ApplyPatchFileUpdate { + unified_diff: expected_diff.to_string(), + original_content: "a\nb\nc\nd\ne\nf\n".to_string(), + content: "a\nB\nc\nd\nE\nf\ng\n".to_string(), + }; + + assert_eq!(expected, diff); + + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + apply_patch( + &patch, + &PathUri::from_host_native_path(dir.path()).expect("absolute test path"), + &mut stdout, + &mut stderr, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .unwrap(); + let contents = fs::read_to_string(path).unwrap(); + assert_eq!( + contents, + r#"a +B +c +d +E +f +g +"# + ); +} diff --git a/codex-rs/apply-patch/src/invocation.rs b/codex-rs/apply-patch/src/invocation.rs new file mode 100644 index 0000000000000000000000000000000000000000..5b805c30c69e208725fd069a51ab31236629af9b --- /dev/null +++ b/codex-rs/apply-patch/src/invocation.rs @@ -0,0 +1,1036 @@ +use std::collections::HashMap; +use std::sync::LazyLock; + +use codex_exec_server::ExecutorFileSystem; +use tree_sitter::Parser; +use tree_sitter::Query; +use tree_sitter::QueryCursor; +use tree_sitter::StreamingIterator; +use tree_sitter_bash::LANGUAGE as BASH; + +use crate::ApplyPatchAction; +use crate::ApplyPatchArgs; +use crate::ApplyPatchError; +use crate::ApplyPatchFileChange; +use crate::ApplyPatchFileUpdate; +use crate::ApplyPatchFileUpdateMode; +use crate::IoError; +use crate::MaybeApplyPatchVerified; +use crate::parser::Hunk; +use crate::parser::ParseError; +use crate::parser::parse_patch; +use crate::unified_diff_from_chunks_with_mode; +use codex_utils_path_uri::PathConvention; +use codex_utils_path_uri::PathUri; +use std::str::Utf8Error; +use tree_sitter::LanguageError; + +const APPLY_PATCH_COMMANDS: [&str; 2] = ["apply_patch", "applypatch"]; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ApplyPatchShell { + Unix, + PowerShell, + Cmd, +} + +#[derive(Debug, PartialEq)] +pub enum MaybeApplyPatch { + Body(ApplyPatchArgs), + ShellParseError(ExtractHeredocError), + PatchParseError(ParseError), + NotApplyPatch, +} + +#[derive(Debug, PartialEq)] +pub enum ExtractHeredocError { + CommandDidNotStartWithApplyPatch, + FailedToLoadBashGrammar(LanguageError), + HeredocNotUtf8(Utf8Error), + FailedToParsePatchIntoAst, + FailedToFindHeredocBody, +} + +fn classify_shell_name(shell: &str, convention: PathConvention) -> Option { + let basename = convention.path_segments(shell).next_back()?; + let stem = basename + .rsplit_once('.') + .and_then(|(stem, _extension)| (!stem.is_empty()).then_some(stem)) + .unwrap_or(basename); + Some(stem.to_ascii_lowercase()) +} + +fn classify_shell(shell: &str, flag: &str, convention: PathConvention) -> Option { + classify_shell_name(shell, convention).and_then(|name| match name.as_str() { + "bash" | "zsh" | "sh" if matches!(flag, "-lc" | "-c") => Some(ApplyPatchShell::Unix), + "pwsh" | "powershell" if flag.eq_ignore_ascii_case("-command") => { + Some(ApplyPatchShell::PowerShell) + } + "cmd" if flag.eq_ignore_ascii_case("/c") => Some(ApplyPatchShell::Cmd), + _ => None, + }) +} + +fn can_skip_flag(shell: &str, flag: &str, convention: PathConvention) -> bool { + classify_shell_name(shell, convention).is_some_and(|name| { + matches!(name.as_str(), "pwsh" | "powershell") && flag.eq_ignore_ascii_case("-noprofile") + }) +} + +fn parse_shell_script<'a>(argv: &'a [String], cwd: &PathUri) -> Option<(ApplyPatchShell, &'a str)> { + let convention = cwd.infer_path_convention()?; + match argv { + [shell, flag, script] => classify_shell(shell, flag, convention).map(|shell_type| { + let script = script.as_str(); + (shell_type, script) + }), + [shell, skip_flag, flag, script] => { + if !can_skip_flag(shell, skip_flag, convention) { + return None; + } + classify_shell(shell, flag, convention).map(|shell_type| { + let script = script.as_str(); + (shell_type, script) + }) + } + _ => None, + } +} + +fn extract_apply_patch_from_shell( + shell: ApplyPatchShell, + script: &str, +) -> std::result::Result<(String, Option), ExtractHeredocError> { + match shell { + ApplyPatchShell::Unix | ApplyPatchShell::PowerShell | ApplyPatchShell::Cmd => { + extract_apply_patch_from_bash(script) + } + } +} + +// TODO: make private once we remove tests in lib.rs +/// `cwd` supplies the path convention used to interpret the shell executable in `argv`. +pub fn maybe_parse_apply_patch(argv: &[String], cwd: &PathUri) -> MaybeApplyPatch { + match argv { + // Direct invocation: apply_patch + [cmd, body] if APPLY_PATCH_COMMANDS.contains(&cmd.as_str()) => match parse_patch(body) { + Ok(source) => MaybeApplyPatch::Body(source), + Err(e) => MaybeApplyPatch::PatchParseError(e), + }, + // Shell heredoc form: (optional `cd &&`) apply_patch <<'EOF' ... + _ => match parse_shell_script(argv, cwd) { + Some((shell, script)) => match extract_apply_patch_from_shell(shell, script) { + Ok((body, workdir)) => match parse_patch(&body) { + Ok(mut source) => { + source.workdir = workdir; + MaybeApplyPatch::Body(source) + } + Err(e) => MaybeApplyPatch::PatchParseError(e), + }, + Err(ExtractHeredocError::CommandDidNotStartWithApplyPatch) => { + MaybeApplyPatch::NotApplyPatch + } + Err(e) => MaybeApplyPatch::ShellParseError(e), + }, + None => MaybeApplyPatch::NotApplyPatch, + }, + } +} + +/// `cwd` must identify an absolute environment-native path so relative patch paths can be +/// resolved without projecting them onto the app-server or exec-server host. +pub async fn maybe_parse_apply_patch_verified( + argv: &[String], + cwd: &PathUri, + fs: &dyn ExecutorFileSystem, + sandbox: Option<&codex_exec_server::FileSystemSandboxContext>, +) -> MaybeApplyPatchVerified { + maybe_parse_apply_patch_verified_with_mode( + argv, + cwd, + ApplyPatchFileUpdateMode::default(), + fs, + sandbox, + ) + .await +} + +/// Parses and verifies an `apply_patch` invocation using the selected +/// file-update mode. +pub async fn maybe_parse_apply_patch_verified_with_mode( + argv: &[String], + cwd: &PathUri, + update_file_mode: ApplyPatchFileUpdateMode, + fs: &dyn ExecutorFileSystem, + sandbox: Option<&codex_exec_server::FileSystemSandboxContext>, +) -> MaybeApplyPatchVerified { + // Detect a raw patch body passed directly as the command or as the body of a shell + // script. In these cases, report an explicit error rather than applying the patch. + if let [body] = argv + && parse_patch(body).is_ok() + { + return MaybeApplyPatchVerified::CorrectnessError(ApplyPatchError::ImplicitInvocation); + } + if let Some((_, script)) = parse_shell_script(argv, cwd) + && parse_patch(script).is_ok() + { + return MaybeApplyPatchVerified::CorrectnessError(ApplyPatchError::ImplicitInvocation); + } + + match maybe_parse_apply_patch(argv, cwd) { + MaybeApplyPatch::Body(args) => { + verify_apply_patch_args_with_mode(args, cwd, update_file_mode, fs, sandbox).await + } + MaybeApplyPatch::ShellParseError(e) => MaybeApplyPatchVerified::ShellParseError(e), + MaybeApplyPatch::PatchParseError(e) => MaybeApplyPatchVerified::CorrectnessError(e.into()), + MaybeApplyPatch::NotApplyPatch => MaybeApplyPatchVerified::NotApplyPatch, + } +} + +pub async fn verify_apply_patch_args( + args: ApplyPatchArgs, + cwd: &PathUri, + fs: &dyn ExecutorFileSystem, + sandbox: Option<&codex_exec_server::FileSystemSandboxContext>, +) -> MaybeApplyPatchVerified { + verify_apply_patch_args_with_mode(args, cwd, ApplyPatchFileUpdateMode::default(), fs, sandbox) + .await +} + +/// Verifies parsed patch arguments using the selected file-update mode. +pub async fn verify_apply_patch_args_with_mode( + args: ApplyPatchArgs, + cwd: &PathUri, + update_file_mode: ApplyPatchFileUpdateMode, + fs: &dyn ExecutorFileSystem, + sandbox: Option<&codex_exec_server::FileSystemSandboxContext>, +) -> MaybeApplyPatchVerified { + match try_verify_apply_patch_args(args, cwd, update_file_mode, fs, sandbox).await { + Ok(action) => MaybeApplyPatchVerified::Body(action), + Err(err) => MaybeApplyPatchVerified::CorrectnessError(err), + } +} + +async fn try_verify_apply_patch_args( + args: ApplyPatchArgs, + cwd: &PathUri, + update_file_mode: ApplyPatchFileUpdateMode, + fs: &dyn ExecutorFileSystem, + sandbox: Option<&codex_exec_server::FileSystemSandboxContext>, +) -> Result { + let ApplyPatchArgs { + patch, + hunks, + workdir, + .. + } = args; + let effective_cwd = workdir + .as_ref() + .map(|dir| cwd.join(dir)) + .transpose()? + .unwrap_or_else(|| cwd.clone()); + let mut changes = HashMap::new(); + for hunk in hunks { + let path = hunk.resolve_path(&effective_cwd)?; + if changes.contains_key(&path) { + return Err(ParseError::InvalidPatchError(format!( + "multiple operations target {}", + path.inferred_native_path_string() + )) + .into()); + } + match hunk { + Hunk::AddFile { contents, .. } => { + changes.insert(path, ApplyPatchFileChange::Add { content: contents }); + } + Hunk::DeleteFile { .. } => { + let content = fs + .read_file_text(&path, Default::default(), sandbox) + .await + .map_err(|source| { + ApplyPatchError::IoError(IoError { + context: format!( + "Failed to read {}", + path.inferred_native_path_string() + ), + source, + }) + })?; + changes.insert(path, ApplyPatchFileChange::Delete { content }); + } + Hunk::UpdateFile { + move_path, chunks, .. + } => { + let ApplyPatchFileUpdate { + unified_diff, + content: contents, + .. + } = unified_diff_from_chunks_with_mode( + &path, + &chunks, + update_file_mode, + fs, + sandbox, + ) + .await?; + changes.insert( + path, + ApplyPatchFileChange::Update { + unified_diff, + move_path: move_path + .map(|path| effective_cwd.join(&path.to_string_lossy())) + .transpose()?, + new_content: contents, + }, + ); + } + } + } + Ok(ApplyPatchAction { + changes, + update_file_mode, + patch, + cwd: effective_cwd, + }) +} + +/// Extract the heredoc body (and optional `cd` workdir) from a `bash -lc` script +/// that invokes the apply_patch tool using a heredoc. +/// +/// Supported top‑level forms (must be the only top‑level statement): +/// - `apply_patch <<'EOF'\n...\nEOF` +/// - `cd && apply_patch <<'EOF'\n...\nEOF` +/// +/// Notes about matching: +/// - Parsed with Tree‑sitter Bash and a strict query that uses anchors so the +/// heredoc‑redirected statement is the only top‑level statement. +/// - The connector between `cd` and `apply_patch` must be `&&` (not `|` or `||`). +/// - Exactly one positional `word` argument is allowed for `cd` (no flags, no quoted +/// strings, no second argument). +/// - The apply command is validated in‑query via `#any-of?` to allow `apply_patch` +/// or `applypatch`. +/// - Preceding or trailing commands (e.g., `echo ...;` or `... && echo done`) do not match. +/// +/// Returns `(heredoc_body, Some(path))` when the `cd` variant matches, or +/// `(heredoc_body, None)` for the direct form. Errors are returned if the script +/// cannot be parsed or does not match the allowed patterns. +fn extract_apply_patch_from_bash( + src: &str, +) -> std::result::Result<(String, Option), ExtractHeredocError> { + // This function uses a Tree-sitter query to recognize one of two + // whole-script forms, each expressed as a single top-level statement: + // + // 1. apply_patch <<'EOF'\n...\nEOF + // 2. cd && apply_patch <<'EOF'\n...\nEOF + // + // Key ideas when reading the query: + // - dots (`.`) between named nodes enforces adjacency among named children and + // anchor to the start/end of the expression. + // - we match a single redirected_statement directly under program with leading + // and trailing anchors (`.`). This ensures it is the only top-level statement + // (so prefixes like `echo ...;` or suffixes like `... && echo done` do not match). + // + // Overall, we want to be conservative and only match the intended forms, as other + // forms are likely to be model errors, or incorrectly interpreted by later code. + // + // If you're editing this query, it's helpful to start by creating a debugging binary + // which will let you see the AST of an arbitrary bash script passed in, and optionally + // also run an arbitrary query against the AST. This is useful for understanding + // how tree-sitter parses the script and whether the query syntax is correct. Be sure + // to test both positive and negative cases. + static APPLY_PATCH_QUERY: LazyLock = LazyLock::new(|| { + let language = BASH.into(); + #[expect(clippy::expect_used)] + Query::new( + &language, + r#" + ( + program + . (redirected_statement + body: (command + name: (command_name (word) @apply_name) .) + (#any-of? @apply_name "apply_patch" "applypatch") + redirect: (heredoc_redirect + . (heredoc_start) + . (heredoc_body) @heredoc + . (heredoc_end) + .)) + .) + + ( + program + . (redirected_statement + body: (list + . (command + name: (command_name (word) @cd_name) . + argument: [ + (word) @cd_path + (string (string_content) @cd_path) + (raw_string) @cd_raw_string + ] .) + "&&" + . (command + name: (command_name (word) @apply_name)) + .) + (#eq? @cd_name "cd") + (#any-of? @apply_name "apply_patch" "applypatch") + redirect: (heredoc_redirect + . (heredoc_start) + . (heredoc_body) @heredoc + . (heredoc_end) + .)) + .) + "#, + ) + .expect("valid bash query") + }); + + let lang = BASH.into(); + let mut parser = Parser::new(); + parser + .set_language(&lang) + .map_err(ExtractHeredocError::FailedToLoadBashGrammar)?; + let tree = parser + .parse(src, None) + .ok_or(ExtractHeredocError::FailedToParsePatchIntoAst)?; + + let bytes = src.as_bytes(); + let root = tree.root_node(); + + let mut cursor = QueryCursor::new(); + let mut matches = cursor.matches(&APPLY_PATCH_QUERY, root, bytes); + while let Some(m) = matches.next() { + let mut heredoc_text: Option = None; + let mut cd_path: Option = None; + + for capture in m.captures.iter() { + let name = APPLY_PATCH_QUERY.capture_names()[capture.index as usize]; + match name { + "heredoc" => { + let text = capture + .node + .utf8_text(bytes) + .map_err(ExtractHeredocError::HeredocNotUtf8)? + .trim_end_matches('\n') + .to_string(); + heredoc_text = Some(text); + } + "cd_path" => { + let text = capture + .node + .utf8_text(bytes) + .map_err(ExtractHeredocError::HeredocNotUtf8)? + .to_string(); + cd_path = Some(text); + } + "cd_raw_string" => { + let raw = capture + .node + .utf8_text(bytes) + .map_err(ExtractHeredocError::HeredocNotUtf8)?; + let trimmed = raw + .strip_prefix('\'') + .and_then(|s| s.strip_suffix('\'')) + .unwrap_or(raw); + cd_path = Some(trimmed.to_string()); + } + _ => {} + } + } + + if let Some(heredoc) = heredoc_text { + return Ok((heredoc, cd_path)); + } + } + + Err(ExtractHeredocError::CommandDidNotStartWithApplyPatch) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::unified_diff_from_chunks; + use assert_matches::assert_matches; + use codex_exec_server::LOCAL_FS; + use pretty_assertions::assert_eq; + use std::fs; + use std::path::PathBuf; + use std::string::ToString; + use tempfile::tempdir; + + /// Helper to construct a patch with the given body. + fn wrap_patch(body: &str) -> String { + format!("*** Begin Patch\n{body}\n*** End Patch") + } + + fn strs_to_strings(strs: &[&str]) -> Vec { + strs.iter().map(ToString::to_string).collect() + } + + // Test helpers to reduce repetition when building bash -lc heredoc scripts + fn args_bash(script: &str) -> Vec { + strs_to_strings(&["bash", "-lc", script]) + } + + fn args_powershell(script: &str) -> Vec { + strs_to_strings(&["powershell.exe", "-Command", script]) + } + + fn args_powershell_no_profile(script: &str) -> Vec { + strs_to_strings(&["powershell.exe", "-NoProfile", "-Command", script]) + } + + fn args_pwsh(script: &str) -> Vec { + strs_to_strings(&["pwsh", "-NoProfile", "-Command", script]) + } + + fn args_cmd(script: &str) -> Vec { + strs_to_strings(&["cmd.exe", "/c", script]) + } + + fn heredoc_script(prefix: &str) -> String { + format!( + "{prefix}apply_patch <<'PATCH'\n*** Begin Patch\n*** Add File: foo\n+hi\n*** End Patch\nPATCH" + ) + } + + fn heredoc_script_ps(prefix: &str, suffix: &str) -> String { + format!( + "{prefix}apply_patch <<'PATCH'\n*** Begin Patch\n*** Add File: foo\n+hi\n*** End Patch\nPATCH{suffix}" + ) + } + + fn expected_single_add() -> Vec { + vec![Hunk::AddFile { + path: PathBuf::from("foo"), + contents: "hi\n".to_string(), + }] + } + + #[track_caller] + fn assert_match_args(args: Vec, expected_workdir: Option<&str>) { + assert_match_args_with_cwd( + args, + &PathUri::parse("file:///workspace").expect("valid POSIX test cwd"), + expected_workdir, + ); + } + + #[track_caller] + fn assert_match_args_with_cwd( + args: Vec, + cwd: &PathUri, + expected_workdir: Option<&str>, + ) { + match maybe_parse_apply_patch(&args, cwd) { + MaybeApplyPatch::Body(ApplyPatchArgs { hunks, workdir, .. }) => { + assert_eq!(workdir.as_deref(), expected_workdir); + assert_eq!(hunks, expected_single_add()); + } + result => panic!("expected MaybeApplyPatch::Body got {result:?}"), + } + } + + #[track_caller] + fn assert_match(script: &str, expected_workdir: Option<&str>) { + let args = args_bash(script); + assert_match_args(args, expected_workdir); + } + + fn assert_not_match(script: &str) { + let args = args_bash(script); + assert_matches!( + maybe_parse_apply_patch( + &args, + &PathUri::parse("file:///workspace").expect("valid POSIX test cwd"), + ), + MaybeApplyPatch::NotApplyPatch + ); + } + + #[tokio::test] + async fn test_implicit_patch_single_arg_is_error() { + let patch = "*** Begin Patch\n*** Add File: foo\n+hi\n*** End Patch".to_string(); + let args = vec![patch]; + let dir = tempdir().unwrap(); + assert_matches!( + maybe_parse_apply_patch_verified( + &args, + &PathUri::from_host_native_path(dir.path()).expect("absolute test path"), + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await, + MaybeApplyPatchVerified::CorrectnessError(ApplyPatchError::ImplicitInvocation) + ); + } + + #[tokio::test] + async fn test_implicit_patch_bash_script_is_error() { + let script = "*** Begin Patch\n*** Add File: foo\n+hi\n*** End Patch"; + let args = args_bash(script); + let dir = tempdir().unwrap(); + assert_matches!( + maybe_parse_apply_patch_verified( + &args, + &PathUri::from_host_native_path(dir.path()).expect("absolute test path"), + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await, + MaybeApplyPatchVerified::CorrectnessError(ApplyPatchError::ImplicitInvocation) + ); + } + + #[tokio::test] + async fn test_literal() { + let args = strs_to_strings(&[ + "apply_patch", + r#"*** Begin Patch +*** Add File: foo ++hi +*** End Patch +"#, + ]); + + match maybe_parse_apply_patch( + &args, + &PathUri::parse("file:///workspace").expect("valid POSIX test cwd"), + ) { + MaybeApplyPatch::Body(ApplyPatchArgs { hunks, .. }) => { + assert_eq!( + hunks, + vec![Hunk::AddFile { + path: PathBuf::from("foo"), + contents: "hi\n".to_string() + }] + ); + } + result => panic!("expected MaybeApplyPatch::Body got {result:?}"), + } + } + + #[tokio::test] + async fn test_literal_applypatch() { + let args = strs_to_strings(&[ + "applypatch", + r#"*** Begin Patch +*** Add File: foo ++hi +*** End Patch +"#, + ]); + + match maybe_parse_apply_patch( + &args, + &PathUri::parse("file:///workspace").expect("valid POSIX test cwd"), + ) { + MaybeApplyPatch::Body(ApplyPatchArgs { hunks, .. }) => { + assert_eq!( + hunks, + vec![Hunk::AddFile { + path: PathBuf::from("foo"), + contents: "hi\n".to_string() + }] + ); + } + result => panic!("expected MaybeApplyPatch::Body got {result:?}"), + } + } + + #[tokio::test] + async fn test_heredoc() { + assert_match(&heredoc_script(""), /*expected_workdir*/ None); + } + + #[tokio::test] + async fn test_heredoc_non_login_shell() { + let script = heredoc_script(""); + let args = strs_to_strings(&["bash", "-c", &script]); + assert_match_args(args, /*expected_workdir*/ None); + } + + #[tokio::test] + async fn test_heredoc_applypatch() { + let args = strs_to_strings(&[ + "bash", + "-lc", + r#"applypatch <<'PATCH' +*** Begin Patch +*** Add File: foo ++hi +*** End Patch +PATCH"#, + ]); + + match maybe_parse_apply_patch( + &args, + &PathUri::parse("file:///workspace").expect("valid POSIX test cwd"), + ) { + MaybeApplyPatch::Body(ApplyPatchArgs { hunks, workdir, .. }) => { + assert_eq!(workdir, None); + assert_eq!( + hunks, + vec![Hunk::AddFile { + path: PathBuf::from("foo"), + contents: "hi\n".to_string() + }] + ); + } + result => panic!("expected MaybeApplyPatch::Body got {result:?}"), + } + } + + #[tokio::test] + async fn test_powershell_heredoc() { + let script = heredoc_script(""); + assert_match_args(args_powershell(&script), /*expected_workdir*/ None); + } + #[tokio::test] + async fn test_powershell_heredoc_no_profile() { + let script = heredoc_script(""); + assert_match_args( + args_powershell_no_profile(&script), + /*expected_workdir*/ None, + ); + } + #[tokio::test] + async fn test_pwsh_heredoc() { + let script = heredoc_script(""); + assert_match_args(args_pwsh(&script), /*expected_workdir*/ None); + } + + #[tokio::test] + async fn test_apply_patch_interception_uses_cwd_convention_for_windows_pwsh_path() { + let script = heredoc_script(""); + assert_match_args_with_cwd( + strs_to_strings(&[ + r"C:\Program Files\PowerShell\7\pwsh.exe", + "-NoProfile", + "-Command", + &script, + ]), + &PathUri::parse("file:///C:/windows").expect("valid Windows test cwd"), + /*expected_workdir*/ None, + ); + } + + #[tokio::test] + async fn test_cmd_heredoc_with_cd() { + let script = heredoc_script("cd foo && "); + assert_match_args(args_cmd(&script), Some("foo")); + } + + #[tokio::test] + async fn test_heredoc_with_leading_cd() { + assert_match(&heredoc_script("cd foo && "), Some("foo")); + } + + #[tokio::test] + async fn test_cd_with_semicolon_is_ignored() { + assert_not_match(&heredoc_script("cd foo; ")); + } + + #[tokio::test] + async fn test_cd_or_apply_patch_is_ignored() { + assert_not_match(&heredoc_script("cd bar || ")); + } + + #[tokio::test] + async fn test_cd_pipe_apply_patch_is_ignored() { + assert_not_match(&heredoc_script("cd bar | ")); + } + + #[tokio::test] + async fn test_cd_single_quoted_path_with_spaces() { + assert_match(&heredoc_script("cd 'foo bar' && "), Some("foo bar")); + } + + #[tokio::test] + async fn test_cd_double_quoted_path_with_spaces() { + assert_match(&heredoc_script("cd \"foo bar\" && "), Some("foo bar")); + } + + #[tokio::test] + async fn test_echo_and_apply_patch_is_ignored() { + assert_not_match(&heredoc_script("echo foo && ")); + } + + #[tokio::test] + async fn test_apply_patch_with_arg_is_ignored() { + let script = "apply_patch foo <<'PATCH'\n*** Begin Patch\n*** Add File: foo\n+hi\n*** End Patch\nPATCH"; + assert_not_match(script); + } + + #[tokio::test] + async fn test_double_cd_then_apply_patch_is_ignored() { + assert_not_match(&heredoc_script("cd foo && cd bar && ")); + } + + #[tokio::test] + async fn test_cd_two_args_is_ignored() { + assert_not_match(&heredoc_script("cd foo bar && ")); + } + + #[tokio::test] + async fn test_cd_then_apply_patch_then_extra_is_ignored() { + let script = heredoc_script_ps("cd bar && ", " && echo done"); + assert_not_match(&script); + } + + #[tokio::test] + async fn test_echo_then_cd_and_apply_patch_is_ignored() { + // Ensure preceding commands before the `cd && apply_patch <<...` sequence do not match. + assert_not_match(&heredoc_script("echo foo; cd bar && ")); + } + + #[tokio::test] + async fn test_unified_diff_last_line_replacement() { + // Replace the very last line of the file. + let dir = tempdir().unwrap(); + let path = dir.path().join("last.txt"); + fs::write(&path, "foo\nbar\nbaz\n").unwrap(); + + let patch = wrap_patch(&format!( + r#"*** Update File: {} +@@ + foo + bar +-baz ++BAZ +"#, + path.display() + )); + + let patch = parse_patch(&patch).unwrap(); + let chunks = match patch.hunks.as_slice() { + [Hunk::UpdateFile { chunks, .. }] => chunks, + _ => panic!("Expected a single UpdateFile hunk"), + }; + + let path_uri = PathUri::from_host_native_path(&path).expect("absolute test path"); + let diff = + unified_diff_from_chunks(&path_uri, chunks, LOCAL_FS.as_ref(), /*sandbox*/ None) + .await + .unwrap(); + let expected_diff = r#"@@ -2,2 +2,2 @@ + bar +-baz ++BAZ +"#; + let expected = ApplyPatchFileUpdate { + unified_diff: expected_diff.to_string(), + original_content: "foo\nbar\nbaz\n".to_string(), + content: "foo\nbar\nBAZ\n".to_string(), + }; + assert_eq!(expected, diff); + } + + #[tokio::test] + async fn test_unified_diff_insert_at_eof() { + // Insert a new line at end‑of‑file. + let dir = tempdir().unwrap(); + let path = dir.path().join("insert.txt"); + fs::write(&path, "foo\nbar\nbaz\n").unwrap(); + + let patch = wrap_patch(&format!( + r#"*** Update File: {} +@@ ++quux +*** End of File +"#, + path.display() + )); + + let patch = parse_patch(&patch).unwrap(); + let chunks = match patch.hunks.as_slice() { + [Hunk::UpdateFile { chunks, .. }] => chunks, + _ => panic!("Expected a single UpdateFile hunk"), + }; + + let path_uri = PathUri::from_host_native_path(&path).expect("absolute test path"); + let diff = + unified_diff_from_chunks(&path_uri, chunks, LOCAL_FS.as_ref(), /*sandbox*/ None) + .await + .unwrap(); + let expected_diff = r#"@@ -3 +3,2 @@ + baz ++quux +"#; + let expected = ApplyPatchFileUpdate { + unified_diff: expected_diff.to_string(), + original_content: "foo\nbar\nbaz\n".to_string(), + content: "foo\nbar\nbaz\nquux\n".to_string(), + }; + assert_eq!(expected, diff); + } + + #[tokio::test] + async fn test_apply_patch_should_resolve_absolute_paths_in_cwd() { + let session_dir = tempdir().unwrap(); + let relative_path = "source.txt"; + + // Note that we need this file to exist for the patch to be "verified" + // and parsed correctly. + let session_file_path = session_dir.path().join(relative_path); + fs::write(&session_file_path, "session directory content\n").unwrap(); + + let argv = vec![ + "apply_patch".to_string(), + r#"*** Begin Patch +*** Update File: source.txt +@@ +-session directory content ++updated session directory content +*** End Patch"# + .to_string(), + ]; + + let result = maybe_parse_apply_patch_verified( + &argv, + &PathUri::from_host_native_path(session_dir.path()).expect("absolute test path"), + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await; + + // Verify the patch contents - as otherwise we may have pulled contents + // from the wrong file (as we're using relative paths) + assert_eq!( + result, + MaybeApplyPatchVerified::Body(ApplyPatchAction { + changes: HashMap::from([( + PathUri::from_host_native_path(session_dir.path().join(relative_path)) + .expect("absolute test path"), + ApplyPatchFileChange::Update { + unified_diff: r#"@@ -1 +1 @@ +-session directory content ++updated session directory content +"# + .to_string(), + move_path: None, + new_content: "updated session directory content\n".to_string(), + }, + )]), + update_file_mode: ApplyPatchFileUpdateMode::default(), + patch: argv[1].clone(), + cwd: PathUri::from_host_native_path(session_dir.path()) + .expect("absolute test path"), + }) + ); + } + + #[tokio::test] + async fn test_apply_patch_resolves_move_path_with_effective_cwd() { + let session_dir = tempdir().unwrap(); + let worktree_rel = "alt"; + let worktree_dir = session_dir.path().join(worktree_rel); + fs::create_dir_all(&worktree_dir).unwrap(); + + let source_name = "old.txt"; + let dest_name = "renamed.txt"; + let source_path = worktree_dir.join(source_name); + fs::write(&source_path, "before\n").unwrap(); + + let patch = wrap_patch(&format!( + r#"*** Update File: {source_name} +*** Move to: {dest_name} +@@ +-before ++after"# + )); + + let shell_script = format!("cd {worktree_rel} && apply_patch <<'PATCH'\n{patch}\nPATCH"); + let argv = vec!["bash".into(), "-lc".into(), shell_script]; + + let result = maybe_parse_apply_patch_verified( + &argv, + &PathUri::from_host_native_path(session_dir.path()).expect("absolute test path"), + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await; + let action = match result { + MaybeApplyPatchVerified::Body(action) => action, + other => panic!("expected verified body, got {other:?}"), + }; + + assert_eq!( + action.cwd.to_abs_path().unwrap().as_path(), + worktree_dir.as_path() + ); + + let source_path = PathUri::from_host_native_path(worktree_dir.join(source_name)) + .expect("absolute test path"); + let change = action + .changes() + .get(&source_path) + .expect("source file change present"); + + match change { + ApplyPatchFileChange::Update { move_path, .. } => { + let expected_move_path = + PathUri::from_host_native_path(worktree_dir.join(dest_name)) + .expect("absolute test path"); + assert_eq!(move_path.as_ref(), Some(&expected_move_path)); + } + other => panic!("expected update change, got {other:?}"), + } + } + + #[tokio::test] + async fn test_unreadable_destinations_still_verify() { + let session_dir = tempdir().unwrap(); + fs::write(session_dir.path().join("binary.dat"), [0xff, 0xfe, 0xfd]).unwrap(); + let cwd = PathUri::from_host_native_path(session_dir.path()).expect("absolute test path"); + let add_argv = vec![ + "apply_patch".to_string(), + "*** Begin Patch\n*** Add File: binary.dat\n+text\n*** End Patch".to_string(), + ]; + fs::write(session_dir.path().join("source.txt"), "before\n").unwrap(); + let move_argv = vec![ + "apply_patch".to_string(), + "*** Begin Patch\n*** Update File: source.txt\n*** Move to: binary.dat\n@@\n-before\n+after\n*** End Patch".to_string(), + ]; + + for argv in [add_argv, move_argv] { + let result = maybe_parse_apply_patch_verified( + &argv, + &cwd, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await; + + assert!(matches!(result, MaybeApplyPatchVerified::Body(_))); + } + } + + #[cfg(unix)] + #[tokio::test] + async fn test_delete_symlink_still_verifies() { + use std::os::unix::fs::symlink; + + let session_dir = tempdir().unwrap(); + fs::write(session_dir.path().join("target.txt"), "target\n").unwrap(); + symlink( + session_dir.path().join("target.txt"), + session_dir.path().join("link.txt"), + ) + .unwrap(); + let argv = vec![ + "apply_patch".to_string(), + "*** Begin Patch\n*** Delete File: link.txt\n*** End Patch".to_string(), + ]; + + let result = maybe_parse_apply_patch_verified( + &argv, + &PathUri::from_host_native_path(session_dir.path()).expect("absolute test path"), + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await; + + assert!(matches!(result, MaybeApplyPatchVerified::Body(_))); + } +} diff --git a/codex-rs/apply-patch/src/lib.rs b/codex-rs/apply-patch/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..0ee7a4e7ec811751c68984b7e05a0d969d7ae089 --- /dev/null +++ b/codex-rs/apply-patch/src/lib.rs @@ -0,0 +1,1444 @@ +mod file_update; +mod invocation; +mod parser; +mod seek_sequence; +mod standalone_executable; +mod streaming_parser; +mod text_file; + +use std::collections::HashMap; +use std::io; +use std::path::PathBuf; + +use anyhow::Context; +use anyhow::Result; +use codex_exec_server::CreateDirectoryOptions; +use codex_exec_server::ExecutorFileSystem; +use codex_exec_server::FileSystemSandboxContext; +use codex_exec_server::GetMetadataOptions; +use codex_exec_server::ReadFileOptions; +use codex_exec_server::RemoveOptions; +use codex_exec_server::WriteFileOptions; +use codex_utils_path_uri::PathUri; +use codex_utils_path_uri::PathUriParseError; +pub use parser::Hunk; +pub use parser::ParseError; +use parser::ParseError::*; +pub use parser::UpdateFileChunk; +pub use parser::parse_patch; +pub use streaming_parser::StreamingPatchParser; +use thiserror::Error; + +use file_update::AppliedPatch; +pub use file_update::ApplyPatchFileUpdate; +use file_update::derive_new_contents_from_chunks; +pub use file_update::unified_diff_from_chunks; +pub use file_update::unified_diff_from_chunks_with_context; +pub(crate) use file_update::unified_diff_from_chunks_with_mode; +pub use invocation::MaybeApplyPatch; +pub use invocation::maybe_parse_apply_patch; +pub use invocation::maybe_parse_apply_patch_verified; +pub use invocation::maybe_parse_apply_patch_verified_with_mode; +pub use invocation::verify_apply_patch_args; +pub use invocation::verify_apply_patch_args_with_mode; +pub use standalone_executable::main; + +use crate::invocation::ExtractHeredocError; + +/// Special argv[1] flag used when the Codex executable self-invokes to run the +/// internal `apply_patch` path. +/// +/// Although this constant lives in `codex-apply-patch` (to avoid forcing +/// `codex-arg0` to depend on `codex-core`), it remains part of the "codex core" +/// process-invocation contract for the standalone `apply_patch` command +/// surface. +pub const CODEX_CORE_APPLY_PATCH_ARG1: &str = "--codex-run-as-apply-patch"; + +/// Internal environment variable used to carry the selected update mode +/// through the arg0-dispatched standalone executable. +pub const CODEX_APPLY_PATCH_PRESERVE_LINE_ENDINGS_ENV_VAR: &str = + "CODEX_APPLY_PATCH_PRESERVE_LINE_ENDINGS"; + +/// Controls how updates reconstruct the target file after matching a patch. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum ApplyPatchFileUpdateMode { + /// Preserve the historical behavior of normalizing updated files to LF. + #[default] + NormalizeToLf, + /// Preserve existing line endings and use the file's preferred ending for new lines. + PreserveLineEndings, +} + +/// Policy for one patch application. Standalone callers follow symlinks by default. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct ApplyPatchOptions { + pub update_file_mode: ApplyPatchFileUpdateMode, + /// Whether filesystem operations may resolve symlinks in any path component. + pub follow_symlinks: bool, +} + +impl Default for ApplyPatchOptions { + fn default() -> Self { + Self { + update_file_mode: ApplyPatchFileUpdateMode::default(), + follow_symlinks: true, + } + } +} + +/// Reads the update mode selected for an arg0-dispatched `apply_patch` process. +#[doc(hidden)] +pub fn apply_patch_file_update_mode_from_env() -> ApplyPatchFileUpdateMode { + match std::env::var(CODEX_APPLY_PATCH_PRESERVE_LINE_ENDINGS_ENV_VAR).as_deref() { + Ok("1") => ApplyPatchFileUpdateMode::PreserveLineEndings, + _ => ApplyPatchFileUpdateMode::NormalizeToLf, + } +} + +#[derive(Debug, Error, PartialEq)] +pub enum ApplyPatchError { + #[error(transparent)] + ParseError(#[from] ParseError), + #[error(transparent)] + IoError(#[from] IoError), + /// Error that occurs while computing replacements when applying patch chunks + #[error("{0}")] + ComputeReplacements(String), + /// A patch path could not be resolved as a path URI. + #[error(transparent)] + PathUri(#[from] PathUriParseError), + /// A raw patch body was provided without an explicit `apply_patch` invocation. + #[error( + "patch detected without explicit call to apply_patch. Rerun as [\"apply_patch\", \"\"]" + )] + ImplicitInvocation, +} + +impl From for ApplyPatchError { + fn from(err: std::io::Error) -> Self { + ApplyPatchError::IoError(IoError { + context: "I/O error".to_string(), + source: err, + }) + } +} + +impl From<&std::io::Error> for ApplyPatchError { + fn from(err: &std::io::Error) -> Self { + ApplyPatchError::IoError(IoError { + context: "I/O error".to_string(), + source: std::io::Error::new(err.kind(), err.to_string()), + }) + } +} + +#[derive(Debug, Error)] +#[error("{context}: {source}")] +pub struct IoError { + context: String, + #[source] + source: std::io::Error, +} + +impl PartialEq for IoError { + fn eq(&self, other: &Self) -> bool { + self.context == other.context && self.source.to_string() == other.source.to_string() + } +} + +/// Both the raw PATCH argument to `apply_patch` as well as the PATCH argument +/// parsed into hunks. +#[derive(Debug, PartialEq)] +pub struct ApplyPatchArgs { + pub patch: String, + pub hunks: Vec, + pub workdir: Option, + pub environment_id: Option, +} + +#[derive(Debug, PartialEq)] +pub enum ApplyPatchFileChange { + Add { + content: String, + }, + Delete { + content: String, + }, + Update { + unified_diff: String, + move_path: Option, + /// new_content that will result after the unified_diff is applied. + new_content: String, + }, +} + +#[derive(Debug, PartialEq)] +pub enum MaybeApplyPatchVerified { + /// `argv` corresponded to an `apply_patch` invocation, and these are the + /// resulting proposed file changes. + Body(ApplyPatchAction), + /// `argv` could not be parsed to determine whether it corresponds to an + /// `apply_patch` invocation. + ShellParseError(ExtractHeredocError), + /// `argv` corresponded to an `apply_patch` invocation, but it could not + /// be fulfilled due to the specified error. + CorrectnessError(ApplyPatchError), + /// `argv` decidedly did not correspond to an `apply_patch` invocation. + NotApplyPatch, +} + +/// ApplyPatchAction is the result of parsing an `apply_patch` command. By +/// construction, all paths should be absolute paths. +#[derive(Debug, PartialEq)] +pub struct ApplyPatchAction { + changes: HashMap, + + update_file_mode: ApplyPatchFileUpdateMode, + + /// The raw patch argument that can be used to apply the patch. i.e., if the + /// original arg was parsed in "lenient" mode with a + /// heredoc, this should be the value without the heredoc wrapper. + pub patch: String, + + /// The working directory that was used to resolve relative paths in the patch. + pub cwd: PathUri, +} + +impl ApplyPatchAction { + pub fn is_empty(&self) -> bool { + self.changes.is_empty() + } + + /// Returns the changes that would be made by applying the patch. + pub fn changes(&self) -> &HashMap { + &self.changes + } + + /// Returns the update mode selected while the patch was verified. + pub fn update_file_mode(&self) -> ApplyPatchFileUpdateMode { + self.update_file_mode + } + + /// Should be used exclusively for testing. (Not worth the overhead of + /// creating a feature flag for this.) + pub fn new_add_for_test(path: &PathUri, content: String) -> Self { + #[expect(clippy::expect_used)] + let filename = path.basename().expect("path should not be empty"); + let patch = format!( + r#"*** Begin Patch +*** Update File: {filename} +@@ ++ {content} +*** End Patch"#, + ); + let changes = HashMap::from([(path.clone(), ApplyPatchFileChange::Add { content })]); + #[expect(clippy::expect_used)] + Self { + changes, + update_file_mode: ApplyPatchFileUpdateMode::default(), + cwd: path.parent().expect("path should have parent"), + patch, + } + } +} + +/// Textual file changes that were actually committed while applying a patch. +#[derive(Clone, Debug, PartialEq)] +pub struct AppliedPatchDelta { + changes: Vec, + exact: bool, +} + +impl AppliedPatchDelta { + fn new(changes: Vec, exact: bool) -> Self { + Self { changes, exact } + } + + fn empty() -> Self { + Self::new(Vec::new(), /*exact*/ true) + } + + pub fn changes(&self) -> &[AppliedPatchChange] { + &self.changes + } + + pub fn is_empty(&self) -> bool { + self.changes.is_empty() + } + + pub fn is_exact(&self) -> bool { + self.exact + } + + /// Appends a later committed prefix while preserving the aggregate exactness. + pub fn append(&mut self, other: Self) { + self.changes.extend(other.changes); + self.exact &= other.exact; + } +} + +impl Default for AppliedPatchDelta { + fn default() -> Self { + Self::empty() + } +} + +/// A committed file change, preserved in the order it was applied. +#[derive(Clone, Debug, PartialEq)] +pub struct AppliedPatchChange { + pub path: PathUri, + pub change: AppliedPatchFileChange, +} + +#[derive(Clone, Debug, PartialEq)] +pub enum AppliedPatchFileChange { + Add { + content: String, + overwritten_content: Option, + }, + Delete { + content: String, + }, + Update { + move_path: Option, + old_content: String, + overwritten_move_content: Option, + new_content: String, + }, +} + +/// A failed patch application together with the textual mutations that were +/// definitely committed before the failure was observed. +#[derive(Debug, Error)] +#[error("{error}")] +pub struct ApplyPatchFailure { + #[source] + error: ApplyPatchError, + delta: AppliedPatchDelta, +} + +impl ApplyPatchFailure { + fn new(error: ApplyPatchError, delta: AppliedPatchDelta) -> Self { + Self { error, delta } + } + + fn without_delta(error: ApplyPatchError) -> Self { + Self::new(error, AppliedPatchDelta::empty()) + } + + pub fn delta(&self) -> &AppliedPatchDelta { + &self.delta + } + + pub fn into_parts(self) -> (ApplyPatchError, AppliedPatchDelta) { + (self.error, self.delta) + } +} + +/// Applies the patch and prints the result to stdout/stderr. +pub async fn apply_patch( + patch: &str, + cwd: &PathUri, + stdout: &mut impl std::io::Write, + stderr: &mut impl std::io::Write, + fs: &dyn ExecutorFileSystem, + sandbox: Option<&FileSystemSandboxContext>, +) -> Result { + apply_patch_with_options( + patch, + ApplyPatchOptions::default(), + cwd, + stdout, + stderr, + fs, + sandbox, + ) + .await +} + +/// Applies the patch using the selected options and prints the result +/// to stdout/stderr. +pub async fn apply_patch_with_options( + patch: &str, + options: ApplyPatchOptions, + cwd: &PathUri, + stdout: &mut impl std::io::Write, + stderr: &mut impl std::io::Write, + fs: &dyn ExecutorFileSystem, + sandbox: Option<&FileSystemSandboxContext>, +) -> Result { + let hunks = match parse_patch(patch) { + Ok(source) => source.hunks, + Err(e) => { + match &e { + InvalidPatchError(message) => { + writeln!(stderr, "Invalid patch: {message}") + .map_err(ApplyPatchError::from) + .map_err(ApplyPatchFailure::without_delta)?; + } + InvalidHunkError { + message, + line_number, + } => { + writeln!( + stderr, + "Invalid patch hunk on line {line_number}: {message}" + ) + .map_err(ApplyPatchError::from) + .map_err(ApplyPatchFailure::without_delta)?; + } + } + return Err(ApplyPatchFailure::without_delta( + ApplyPatchError::ParseError(e), + )); + } + }; + + apply_hunks_with_options(&hunks, options, cwd, stdout, stderr, fs, sandbox).await +} + +/// Applies hunks and continues to update stdout/stderr +pub async fn apply_hunks( + hunks: &[Hunk], + cwd: &PathUri, + stdout: &mut impl std::io::Write, + stderr: &mut impl std::io::Write, + fs: &dyn ExecutorFileSystem, + sandbox: Option<&FileSystemSandboxContext>, +) -> Result { + apply_hunks_with_options( + hunks, + ApplyPatchOptions::default(), + cwd, + stdout, + stderr, + fs, + sandbox, + ) + .await +} + +/// Applies hunks using the selected file-update mode and continues to update +/// stdout/stderr. +async fn apply_hunks_with_options( + hunks: &[Hunk], + options: ApplyPatchOptions, + cwd: &PathUri, + stdout: &mut impl std::io::Write, + stderr: &mut impl std::io::Write, + fs: &dyn ExecutorFileSystem, + sandbox: Option<&FileSystemSandboxContext>, +) -> Result { + let mut delta = AppliedPatchDelta::empty(); + match apply_hunks_to_files(hunks, options, cwd, fs, sandbox, &mut delta).await { + Ok(affected_paths) => { + print_summary(&affected_paths, stdout).map_err(|error| { + ApplyPatchFailure::new(ApplyPatchError::from(error), delta.clone()) + })?; + Ok(delta) + } + Err(error) => { + let msg = error.to_string(); + writeln!(stderr, "{msg}").map_err(|error| { + ApplyPatchFailure::new(ApplyPatchError::from(error), delta.clone()) + })?; + let error = if let Some(io) = error.downcast_ref::() { + ApplyPatchError::from(io) + } else { + ApplyPatchError::IoError(IoError { + context: msg, + source: std::io::Error::other(error), + }) + }; + Err(ApplyPatchFailure::new(error, delta)) + } + } +} + +/// Applies each parsed patch hunk to the filesystem. +/// Returns an error if any of the changes could not be applied. +/// Tracks file paths affected by applying a patch, preserving the path spelling +/// from the patch for user-facing summaries. +pub struct AffectedPaths { + pub added: Vec, + pub modified: Vec, + pub deleted: Vec, +} + +/// Apply the hunks to the filesystem, returning which files were added, modified, or deleted. +/// Returns an error if the patch could not be applied. +async fn apply_hunks_to_files( + hunks: &[Hunk], + options: ApplyPatchOptions, + cwd: &PathUri, + fs: &dyn ExecutorFileSystem, + sandbox: Option<&FileSystemSandboxContext>, + delta: &mut AppliedPatchDelta, +) -> anyhow::Result { + let ApplyPatchOptions { + update_file_mode, + follow_symlinks, + } = options; + if hunks.is_empty() { + anyhow::bail!("No files were modified."); + } + + let mut added: Vec = Vec::new(); + let mut modified: Vec = Vec::new(); + let mut deleted: Vec = Vec::new(); + // A failed write can still have modified the target before surfacing an + // error (for example by truncating before ENOSPC), so the accumulated + // delta is no longer exact when a write fails. + macro_rules! try_write { + ($result:expr) => { + match $result { + Ok(value) => value, + Err(error) => { + delta.exact = false; + return Err(anyhow::Error::from(error)); + } + } + }; + } + + for hunk in hunks { + let affected_path = hunk.path().to_path_buf(); + let path_uri = hunk.resolve_path(cwd)?; + match hunk { + Hunk::AddFile { contents, .. } => { + let overwritten_content = read_optional_file_text_for_delta( + &path_uri, + fs, + follow_symlinks, + sandbox, + &mut delta.exact, + ) + .await; + try_write!( + write_file_with_missing_parent_retry( + fs, + &path_uri, + contents.clone().into_bytes(), + follow_symlinks, + sandbox, + ) + .await + ); + delta.changes.push(AppliedPatchChange { + path: path_uri, + change: AppliedPatchFileChange::Add { + content: contents.clone(), + overwritten_content, + }, + }); + added.push(affected_path); + } + Hunk::DeleteFile { .. } => { + note_existing_path_delta_support( + &path_uri, + fs, + follow_symlinks, + sandbox, + &mut delta.exact, + ) + .await; + let deleted_content = fs + .read_file_text(&path_uri, ReadFileOptions { follow_symlinks }, sandbox) + .await + .ok(); + if deleted_content.is_none() { + delta.exact = false; + } + ensure_not_directory(&path_uri, fs, follow_symlinks, sandbox) + .await + .with_context(|| { + format!( + "Failed to delete file {}", + path_uri.inferred_native_path_string() + ) + })?; + if let Err(error) = fs + .remove( + &path_uri, + RemoveOptions { + recursive: false, + force: false, + follow_symlinks, + }, + sandbox, + ) + .await + .with_context(|| { + format!( + "Failed to delete file {}", + path_uri.inferred_native_path_string() + ) + }) + { + delta.exact &= remove_failure_was_side_effect_free( + &path_uri, + deleted_content.as_deref(), + fs, + follow_symlinks, + sandbox, + ) + .await; + return Err(error); + } + if let Some(content) = deleted_content { + delta.changes.push(AppliedPatchChange { + path: path_uri, + change: AppliedPatchFileChange::Delete { content }, + }); + } + deleted.push(affected_path); + } + Hunk::UpdateFile { + move_path, chunks, .. + } => { + note_existing_path_delta_support( + &path_uri, + fs, + follow_symlinks, + sandbox, + &mut delta.exact, + ) + .await; + let AppliedPatch { + original_contents, + new_contents, + } = derive_new_contents_from_chunks( + &path_uri, + chunks, + update_file_mode, + fs, + follow_symlinks, + sandbox, + ) + .await?; + if let Some(dest) = move_path { + let dest_uri = cwd.join(&dest.to_string_lossy())?; + let overwritten_move_content = read_optional_file_text_for_delta( + &dest_uri, + fs, + follow_symlinks, + sandbox, + &mut delta.exact, + ) + .await; + try_write!( + write_file_with_missing_parent_retry( + fs, + &dest_uri, + new_contents.clone().into_bytes(), + follow_symlinks, + sandbox, + ) + .await + ); + let dest_write_change_index = delta.changes.len(); + delta.changes.push(AppliedPatchChange { + path: dest_uri.clone(), + change: AppliedPatchFileChange::Add { + content: new_contents.clone(), + overwritten_content: overwritten_move_content.clone(), + }, + }); + ensure_not_directory(&path_uri, fs, follow_symlinks, sandbox) + .await + .with_context(|| { + format!( + "Failed to remove original {}", + path_uri.inferred_native_path_string() + ) + })?; + if let Err(error) = fs + .remove( + &path_uri, + RemoveOptions { + recursive: false, + force: false, + follow_symlinks, + }, + sandbox, + ) + .await + .with_context(|| { + format!( + "Failed to remove original {}", + path_uri.inferred_native_path_string() + ) + }) + { + delta.exact &= remove_failure_was_side_effect_free( + &path_uri, + Some(&original_contents), + fs, + follow_symlinks, + sandbox, + ) + .await; + return Err(error); + } + delta.changes[dest_write_change_index] = AppliedPatchChange { + path: path_uri, + change: AppliedPatchFileChange::Update { + move_path: Some(dest_uri), + old_content: original_contents, + overwritten_move_content, + new_content: new_contents, + }, + }; + modified.push(affected_path); + } else { + try_write!( + fs.write_file( + &path_uri, + new_contents.clone().into_bytes(), + WriteFileOptions { follow_symlinks }, + sandbox, + ) + .await + .with_context(|| format!( + "Failed to write file {}", + path_uri.inferred_native_path_string() + )) + ); + delta.changes.push(AppliedPatchChange { + path: path_uri, + change: AppliedPatchFileChange::Update { + move_path: None, + old_content: original_contents, + overwritten_move_content: None, + new_content: new_contents, + }, + }); + modified.push(affected_path); + } + } + } + } + Ok(AffectedPaths { + added, + modified, + deleted, + }) +} + +async fn ensure_not_directory( + path: &PathUri, + fs: &dyn ExecutorFileSystem, + follow_symlinks: bool, + sandbox: Option<&FileSystemSandboxContext>, +) -> io::Result<()> { + let metadata = fs + .get_metadata(path, GetMetadataOptions { follow_symlinks }, sandbox) + .await?; + if metadata.is_directory { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "path is a directory", + )); + } + Ok(()) +} + +async fn remove_failure_was_side_effect_free( + path: &PathUri, + expected_content: Option<&str>, + fs: &dyn ExecutorFileSystem, + follow_symlinks: bool, + sandbox: Option<&FileSystemSandboxContext>, +) -> bool { + match expected_content { + Some(expected_content) => fs + .read_file_text(path, ReadFileOptions { follow_symlinks }, sandbox) + .await + .is_ok_and(|content| content == expected_content), + None => false, + } +} + +async fn read_optional_file_text_for_delta( + path: &PathUri, + fs: &dyn ExecutorFileSystem, + follow_symlinks: bool, + sandbox: Option<&FileSystemSandboxContext>, + exact: &mut bool, +) -> Option { + note_existing_path_delta_support(path, fs, follow_symlinks, sandbox, exact).await; + match fs + .read_file_text(path, ReadFileOptions { follow_symlinks }, sandbox) + .await + { + Ok(content) => Some(content), + Err(source) if source.kind() == io::ErrorKind::NotFound => None, + Err(_) => { + *exact = false; + None + } + } +} + +async fn note_existing_path_delta_support( + path: &PathUri, + fs: &dyn ExecutorFileSystem, + follow_symlinks: bool, + sandbox: Option<&FileSystemSandboxContext>, + exact: &mut bool, +) { + match fs + .get_metadata(path, GetMetadataOptions { follow_symlinks }, sandbox) + .await + { + Ok(metadata) if metadata.is_file && !metadata.is_symlink => {} + Ok(_) => *exact = false, + Err(source) if source.kind() == io::ErrorKind::NotFound => {} + Err(_) => *exact = false, + } +} + +async fn write_file_with_missing_parent_retry( + fs: &dyn ExecutorFileSystem, + path: &PathUri, + contents: Vec, + follow_symlinks: bool, + sandbox: Option<&FileSystemSandboxContext>, +) -> anyhow::Result<()> { + match fs + .write_file( + path, + contents.clone(), + WriteFileOptions { follow_symlinks }, + sandbox, + ) + .await + { + Ok(()) => Ok(()), + Err(err) if err.kind() == io::ErrorKind::NotFound => { + if let Some(parent) = path.parent() { + fs.create_directory( + &parent, + CreateDirectoryOptions { + recursive: true, + follow_symlinks, + }, + sandbox, + ) + .await + .with_context(|| { + format!( + "Failed to create parent directories for {}", + path.inferred_native_path_string() + ) + })?; + } + fs.write_file( + path, + contents, + WriteFileOptions { follow_symlinks }, + sandbox, + ) + .await + .with_context(|| { + format!( + "Failed to write file {}", + path.inferred_native_path_string() + ) + })?; + Ok(()) + } + Err(err) => Err(err).with_context(|| { + format!( + "Failed to write file {}", + path.inferred_native_path_string() + ) + }), + } +} + +/// Print the summary of changes in git-style format. +/// Write a summary of changes to the given writer. +pub fn print_summary( + affected: &AffectedPaths, + out: &mut impl std::io::Write, +) -> std::io::Result<()> { + writeln!(out, "Success. Updated the following files:")?; + for path in &affected.added { + writeln!(out, "A {}", path.display())?; + } + for path in &affected.modified { + writeln!(out, "M {}", path.display())?; + } + for path in &affected.deleted { + writeln!(out, "D {}", path.display())?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use codex_exec_server::LOCAL_FS; + use pretty_assertions::assert_eq; + use std::fs; + use tempfile::tempdir; + + /// Helper to construct a patch with the given body. + fn wrap_patch(body: &str) -> String { + format!("*** Begin Patch\n{body}\n*** End Patch") + } + + #[tokio::test] + async fn test_add_file_hunk_creates_file_with_contents() { + let dir = tempdir().unwrap(); + let path = dir.path().join("add.txt"); + let patch = wrap_patch(&format!( + r#"*** Add File: {} ++ab ++cd"#, + path.display() + )); + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + apply_patch( + &patch, + &PathUri::from_host_native_path(dir.path()).expect("absolute test path"), + &mut stdout, + &mut stderr, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .unwrap(); + // Verify expected stdout and stderr outputs. + let stdout_str = String::from_utf8(stdout).unwrap(); + let stderr_str = String::from_utf8(stderr).unwrap(); + let expected_out = format!( + "Success. Updated the following files:\nA {}\n", + path.display() + ); + assert_eq!(stdout_str, expected_out); + assert_eq!(stderr_str, ""); + let contents = fs::read_to_string(path).unwrap(); + assert_eq!(contents, "ab\ncd\n"); + } + + #[tokio::test] + async fn test_apply_patch_hunks_accept_relative_and_absolute_paths() { + let dir = tempdir().unwrap(); + let cwd = PathUri::from_host_native_path(dir.path()).expect("absolute test path"); + let relative_add = dir.path().join("relative-add.txt"); + let absolute_add = dir.path().join("absolute-add.txt"); + let relative_delete = dir.path().join("relative-delete.txt"); + let absolute_delete = dir.path().join("absolute-delete.txt"); + let relative_update = dir.path().join("relative-update.txt"); + let absolute_update = dir.path().join("absolute-update.txt"); + fs::write(&relative_delete, "delete relative\n").unwrap(); + fs::write(&absolute_delete, "delete absolute\n").unwrap(); + fs::write(&relative_update, "relative old\n").unwrap(); + fs::write(&absolute_update, "absolute old\n").unwrap(); + + let patch = wrap_patch(&format!( + r#"*** Add File: relative-add.txt ++relative add +*** Add File: {} ++absolute add +*** Delete File: relative-delete.txt +*** Delete File: {} +*** Update File: relative-update.txt +@@ +-relative old ++relative new +*** Update File: {} +@@ +-absolute old ++absolute new"#, + absolute_add.display(), + absolute_delete.display(), + absolute_update.display(), + )); + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + + apply_patch( + &patch, + &cwd, + &mut stdout, + &mut stderr, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .unwrap(); + + assert_eq!(fs::read_to_string(&relative_add).unwrap(), "relative add\n"); + assert_eq!(fs::read_to_string(&absolute_add).unwrap(), "absolute add\n"); + assert!(!relative_delete.exists()); + assert!(!absolute_delete.exists()); + assert_eq!( + fs::read_to_string(&relative_update).unwrap(), + "relative new\n" + ); + assert_eq!( + fs::read_to_string(&absolute_update).unwrap(), + "absolute new\n" + ); + assert_eq!(String::from_utf8(stderr).unwrap(), ""); + assert_eq!( + String::from_utf8(stdout).unwrap(), + format!( + "Success. Updated the following files:\nA relative-add.txt\nA {}\nM relative-update.txt\nM {}\nD relative-delete.txt\nD {}\n", + absolute_add.display(), + absolute_update.display(), + absolute_delete.display(), + ) + ); + } + + #[tokio::test] + async fn test_delete_file_hunk_removes_file() { + let dir = tempdir().unwrap(); + let path = dir.path().join("del.txt"); + fs::write(&path, "x").unwrap(); + let patch = wrap_patch(&format!("*** Delete File: {}", path.display())); + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + apply_patch( + &patch, + &PathUri::from_host_native_path(dir.path()).expect("absolute test path"), + &mut stdout, + &mut stderr, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .unwrap(); + let stdout_str = String::from_utf8(stdout).unwrap(); + let stderr_str = String::from_utf8(stderr).unwrap(); + let expected_out = format!( + "Success. Updated the following files:\nD {}\n", + path.display() + ); + assert_eq!(stdout_str, expected_out); + assert_eq!(stderr_str, ""); + assert!(!path.exists()); + } + + #[tokio::test] + async fn test_update_file_hunk_modifies_content() { + let dir = tempdir().unwrap(); + let path = dir.path().join("update.txt"); + fs::write(&path, "foo\nbar\n").unwrap(); + let patch = wrap_patch(&format!( + r#"*** Update File: {} +@@ + foo +-bar ++baz"#, + path.display() + )); + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + apply_patch( + &patch, + &PathUri::from_host_native_path(dir.path()).expect("absolute test path"), + &mut stdout, + &mut stderr, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .unwrap(); + // Validate modified file contents and expected stdout/stderr. + let stdout_str = String::from_utf8(stdout).unwrap(); + let stderr_str = String::from_utf8(stderr).unwrap(); + let expected_out = format!( + "Success. Updated the following files:\nM {}\n", + path.display() + ); + assert_eq!(stdout_str, expected_out); + assert_eq!(stderr_str, ""); + let contents = fs::read_to_string(&path).unwrap(); + assert_eq!(contents, "foo\nbaz\n"); + } + + #[tokio::test] + async fn test_update_file_hunk_can_move_file() { + let dir = tempdir().unwrap(); + let src = dir.path().join("src.txt"); + let dest = dir.path().join("dst.txt"); + fs::write(&src, "line\n").unwrap(); + let patch = wrap_patch(&format!( + r#"*** Update File: {} +*** Move to: {} +@@ +-line ++line2"#, + src.display(), + dest.display() + )); + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + apply_patch( + &patch, + &PathUri::from_host_native_path(dir.path()).expect("absolute test path"), + &mut stdout, + &mut stderr, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .unwrap(); + // Validate move semantics and expected stdout/stderr. + let stdout_str = String::from_utf8(stdout).unwrap(); + let stderr_str = String::from_utf8(stderr).unwrap(); + let expected_out = format!( + "Success. Updated the following files:\nM {}\n", + dest.display() + ); + assert_eq!(stdout_str, expected_out); + assert_eq!(stderr_str, ""); + assert!(!src.exists()); + let contents = fs::read_to_string(&dest).unwrap(); + assert_eq!(contents, "line2\n"); + } + + #[cfg(unix)] + #[tokio::test] + async fn test_failed_move_returns_committed_destination_delta() { + use std::os::unix::fs::PermissionsExt; + + let dir = tempdir().unwrap(); + let source_dir = dir.path().join("locked"); + let dest_dir = dir.path().join("out"); + fs::create_dir(&source_dir).unwrap(); + fs::create_dir(&dest_dir).unwrap(); + let src = source_dir.join("src.txt"); + let dest = dest_dir.join("dst.txt"); + fs::write(&src, "line\n").unwrap(); + fs::set_permissions(&source_dir, fs::Permissions::from_mode(0o555)).unwrap(); + + let patch = wrap_patch( + "*** Update File: locked/src.txt\n*** Move to: out/dst.txt\n@@\n-line\n+line2", + ); + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + let failure = apply_patch( + &patch, + &PathUri::from_host_native_path(dir.path()).expect("absolute test path"), + &mut stdout, + &mut stderr, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .expect_err("source removal should fail after destination write"); + + fs::set_permissions(&source_dir, fs::Permissions::from_mode(0o755)).unwrap(); + + assert!( + String::from_utf8(stderr) + .unwrap() + .contains(&format!("Failed to remove original {}", src.display())) + ); + assert_eq!( + failure.delta(), + &AppliedPatchDelta::new( + vec![AppliedPatchChange { + path: PathUri::from_host_native_path(&dest).expect("absolute destination path"), + change: AppliedPatchFileChange::Add { + content: "line2\n".to_string(), + overwritten_content: None, + }, + }], + /*exact*/ true, + ) + ); + assert_eq!(fs::read_to_string(src).unwrap(), "line\n"); + assert_eq!(fs::read_to_string(dest).unwrap(), "line2\n"); + } + + /// Verify that a single `Update File` hunk with multiple change chunks can update different + /// parts of a file and that the file is listed only once in the summary. + #[tokio::test] + async fn test_multiple_update_chunks_apply_to_single_file() { + // Start with a file containing four lines. + let dir = tempdir().unwrap(); + let path = dir.path().join("multi.txt"); + fs::write(&path, "foo\nbar\nbaz\nqux\n").unwrap(); + // Construct an update patch with two separate change chunks. + // The first chunk uses the line `foo` as context and transforms `bar` into `BAR`. + // The second chunk uses `baz` as context and transforms `qux` into `QUX`. + let patch = wrap_patch(&format!( + r#"*** Update File: {} +@@ + foo +-bar ++BAR +@@ + baz +-qux ++QUX"#, + path.display() + )); + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + apply_patch( + &patch, + &PathUri::from_host_native_path(dir.path()).expect("absolute test path"), + &mut stdout, + &mut stderr, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .unwrap(); + let stdout_str = String::from_utf8(stdout).unwrap(); + let stderr_str = String::from_utf8(stderr).unwrap(); + let expected_out = format!( + "Success. Updated the following files:\nM {}\n", + path.display() + ); + assert_eq!(stdout_str, expected_out); + assert_eq!(stderr_str, ""); + let contents = fs::read_to_string(&path).unwrap(); + assert_eq!(contents, "foo\nBAR\nbaz\nQUX\n"); + } + + /// A more involved `Update File` hunk that exercises additions, deletions and + /// replacements in separate chunks that appear in non‑adjacent parts of the + /// file. Verifies that all edits are applied and that the summary lists the + /// file only once. + #[tokio::test] + async fn test_update_file_hunk_interleaved_changes() { + let dir = tempdir().unwrap(); + let path = dir.path().join("interleaved.txt"); + + // Original file: six numbered lines. + fs::write(&path, "a\nb\nc\nd\ne\nf\n").unwrap(); + + // Patch performs: + // • Replace `b` → `B` + // • Replace `e` → `E` (using surrounding context) + // • Append new line `g` at the end‑of‑file + let patch = wrap_patch(&format!( + r#"*** Update File: {} +@@ + a +-b ++B +@@ + c + d +-e ++E +@@ + f ++g +*** End of File"#, + path.display() + )); + + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + apply_patch( + &patch, + &PathUri::from_host_native_path(dir.path()).expect("absolute test path"), + &mut stdout, + &mut stderr, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .unwrap(); + + let stdout_str = String::from_utf8(stdout).unwrap(); + let stderr_str = String::from_utf8(stderr).unwrap(); + + let expected_out = format!( + "Success. Updated the following files:\nM {}\n", + path.display() + ); + assert_eq!(stdout_str, expected_out); + assert_eq!(stderr_str, ""); + + let contents = fs::read_to_string(&path).unwrap(); + assert_eq!(contents, "a\nB\nc\nd\nE\nf\ng\n"); + } + + #[tokio::test] + async fn test_pure_addition_chunk_followed_by_removal() { + let dir = tempdir().unwrap(); + let path = dir.path().join("panic.txt"); + fs::write(&path, "line1\nline2\nline3\n").unwrap(); + let patch = wrap_patch(&format!( + r#"*** Update File: {} +@@ ++after-context ++second-line +@@ + line1 +-line2 +-line3 ++line2-replacement"#, + path.display() + )); + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + apply_patch( + &patch, + &PathUri::from_host_native_path(dir.path()).expect("absolute test path"), + &mut stdout, + &mut stderr, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .unwrap(); + let contents = fs::read_to_string(path).unwrap(); + assert_eq!( + contents, + "line1\nline2-replacement\nafter-context\nsecond-line\n" + ); + } + + /// Ensure that patches authored with ASCII characters can update lines that + /// contain typographic Unicode punctuation (e.g. EN DASH, NON-BREAKING + /// HYPHEN). Historically `git apply` succeeds in such scenarios but our + /// internal matcher failed requiring an exact byte-for-byte match. The + /// fuzzy-matching pass that normalises common punctuation should now bridge + /// the gap. + #[tokio::test] + async fn test_update_line_with_unicode_dash() { + let dir = tempdir().unwrap(); + let path = dir.path().join("unicode.py"); + + // Original line contains EN DASH (\u{2013}) and NON-BREAKING HYPHEN (\u{2011}). + let original = "import asyncio # local import \u{2013} avoids top\u{2011}level dep\n"; + std::fs::write(&path, original).unwrap(); + + // Patch uses plain ASCII dash / hyphen. + let patch = wrap_patch(&format!( + r#"*** Update File: {} +@@ +-import asyncio # local import - avoids top-level dep ++import asyncio # HELLO"#, + path.display() + )); + + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + apply_patch( + &patch, + &PathUri::from_host_native_path(dir.path()).expect("absolute test path"), + &mut stdout, + &mut stderr, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .unwrap(); + + // File should now contain the replaced comment. + let expected = "import asyncio # HELLO\n"; + let contents = std::fs::read_to_string(&path).unwrap(); + assert_eq!(contents, expected); + + // Ensure success summary lists the file as modified. + let stdout_str = String::from_utf8(stdout).unwrap(); + let expected_out = format!( + "Success. Updated the following files:\nM {}\n", + path.display() + ); + assert_eq!(stdout_str, expected_out); + + // No stderr expected. + assert_eq!(String::from_utf8(stderr).unwrap(), ""); + } + + #[cfg(unix)] + #[tokio::test] + async fn test_apply_patch_fails_on_write_error() { + use std::os::unix::fs::PermissionsExt; + + let dir = tempdir().unwrap(); + let locked_dir = dir.path().join("locked"); + fs::create_dir(&locked_dir).unwrap(); + fs::set_permissions(&locked_dir, fs::Permissions::from_mode(0o555)).unwrap(); + + let patch = wrap_patch("*** Add File: locked/new.txt\n+after"); + + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + let result = apply_patch( + &patch, + &PathUri::from_host_native_path(dir.path()).expect("absolute test path"), + &mut stdout, + &mut stderr, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await; + let failure = result.expect_err("write should fail"); + + fs::set_permissions(&locked_dir, fs::Permissions::from_mode(0o755)).unwrap(); + + assert!(!failure.delta().is_exact()); + } + + #[tokio::test] + async fn test_unreadable_destinations_return_inexact_delta() { + let dir = tempdir().unwrap(); + let path = dir.path().join("binary.dat"); + fs::write(dir.path().join("source.txt"), "before\n").unwrap(); + let cwd = PathUri::from_host_native_path(dir.path()).expect("absolute test path"); + + for patch in [ + wrap_patch("*** Add File: binary.dat\n+text"), + wrap_patch("*** Update File: source.txt\n*** Move to: binary.dat\n@@\n-before\n+after"), + ] { + fs::write(&path, [0xff, 0xfe, 0xfd]).unwrap(); + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + let delta = apply_patch( + &patch, + &cwd, + &mut stdout, + &mut stderr, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .unwrap(); + + assert!(!delta.is_exact()); + } + } + + #[cfg(unix)] + #[tokio::test] + async fn test_delete_symlink_returns_inexact_delta() { + use std::os::unix::fs::symlink; + + let dir = tempdir().unwrap(); + fs::write(dir.path().join("target.txt"), "target\n").unwrap(); + symlink(dir.path().join("target.txt"), dir.path().join("link.txt")).unwrap(); + let patch = wrap_patch("*** Delete File: link.txt"); + + let mut stdout = Vec::new(); + let mut stderr = Vec::new(); + let delta = apply_patch( + &patch, + &PathUri::from_host_native_path(dir.path()).expect("absolute test path"), + &mut stdout, + &mut stderr, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .unwrap(); + + assert!(!delta.is_exact()); + } +} diff --git a/codex-rs/apply-patch/src/main.rs b/codex-rs/apply-patch/src/main.rs new file mode 100644 index 0000000000000000000000000000000000000000..9d3ed03361f94f22049e1e2c8b52c39c2a0b541e --- /dev/null +++ b/codex-rs/apply-patch/src/main.rs @@ -0,0 +1,3 @@ +pub fn main() -> ! { + codex_apply_patch::main() +} diff --git a/codex-rs/apply-patch/src/parser.rs b/codex-rs/apply-patch/src/parser.rs new file mode 100644 index 0000000000000000000000000000000000000000..c400d075a684fda29c10269aef1958a27faf90aa --- /dev/null +++ b/codex-rs/apply-patch/src/parser.rs @@ -0,0 +1,682 @@ +//! This module is responsible for parsing & validating a patch into a list of "hunks". +//! (It does not attempt to actually check that the patch can be applied to the filesystem.) +//! +//! The official Lark grammar for the apply-patch format is: +//! +//! start: begin_patch environment_id? hunk+ end_patch +//! begin_patch: "*** Begin Patch" LF +//! environment_id: "*** Environment ID: " filename LF +//! end_patch: "*** End Patch" LF? +//! +//! hunk: add_hunk | delete_hunk | update_hunk +//! add_hunk: "*** Add File: " filename LF add_line+ +//! delete_hunk: "*** Delete File: " filename LF +//! update_hunk: "*** Update File: " filename LF change_move? change? +//! filename: /(.+)/ +//! add_line: "+" /(.+)/ LF -> line +//! +//! change_move: "*** Move to: " filename LF +//! change: (change_context | change_line)+ eof_line? +//! change_context: ("@@" | "@@ " /(.+)/) LF +//! change_line: ("+" | "-" | " ") /(.+)/ LF +//! eof_line: "*** End of File" LF +//! +//! The parser below is a little more lenient than the explicit spec and allows for +//! leading/trailing whitespace around patch markers. +use crate::ApplyPatchArgs; +use crate::streaming_parser::StreamingPatchParser; +#[cfg(test)] +use codex_utils_absolute_path::test_support::PathBufExt; +use codex_utils_path_uri::PathUri; +use codex_utils_path_uri::PathUriParseError; +use std::path::Path; +use std::path::PathBuf; + +use thiserror::Error; + +pub(crate) const BEGIN_PATCH_MARKER: &str = "*** Begin Patch"; +pub(crate) const END_PATCH_MARKER: &str = "*** End Patch"; +pub(crate) const ADD_FILE_MARKER: &str = "*** Add File: "; +pub(crate) const DELETE_FILE_MARKER: &str = "*** Delete File: "; +pub(crate) const UPDATE_FILE_MARKER: &str = "*** Update File: "; +pub(crate) const MOVE_TO_MARKER: &str = "*** Move to: "; +pub(crate) const EOF_MARKER: &str = "*** End of File"; +pub(crate) const CHANGE_CONTEXT_MARKER: &str = "@@ "; +pub(crate) const EMPTY_CHANGE_CONTEXT_MARKER: &str = "@@"; + +/// Currently, the only OpenAI model that knowingly requires lenient parsing is +/// gpt-4.1. While we could try to require everyone to pass in a strictness +/// param when invoking apply_patch, it is a pain to thread it through all of +/// the call sites, so we resign ourselves allowing lenient parsing for all +/// models. See [`ParseMode::Lenient`] for details on the exceptions we make for +/// gpt-4.1. +const PARSE_IN_STRICT_MODE: bool = false; + +#[derive(Debug, PartialEq, Error, Clone)] +pub enum ParseError { + #[error("invalid patch: {0}")] + InvalidPatchError(String), + #[error("invalid hunk at line {line_number}, {message}")] + InvalidHunkError { message: String, line_number: usize }, +} +use ParseError::*; + +#[derive(Debug, PartialEq, Clone)] +#[allow(clippy::enum_variant_names)] +pub enum Hunk { + AddFile { + path: PathBuf, + contents: String, + }, + DeleteFile { + path: PathBuf, + }, + UpdateFile { + path: PathBuf, + move_path: Option, + + /// Chunks should be in order, i.e. the `change_context` of one chunk + /// should occur later in the file than the previous chunk. + chunks: Vec, + }, +} + +impl Hunk { + pub fn resolve_path(&self, cwd: &PathUri) -> Result { + let path = match self { + Hunk::UpdateFile { path, .. } => path, + Hunk::AddFile { .. } | Hunk::DeleteFile { .. } => self.path(), + }; + cwd.join(&path.to_string_lossy()) + } + + /// Returns the path affected by this hunk, using the move destination for rename hunks. + pub fn path(&self) -> &Path { + match self { + Hunk::AddFile { path, .. } => path, + Hunk::DeleteFile { path } => path, + Hunk::UpdateFile { + move_path: Some(path), + .. + } => path, + Hunk::UpdateFile { + path, + move_path: None, + .. + } => path, + } + } +} + +#[cfg(test)] +use Hunk::*; + +#[derive(Debug, Default, PartialEq, Clone)] +pub struct UpdateFileChunk { + /// A single line of context used to narrow down the position of the chunk + /// (this is usually a class, method, or function definition.) + pub change_context: Option, + + /// A contiguous block of lines that should be replaced with `new_lines`. + /// `old_lines` must occur strictly after `change_context`. + pub old_lines: Vec, + pub new_lines: Vec, + + /// Pairs of indices into `old_lines` and `new_lines` that identify lines + /// parsed as context rather than inferred to be equal by their contents. + pub context_line_indices: Vec<(usize, usize)>, + + /// If set to true, `old_lines` must occur at the end of the source file. + /// (Tolerance around trailing newlines should be encouraged.) + pub is_end_of_file: bool, +} + +impl UpdateFileChunk { + /// Adds a context line to both sides while recording its corresponding + /// indices so it remains distinguishable from identical changed lines. + pub(crate) fn push_context_line(&mut self, line: String) { + self.context_line_indices + .push((self.old_lines.len(), self.new_lines.len())); + self.old_lines.push(line.clone()); + self.new_lines.push(line); + } +} + +pub fn parse_patch(patch: &str) -> Result { + let mode = if PARSE_IN_STRICT_MODE { + ParseMode::Strict + } else { + ParseMode::Lenient + }; + parse_patch_text(patch, mode) +} + +enum ParseMode { + /// Parse the patch text argument as is. + Strict, + + /// GPT-4.1 is known to formulate the `command` array for the `local_shell` + /// tool call for `apply_patch` call using something like the following: + /// + /// ```json + /// [ + /// "apply_patch", + /// "<<'EOF'\n*** Begin Patch\n*** Update File: README.md\n@@...\n*** End Patch\nEOF\n", + /// ] + /// ``` + /// + /// This is a problem because `local_shell` is a bit of a misnomer: the + /// `command` is not invoked by passing the arguments to a shell like Bash, + /// but are invoked using something akin to `execvpe(3)`. + /// + /// This is significant in this case because where a shell would interpret + /// `<<'EOF'...` as a heredoc and pass the contents via stdin (which is + /// fine, as `apply_patch` is specified to read from stdin if no argument is + /// passed), `execvpe(3)` interprets the heredoc as a literal string. To get + /// the `local_shell` tool to run a command the way shell would, the + /// `command` array must be something like: + /// + /// ```json + /// [ + /// "bash", + /// "-lc", + /// "apply_patch <<'EOF'\n*** Begin Patch\n*** Update File: README.md\n@@...\n*** End Patch\nEOF\n", + /// ] + /// ``` + /// + /// In lenient mode, we check if the argument to `apply_patch` starts with + /// `<<'EOF'` and ends with `EOF\n`. If so, we strip off these markers, + /// trim() the result, and treat what is left as the patch text. + Lenient, +} + +fn parse_patch_text(patch: &str, mode: ParseMode) -> Result { + let lines: Vec<&str> = patch.trim().lines().collect(); + let patch_lines = match mode { + ParseMode::Strict => check_patch_boundaries_strict(&lines)?, + ParseMode::Lenient => check_patch_boundaries_lenient(&lines)?, + }; + + let patch = patch_lines.join("\n"); + let mut parser = StreamingPatchParser::default(); + parser.push_delta(&patch)?; + let hunks = parser.finish()?; + let environment_id = parser.environment_id().map(str::to_owned); + Ok(ApplyPatchArgs { + hunks, + patch, + workdir: None, + environment_id, + }) +} + +/// Checks the start and end lines of the patch text for `apply_patch`, +/// returning an error if they do not match the expected markers. +fn check_patch_boundaries_strict<'a>(lines: &'a [&'a str]) -> Result<&'a [&'a str], ParseError> { + let (first_line, last_line) = match lines { + [] => (None, None), + [first] => (Some(first), Some(first)), + [first, .., last] => (Some(first), Some(last)), + }; + check_start_and_end_lines_strict(first_line, last_line)?; + Ok(lines) +} + +/// If we are in lenient mode, we check if the first line starts with `<( + original_lines: &'a [&'a str], +) -> Result<&'a [&'a str], ParseError> { + let original_parse_error = match check_patch_boundaries_strict(original_lines) { + Ok(lines) => return Ok(lines), + Err(e) => e, + }; + + match original_lines { + [first, .., last] => { + if (first == &"<= 4 + { + let inner_lines = &original_lines[1..original_lines.len() - 1]; + check_patch_boundaries_strict(inner_lines) + } else { + Err(original_parse_error) + } + } + _ => Err(original_parse_error), + } +} + +fn check_start_and_end_lines_strict( + first_line: Option<&&str>, + last_line: Option<&&str>, +) -> Result<(), ParseError> { + let first_line = first_line.map(|line| line.trim()); + let last_line = last_line.map(|line| line.trim()); + + match (first_line, last_line) { + (Some(first), Some(last)) if first == BEGIN_PATCH_MARKER && last == END_PATCH_MARKER => { + Ok(()) + } + (Some(first), _) if first != BEGIN_PATCH_MARKER => Err(InvalidPatchError(String::from( + "The first line of the patch must be '*** Begin Patch'", + ))), + _ => Err(InvalidPatchError(String::from( + "The last line of the patch must be '*** End Patch'", + ))), + } +} + +#[test] +fn test_parse_patch() { + assert_eq!( + parse_patch_text("bad", ParseMode::Strict), + Err(InvalidPatchError( + "The first line of the patch must be '*** Begin Patch'".to_string() + )) + ); + assert_eq!( + parse_patch_text("*** Begin Patch\nbad", ParseMode::Strict), + Err(InvalidPatchError( + "The last line of the patch must be '*** End Patch'".to_string() + )) + ); + + assert_eq!( + parse_patch_text( + concat!( + "*** Begin Patch", + " ", + "\n*** Add File: foo\n+hi\n", + " ", + "*** End Patch" + ), + ParseMode::Strict + ) + .unwrap() + .hunks, + vec![AddFile { + path: PathBuf::from("foo"), + contents: "hi\n".to_string() + }] + ); + assert_eq!( + parse_patch_text( + "*** Begin Patch\n\ + *** Update File: test.py\n\ + *** End Patch", + ParseMode::Strict + ), + Err(InvalidHunkError { + message: "Update file hunk for path 'test.py' is empty".to_string(), + line_number: 2, + }) + ); + assert_eq!( + parse_patch_text( + "*** Begin Patch\n\ + *** End Patch", + ParseMode::Strict + ) + .unwrap() + .hunks, + Vec::new() + ); + assert_eq!( + parse_patch_text( + "*** Begin Patch\n\ + *** Add File: path/add.py\n\ + +abc\n\ + +def\n\ + *** Delete File: path/delete.py\n\ + *** Update File: path/update.py\n\ + *** Move to: path/update2.py\n\ + @@ def f():\n\ + - pass\n\ + + return 123\n\ + *** End Patch", + ParseMode::Strict + ) + .unwrap() + .hunks, + vec![ + AddFile { + path: PathBuf::from("path/add.py"), + contents: "abc\ndef\n".to_string() + }, + DeleteFile { + path: PathBuf::from("path/delete.py") + }, + UpdateFile { + path: PathBuf::from("path/update.py"), + move_path: Some(PathBuf::from("path/update2.py")), + chunks: vec![UpdateFileChunk { + change_context: Some("def f():".to_string()), + old_lines: vec![" pass".to_string()], + new_lines: vec![" return 123".to_string()], + context_line_indices: vec![], + is_end_of_file: false + }] + } + ] + ); + // Update hunk followed by another hunk (Add File). + assert_eq!( + parse_patch_text( + "*** Begin Patch\n\ + *** Update File: file.py\n\ + @@\n\ + +line\n\ + *** Add File: other.py\n\ + +content\n\ + *** End Patch", + ParseMode::Strict + ) + .unwrap() + .hunks, + vec![ + UpdateFile { + path: PathBuf::from("file.py"), + move_path: None, + chunks: vec![UpdateFileChunk { + change_context: None, + old_lines: vec![], + new_lines: vec!["line".to_string()], + context_line_indices: vec![], + is_end_of_file: false + }], + }, + AddFile { + path: PathBuf::from("other.py"), + contents: "content\n".to_string() + } + ] + ); + + // Update hunk without an explicit @@ header for the first chunk should parse. + // Use a raw string to preserve the leading space diff marker on the context line. + assert_eq!( + parse_patch_text( + r#"*** Begin Patch +*** Update File: file2.py + import foo ++bar +*** End Patch"#, + ParseMode::Strict + ) + .unwrap() + .hunks, + vec![UpdateFile { + path: PathBuf::from("file2.py"), + move_path: None, + chunks: vec![UpdateFileChunk { + change_context: None, + old_lines: vec!["import foo".to_string()], + new_lines: vec!["import foo".to_string(), "bar".to_string()], + context_line_indices: vec![(0, 0)], + is_end_of_file: false, + }], + }] + ); +} + +#[test] +fn test_parse_patch_preserves_end_of_file_marker() { + let patch = + "*** Begin Patch\n*** Update File: file.txt\n@@\n+quux\n*** End of File\n\n*** End Patch"; + assert_eq!( + parse_patch(patch), + Ok(ApplyPatchArgs { + hunks: vec![UpdateFile { + path: PathBuf::from("file.txt"), + move_path: None, + chunks: vec![UpdateFileChunk { + change_context: None, + old_lines: Vec::new(), + new_lines: vec!["quux".to_string()], + context_line_indices: vec![], + is_end_of_file: true, + }], + }], + patch: patch.to_string(), + workdir: None, + environment_id: None, + }) + ); +} + +#[test] +fn test_parse_patch_accepts_relative_and_absolute_hunk_paths() { + let dir = tempfile::tempdir().unwrap(); + let absolute_delete = dir.path().join("absolute-delete.py").abs(); + let absolute_update = dir.path().join("absolute-update.py").abs(); + let patch_text = format!( + r#"*** Begin Patch +*** Add File: relative-add.py ++content +*** Delete File: {} +*** Update File: {} +@@ +-old ++new +*** End Patch"#, + absolute_delete.display(), + absolute_update.display() + ); + + assert_eq!( + parse_patch_text(&patch_text, ParseMode::Strict) + .unwrap() + .hunks, + vec![ + AddFile { + path: PathBuf::from("relative-add.py"), + contents: "content\n".to_string() + }, + DeleteFile { + path: absolute_delete.to_path_buf() + }, + UpdateFile { + path: absolute_update.to_path_buf(), + move_path: None, + chunks: vec![UpdateFileChunk { + change_context: None, + old_lines: vec!["old".to_string()], + new_lines: vec!["new".to_string()], + context_line_indices: vec![], + is_end_of_file: false + }] + }, + ] + ); +} + +#[test] +fn test_hunk_resolve_path_accepts_relative_and_absolute_paths() { + let cwd_dir = tempfile::tempdir().unwrap(); + let cwd = PathUri::from_host_native_path(cwd_dir.path()).unwrap(); + let absolute_dir = tempfile::tempdir().unwrap(); + let absolute_add = absolute_dir.path().join("absolute-add.py").abs(); + let absolute_delete = absolute_dir.path().join("absolute-delete.py").abs(); + let absolute_update = absolute_dir.path().join("absolute-update.py").abs(); + + for (hunk, expected_path) in [ + ( + AddFile { + path: PathBuf::from("relative-add.py"), + contents: String::new(), + }, + cwd.join("relative-add.py").unwrap(), + ), + ( + DeleteFile { + path: PathBuf::from("relative-delete.py"), + }, + cwd.join("relative-delete.py").unwrap(), + ), + ( + UpdateFile { + path: PathBuf::from("relative-update.py"), + move_path: None, + chunks: Vec::new(), + }, + cwd.join("relative-update.py").unwrap(), + ), + ( + AddFile { + path: absolute_add.to_path_buf(), + contents: String::new(), + }, + PathUri::from_abs_path(&absolute_add), + ), + ( + DeleteFile { + path: absolute_delete.to_path_buf(), + }, + PathUri::from_abs_path(&absolute_delete), + ), + ( + UpdateFile { + path: absolute_update.to_path_buf(), + move_path: None, + chunks: Vec::new(), + }, + PathUri::from_abs_path(&absolute_update), + ), + ] { + assert_eq!(hunk.resolve_path(&cwd), Ok(expected_path)); + } +} + +#[test] +fn test_parse_patch_lenient() { + let patch_text = r#"*** Begin Patch +*** Update File: file2.py + import foo ++bar +*** End Patch"#; + let expected_patch = vec![UpdateFile { + path: PathBuf::from("file2.py"), + move_path: None, + chunks: vec![UpdateFileChunk { + change_context: None, + old_lines: vec!["import foo".to_string()], + new_lines: vec!["import foo".to_string(), "bar".to_string()], + context_line_indices: vec![(0, 0)], + is_end_of_file: false, + }], + }]; + let expected_error = + InvalidPatchError("The first line of the patch must be '*** Begin Patch'".to_string()); + + let patch_text_in_heredoc = format!("< lines.len()` → returns `None` (cannot match, avoids +/// out‑of‑bounds panic that occurred pre‑2025‑04‑12) +pub(crate) fn seek_sequence( + lines: &[String], + pattern: &[String], + start: usize, + eof: bool, + update_file_mode: crate::ApplyPatchFileUpdateMode, +) -> Option { + if pattern.is_empty() { + return Some(start); + } + + // When the pattern is longer than the available input there is no possible + // match. Early‑return to avoid the out‑of‑bounds slice that would occur in + // the search loops below (previously caused a panic when + // `pattern.len() > lines.len()`). + if pattern.len() > lines.len() { + return None; + } + let search_start = if eof && lines.len() >= pattern.len() { + let eof_start = lines.len() - pattern.len(); + match update_file_mode { + crate::ApplyPatchFileUpdateMode::NormalizeToLf => eof_start, + crate::ApplyPatchFileUpdateMode::PreserveLineEndings => eof_start.max(start), + } + } else { + start + }; + // Exact match first. + for i in search_start..=lines.len().saturating_sub(pattern.len()) { + if lines[i..i + pattern.len()] == *pattern { + return Some(i); + } + } + // Then rstrip match. + for i in search_start..=lines.len().saturating_sub(pattern.len()) { + let mut ok = true; + for (p_idx, pat) in pattern.iter().enumerate() { + if lines[i + p_idx].trim_end() != pat.trim_end() { + ok = false; + break; + } + } + if ok { + return Some(i); + } + } + // Finally, trim both sides to allow more lenience. + for i in search_start..=lines.len().saturating_sub(pattern.len()) { + let mut ok = true; + for (p_idx, pat) in pattern.iter().enumerate() { + if lines[i + p_idx].trim() != pat.trim() { + ok = false; + break; + } + } + if ok { + return Some(i); + } + } + + // ------------------------------------------------------------------ + // Final, most permissive pass – attempt to match after *normalising* + // common Unicode punctuation to their ASCII equivalents so that diffs + // authored with plain ASCII characters can still be applied to source + // files that contain typographic dashes / quotes, etc. This mirrors the + // fuzzy behaviour of `git apply` which ignores minor byte-level + // differences when locating context lines. + // ------------------------------------------------------------------ + + fn normalise(s: &str) -> String { + s.trim() + .chars() + .map(|c| match c { + // Various dash / hyphen code-points → ASCII '-' + '\u{2010}' | '\u{2011}' | '\u{2012}' | '\u{2013}' | '\u{2014}' | '\u{2015}' + | '\u{2212}' => '-', + // Fancy single quotes → '\'' + '\u{2018}' | '\u{2019}' | '\u{201A}' | '\u{201B}' => '\'', + // Fancy double quotes → '"' + '\u{201C}' | '\u{201D}' | '\u{201E}' | '\u{201F}' => '"', + // Non-breaking space and other odd spaces → normal space + '\u{00A0}' | '\u{2002}' | '\u{2003}' | '\u{2004}' | '\u{2005}' | '\u{2006}' + | '\u{2007}' | '\u{2008}' | '\u{2009}' | '\u{200A}' | '\u{202F}' | '\u{205F}' + | '\u{3000}' => ' ', + other => other, + }) + .collect::() + } + + for i in search_start..=lines.len().saturating_sub(pattern.len()) { + let mut ok = true; + for (p_idx, pat) in pattern.iter().enumerate() { + if normalise(&lines[i + p_idx]) != normalise(pat) { + ok = false; + break; + } + } + if ok { + return Some(i); + } + } + + None +} + +#[cfg(test)] +mod tests { + use super::seek_sequence; + use crate::ApplyPatchFileUpdateMode; + use std::string::ToString; + + fn to_vec(strings: &[&str]) -> Vec { + strings.iter().map(ToString::to_string).collect() + } + + #[test] + fn test_exact_match_finds_sequence() { + let lines = to_vec(&["foo", "bar", "baz"]); + let pattern = to_vec(&["bar", "baz"]); + assert_eq!( + seek_sequence( + &lines, + &pattern, + /*start*/ 0, + /*eof*/ false, + ApplyPatchFileUpdateMode::NormalizeToLf, + ), + Some(1) + ); + } + + #[test] + fn test_rstrip_match_ignores_trailing_whitespace() { + let lines = to_vec(&["foo ", "bar\t\t"]); + // Pattern omits trailing whitespace. + let pattern = to_vec(&["foo", "bar"]); + assert_eq!( + seek_sequence( + &lines, + &pattern, + /*start*/ 0, + /*eof*/ false, + ApplyPatchFileUpdateMode::NormalizeToLf, + ), + Some(0) + ); + } + + #[test] + fn test_trim_match_ignores_leading_and_trailing_whitespace() { + let lines = to_vec(&[" foo ", " bar\t"]); + // Pattern omits any additional whitespace. + let pattern = to_vec(&["foo", "bar"]); + assert_eq!( + seek_sequence( + &lines, + &pattern, + /*start*/ 0, + /*eof*/ false, + ApplyPatchFileUpdateMode::NormalizeToLf, + ), + Some(0) + ); + } + + #[test] + fn test_pattern_longer_than_input_returns_none() { + let lines = to_vec(&["just one line"]); + let pattern = to_vec(&["too", "many", "lines"]); + // Should not panic – must return None when pattern cannot possibly fit. + assert_eq!( + seek_sequence( + &lines, + &pattern, + /*start*/ 0, + /*eof*/ false, + ApplyPatchFileUpdateMode::NormalizeToLf, + ), + None + ); + } +} diff --git a/codex-rs/apply-patch/src/standalone_executable.rs b/codex-rs/apply-patch/src/standalone_executable.rs new file mode 100644 index 0000000000000000000000000000000000000000..7f2627beca0c2d0c18c09ab02988f0d872a74adc --- /dev/null +++ b/codex-rs/apply-patch/src/standalone_executable.rs @@ -0,0 +1,90 @@ +use std::io::Read; +use std::io::Write; + +pub fn main() -> ! { + let exit_code = run_main(); + std::process::exit(exit_code); +} + +/// We would prefer to return `std::process::ExitCode`, but its `exit_process()` +/// method is still a nightly API and we want main() to return !. +pub fn run_main() -> i32 { + // Expect either one argument (the full apply_patch payload) or read it from stdin. + let mut args = std::env::args_os(); + let _argv0 = args.next(); + + let patch_arg = match args.next() { + Some(arg) => match arg.into_string() { + Ok(s) => s, + Err(_) => { + eprintln!("Error: apply_patch requires a UTF-8 PATCH argument."); + return 1; + } + }, + None => { + // No argument provided; attempt to read the patch from stdin. + let mut buf = String::new(); + match std::io::stdin().read_to_string(&mut buf) { + Ok(_) => { + if buf.is_empty() { + eprintln!("Usage: apply_patch 'PATCH'\n echo 'PATCH' | apply_patch"); + return 2; + } + buf + } + Err(err) => { + eprintln!("Error: Failed to read PATCH from stdin.\n{err}"); + return 1; + } + } + } + }; + + // Refuse extra args to avoid ambiguity. + if args.next().is_some() { + eprintln!("Error: apply_patch accepts exactly one argument."); + return 2; + } + + let mut stdout = std::io::stdout(); + let mut stderr = std::io::stderr(); + let cwd = match codex_utils_absolute_path::AbsolutePathBuf::current_dir() { + Ok(cwd) => cwd, + Err(err) => { + eprintln!("Error: Failed to determine current directory.\n{err}"); + return 1; + } + }; + let runtime = match tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + { + Ok(runtime) => runtime, + Err(err) => { + eprintln!("Error: Failed to initialize runtime.\n{err}"); + return 1; + } + }; + // TODO(anp): Discover the standalone executable cwd as PathUri directly. + let cwd = codex_utils_path_uri::PathUri::from_abs_path(&cwd); + let update_file_mode = crate::apply_patch_file_update_mode_from_env(); + match runtime.block_on(crate::apply_patch_with_options( + &patch_arg, + crate::ApplyPatchOptions { + update_file_mode, + ..Default::default() + }, + &cwd, + &mut stdout, + &mut stderr, + codex_exec_server::LOCAL_FS.as_ref(), + /*sandbox*/ None, + )) { + Ok(_) => { + // Flush to ensure output ordering when used in pipelines. + let _ = stdout.flush(); + 0 + } + Err(_) => 1, + } +} diff --git a/codex-rs/apply-patch/src/streaming_parser.rs b/codex-rs/apply-patch/src/streaming_parser.rs new file mode 100644 index 0000000000000000000000000000000000000000..ff1b2f82feec7bdf454f6cba61ad5f15611f5573 --- /dev/null +++ b/codex-rs/apply-patch/src/streaming_parser.rs @@ -0,0 +1,924 @@ +use std::path::PathBuf; + +use crate::parser::ADD_FILE_MARKER; +use crate::parser::BEGIN_PATCH_MARKER; +use crate::parser::CHANGE_CONTEXT_MARKER; +use crate::parser::DELETE_FILE_MARKER; +use crate::parser::EMPTY_CHANGE_CONTEXT_MARKER; +use crate::parser::END_PATCH_MARKER; +use crate::parser::EOF_MARKER; +use crate::parser::Hunk; +use crate::parser::MOVE_TO_MARKER; +use crate::parser::ParseError; +use crate::parser::UPDATE_FILE_MARKER; +use crate::parser::UpdateFileChunk; + +use Hunk::*; +use ParseError::*; + +const ENVIRONMENT_ID_MARKER: &str = "*** Environment ID:"; + +#[derive(Debug, Default, Clone)] +pub struct StreamingPatchParser { + line_buffer: String, + state: StreamingParserState, + line_number: usize, +} + +#[derive(Debug, Default, Clone)] +struct StreamingParserState { + mode: StreamingParserMode, + hunks: Vec, + environment_id: Option, +} + +#[derive(Debug, Default, Clone, Copy)] +enum StreamingParserMode { + #[default] + NotStarted, + StartedPatch, + AddFile, + DeleteFile, + UpdateFile { + hunk_line_number: usize, + }, + EndedPatch, +} + +impl StreamingPatchParser { + pub fn environment_id(&self) -> Option<&str> { + self.state.environment_id.as_deref() + } + + fn ensure_update_hunk_is_not_empty(&self, line: &str) -> Result<(), ParseError> { + if let Some(UpdateFile { path, chunks, .. }) = self.state.hunks.last() { + if chunks.is_empty() + && let StreamingParserMode::UpdateFile { hunk_line_number } = self.state.mode + { + return Err(InvalidHunkError { + message: format!("Update file hunk for path '{}' is empty", path.display()), + line_number: hunk_line_number, + }); + } + if chunks + .last() + .is_some_and(|chunk| chunk.old_lines.is_empty() && chunk.new_lines.is_empty()) + { + if line == END_PATCH_MARKER { + return Err(InvalidHunkError { + message: "Update hunk does not contain any lines".to_string(), + line_number: self.line_number, + }); + } + return Err(InvalidHunkError { + message: format!( + "Unexpected line found in update hunk: '{line}'. Every line should start with ' ' (context line), '+' (added line), or '-' (removed line)" + ), + line_number: self.line_number, + }); + } + } + Ok(()) + } + + fn handle_hunk_headers_and_end_patch(&mut self, trimmed: &str) -> Result { + if matches!(self.state.mode, StreamingParserMode::StartedPatch) + && let Some(environment_id) = trimmed.strip_prefix(ENVIRONMENT_ID_MARKER) + { + if self.state.environment_id.is_some() { + return Err(InvalidPatchError( + "apply_patch environment_id cannot be specified more than once".to_string(), + )); + } + let environment_id = environment_id.trim(); + if environment_id.is_empty() { + return Err(InvalidPatchError( + "apply_patch environment_id cannot be empty".to_string(), + )); + } + self.state.environment_id = Some(environment_id.to_string()); + return Ok(true); + } + if trimmed == END_PATCH_MARKER { + self.ensure_update_hunk_is_not_empty(trimmed)?; + self.state.mode = StreamingParserMode::EndedPatch; + return Ok(true); + } + if let Some(path) = trimmed.strip_prefix(ADD_FILE_MARKER) { + self.ensure_update_hunk_is_not_empty(trimmed)?; + self.state.hunks.push(AddFile { + path: PathBuf::from(path), + contents: String::new(), + }); + self.state.mode = StreamingParserMode::AddFile; + return Ok(true); + } + if let Some(path) = trimmed.strip_prefix(DELETE_FILE_MARKER) { + self.ensure_update_hunk_is_not_empty(trimmed)?; + self.state.hunks.push(DeleteFile { + path: PathBuf::from(path), + }); + self.state.mode = StreamingParserMode::DeleteFile; + return Ok(true); + } + if let Some(path) = trimmed.strip_prefix(UPDATE_FILE_MARKER) { + self.ensure_update_hunk_is_not_empty(trimmed)?; + self.state.hunks.push(UpdateFile { + path: PathBuf::from(path), + move_path: None, + chunks: Vec::new(), + }); + self.state.mode = StreamingParserMode::UpdateFile { + hunk_line_number: self.line_number, + }; + return Ok(true); + } + Ok(false) + } + + pub fn push_delta(&mut self, delta: &str) -> Result, ParseError> { + for ch in delta.chars() { + if ch == '\n' { + let mut line = std::mem::take(&mut self.line_buffer); + line.truncate(line.strip_suffix('\r').map_or(line.len(), str::len)); + self.line_number += 1; + self.process_line(&line)?; + } else { + self.line_buffer.push(ch); + } + } + + Ok(self.state.hunks.clone()) + } + + pub fn finish(&mut self) -> Result, ParseError> { + if !self.line_buffer.is_empty() { + let line = std::mem::take(&mut self.line_buffer); + self.line_number += 1; + if line.trim() == END_PATCH_MARKER { + self.ensure_update_hunk_is_not_empty(line.trim())?; + self.state.mode = StreamingParserMode::EndedPatch; + } else { + self.process_line(&line)?; + } + } + + if !matches!(self.state.mode, StreamingParserMode::EndedPatch) { + return Err(InvalidPatchError( + "The last line of the patch must be '*** End Patch'".to_string(), + )); + } + + Ok(self.state.hunks.clone()) + } + + fn process_line(&mut self, line: &str) -> Result<(), ParseError> { + let trimmed = line.trim(); + match self.state.mode { + StreamingParserMode::NotStarted => { + if trimmed == BEGIN_PATCH_MARKER { + self.state.mode = StreamingParserMode::StartedPatch; + return Ok(()); + } + Err(InvalidPatchError( + "The first line of the patch must be '*** Begin Patch'".to_string(), + )) + } + StreamingParserMode::StartedPatch => { + if self.handle_hunk_headers_and_end_patch(trimmed)? { + return Ok(()); + } + Err(InvalidHunkError { + message: format!( + "'{trimmed}' is not a valid hunk header. Valid hunk headers: '*** Add File: {{path}}', '*** Delete File: {{path}}', '*** Update File: {{path}}'" + ), + line_number: self.line_number, + }) + } + StreamingParserMode::AddFile => { + if self.handle_hunk_headers_and_end_patch(trimmed)? { + return Ok(()); + } + if let Some(line_to_add) = line.strip_prefix('+') + && let Some(AddFile { contents, .. }) = self.state.hunks.last_mut() + { + contents.push_str(line_to_add); + contents.push('\n'); + return Ok(()); + } + Err(InvalidHunkError { + message: format!( + "'{trimmed}' is not a valid hunk header. Valid hunk headers: '*** Add File: {{path}}', '*** Delete File: {{path}}', '*** Update File: {{path}}'" + ), + line_number: self.line_number, + }) + } + StreamingParserMode::DeleteFile => { + if self.handle_hunk_headers_and_end_patch(trimmed)? { + return Ok(()); + } + Err(InvalidHunkError { + message: format!( + "'{trimmed}' is not a valid hunk header. Valid hunk headers: '*** Add File: {{path}}', '*** Delete File: {{path}}', '*** Update File: {{path}}'" + ), + line_number: self.line_number, + }) + } + StreamingParserMode::UpdateFile { hunk_line_number } => { + let update_line = line.trim_end(); + if self.handle_hunk_headers_and_end_patch(update_line)? { + return Ok(()); + } + + if let Some(UpdateFile { + move_path, chunks, .. + }) = self.state.hunks.last_mut() + { + if chunks.last().is_some_and(|chunk| chunk.is_end_of_file) { + if update_line.is_empty() { + return Ok(()); + } + if update_line != EMPTY_CHANGE_CONTEXT_MARKER + && !update_line.starts_with(CHANGE_CONTEXT_MARKER) + { + return Err(InvalidHunkError { + message: format!( + "Expected update hunk to start with a @@ context marker, got: '{line}'" + ), + line_number: self.line_number, + }); + } + } + + if chunks.is_empty() + && move_path.is_none() + && let Some(move_to_path) = update_line.strip_prefix(MOVE_TO_MARKER) + { + *move_path = Some(PathBuf::from(move_to_path)); + self.state.mode = StreamingParserMode::UpdateFile { hunk_line_number }; + return Ok(()); + } + + if (update_line == EMPTY_CHANGE_CONTEXT_MARKER + || update_line.starts_with(CHANGE_CONTEXT_MARKER)) + && chunks.last().is_some_and(|chunk| { + chunk.old_lines.is_empty() && chunk.new_lines.is_empty() + }) + { + return Err(InvalidHunkError { + message: format!( + "Unexpected line found in update hunk: '{line}'. Every line should start with ' ' (context line), '+' (added line), or '-' (removed line)" + ), + line_number: self.line_number, + }); + } + + if update_line == EMPTY_CHANGE_CONTEXT_MARKER { + chunks.push(UpdateFileChunk::default()); + self.state.mode = StreamingParserMode::UpdateFile { hunk_line_number }; + return Ok(()); + } + + if let Some(change_context) = update_line.strip_prefix(CHANGE_CONTEXT_MARKER) { + chunks.push(UpdateFileChunk { + change_context: Some(change_context.to_string()), + ..UpdateFileChunk::default() + }); + self.state.mode = StreamingParserMode::UpdateFile { hunk_line_number }; + return Ok(()); + } + + if update_line == EOF_MARKER { + if chunks.last().is_some_and(|chunk| { + chunk.old_lines.is_empty() && chunk.new_lines.is_empty() + }) { + return Err(InvalidHunkError { + message: "Update hunk does not contain any lines".to_string(), + line_number: self.line_number, + }); + } + if let Some(chunk) = chunks.last_mut() { + chunk.is_end_of_file = true; + } + self.state.mode = StreamingParserMode::UpdateFile { hunk_line_number }; + return Ok(()); + } + + if line.is_empty() { + if chunks.is_empty() { + chunks.push(UpdateFileChunk::default()); + } + if let Some(chunk) = chunks.last_mut() { + chunk.push_context_line(String::new()); + } + self.state.mode = StreamingParserMode::UpdateFile { hunk_line_number }; + return Ok(()); + } + + if let Some(line_to_add) = line.strip_prefix(' ') { + if chunks.is_empty() { + chunks.push(UpdateFileChunk::default()); + } + if let Some(chunk) = chunks.last_mut() { + chunk.push_context_line(line_to_add.to_string()); + } + self.state.mode = StreamingParserMode::UpdateFile { hunk_line_number }; + return Ok(()); + } + + if let Some(line_to_add) = line.strip_prefix('+') { + if chunks.is_empty() { + chunks.push(UpdateFileChunk::default()); + } + if let Some(chunk) = chunks.last_mut() { + chunk.new_lines.push(line_to_add.to_string()); + } + self.state.mode = StreamingParserMode::UpdateFile { hunk_line_number }; + return Ok(()); + } + + if let Some(line_to_remove) = line.strip_prefix('-') { + if chunks.is_empty() { + chunks.push(UpdateFileChunk::default()); + } + if let Some(chunk) = chunks.last_mut() { + chunk.old_lines.push(line_to_remove.to_string()); + } + self.state.mode = StreamingParserMode::UpdateFile { hunk_line_number }; + return Ok(()); + } + + if chunks.last().is_some_and(|chunk| { + !chunk.old_lines.is_empty() || !chunk.new_lines.is_empty() + }) { + return Err(InvalidHunkError { + message: format!( + "Expected update hunk to start with a @@ context marker, got: '{line}'" + ), + line_number: self.line_number, + }); + } + } + Err(InvalidHunkError { + message: format!( + "Unexpected line found in update hunk: '{line}'. Every line should start with ' ' (context line), '+' (added line), or '-' (removed line)" + ), + line_number: self.line_number, + }) + } + StreamingParserMode::EndedPatch => { + if trimmed.is_empty() { + Ok(()) + } else { + Err(InvalidPatchError( + "The last line of the patch must be '*** End Patch'".to_string(), + )) + } + } + } + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + use std::path::PathBuf; + + use super::*; + + #[test] + fn test_streaming_patch_parser_streams_complete_lines_before_end_patch() { + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta("*** Begin Patch\n*** Add File: src/hello.txt\n+hello\n+wor"), + Ok(vec![AddFile { + path: PathBuf::from("src/hello.txt"), + contents: "hello\n".to_string(), + }]) + ); + assert_eq!( + parser.push_delta("ld\n"), + Ok(vec![AddFile { + path: PathBuf::from("src/hello.txt"), + contents: "hello\nworld\n".to_string(), + }]) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta( + "*** Begin Patch\n*** Update File: src/old.rs\n*** Move to: src/new.rs\n@@\n-old\n+new\n", + ), + Ok(vec![UpdateFile { + path: PathBuf::from("src/old.rs"), + move_path: Some(PathBuf::from("src/new.rs")), + chunks: vec![UpdateFileChunk { + change_context: None, + old_lines: vec!["old".to_string()], + new_lines: vec!["new".to_string()], + context_line_indices: vec![], + is_end_of_file: false, + }], + }]) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta("*** Begin Patch\n*** Delete File: gone.txt"), + Ok(Vec::new()) + ); + assert_eq!( + parser.push_delta("\n"), + Ok(vec![DeleteFile { + path: PathBuf::from("gone.txt"), + }]) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta( + "*** Begin Patch\n*** Add File: src/one.txt\n+one\n*** Delete File: src/two.txt\n", + ), + Ok(vec![ + AddFile { + path: PathBuf::from("src/one.txt"), + contents: "one\n".to_string(), + }, + DeleteFile { + path: PathBuf::from("src/two.txt"), + }, + ]) + ); + } + + #[test] + fn test_streaming_patch_parser_environment_id_mode() { + let patch = "\ +*** Begin Patch +*** Environment ID: remote +*** Add File: src/hello.txt ++hello +*** End Patch +"; + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta(patch), + Ok(vec![AddFile { + path: PathBuf::from("src/hello.txt"), + contents: "hello\n".to_string(), + }]) + ); + assert_eq!(parser.environment_id(), Some("remote")); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta( + "*** Begin Patch\n*** Environment ID: first\n*** Environment ID: second\n", + ), + Err(InvalidPatchError( + "apply_patch environment_id cannot be specified more than once".to_string(), + )) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta("*** Begin Patch\n*** Environment ID: \n"), + Err(InvalidPatchError( + "apply_patch environment_id cannot be empty".to_string(), + )) + ); + } + + #[test] + fn test_streaming_patch_parser_large_patch_split_by_character() { + let patch = "\ +*** Begin Patch +*** Add File: docs/release-notes.md ++# Release notes ++ ++## CLI ++- Surface apply_patch progress while arguments stream. ++- Keep final patch application gated on the completed tool call. ++- Include file summaries in the progress event payload. +*** Update File: src/config.rs +@@ impl Config +- pub apply_patch_progress: bool, ++ pub stream_apply_patch_progress: bool, + pub include_diagnostics: bool, +@@ fn default_progress_interval() +- Duration::from_millis(500) ++ Duration::from_millis(250) +*** Delete File: src/legacy_patch_progress.rs +*** Update File: crates/cli/src/main.rs +*** Move to: crates/cli/src/bin/codex.rs +@@ fn run() +- let args = Args::parse(); +- dispatch(args) ++ let cli = Cli::parse(); ++ dispatch(cli) +*** Add File: tests/fixtures/apply_patch_progress.json ++{ ++ \"type\": \"apply_patch_progress\", ++ \"hunks\": [ ++ { \"operation\": \"add\", \"path\": \"docs/release-notes.md\" }, ++ { \"operation\": \"update\", \"path\": \"src/config.rs\" } ++ ] ++} +*** Update File: README.md +@@ Development workflow + Build the Rust workspace before opening a pull request. ++When touching streamed tool calls, include parser coverage for partial input. ++Prefer tests that exercise the exact event payload shape. +*** Delete File: docs/old-apply-patch-progress.md +*** End Patch"; + + let mut parser = StreamingPatchParser::default(); + let mut max_hunk_count = 0; + let mut saw_hunk_counts = Vec::new(); + let mut hunks = Vec::new(); + for ch in patch.chars() { + let updated_hunks = parser.push_delta(&ch.to_string()).unwrap(); + if !updated_hunks.is_empty() { + let hunk_count = updated_hunks.len(); + assert!( + hunk_count >= max_hunk_count, + "hunk count should never decrease while streaming: {hunk_count} < {max_hunk_count}", + ); + if hunk_count > max_hunk_count { + saw_hunk_counts.push(hunk_count); + max_hunk_count = hunk_count; + } + hunks = updated_hunks; + } + } + + assert_eq!(saw_hunk_counts, vec![1, 2, 3, 4, 5, 6, 7]); + assert_eq!(hunks.len(), 7); + assert_eq!( + hunks + .iter() + .map(|hunk| match hunk { + AddFile { .. } => "add", + DeleteFile { .. } => "delete", + UpdateFile { + move_path: Some(_), .. + } => "move-update", + UpdateFile { + move_path: None, .. + } => "update", + }) + .collect::>(), + vec![ + "add", + "update", + "delete", + "move-update", + "add", + "update", + "delete" + ] + ); + } + + #[test] + fn test_streaming_patch_parser_keeps_indented_update_markers_as_context_lines() { + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta( + "\ +*** Begin Patch +*** Update File: a.txt +@@ +-old a ++new a + *** Update File: b.txt +@@ +-old b ++new b +*** End Patch +", + ), + Ok(vec![UpdateFile { + path: PathBuf::from("a.txt"), + move_path: None, + chunks: vec![ + UpdateFileChunk { + change_context: None, + old_lines: vec!["old a".to_string(), "*** Update File: b.txt".to_string()], + new_lines: vec!["new a".to_string(), "*** Update File: b.txt".to_string()], + context_line_indices: vec![(1, 1)], + is_end_of_file: false, + }, + UpdateFileChunk { + change_context: None, + old_lines: vec!["old b".to_string()], + new_lines: vec!["new b".to_string()], + context_line_indices: vec![], + is_end_of_file: false, + }, + ], + }]) + ); + } + + #[test] + fn test_streaming_patch_parser_preserves_bare_empty_update_lines() { + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta( + "\ +*** Begin Patch +*** Update File: file.txt +@@ + context before + + context after +*** End Patch +", + ), + Ok(vec![UpdateFile { + path: PathBuf::from("file.txt"), + move_path: None, + chunks: vec![UpdateFileChunk { + change_context: None, + // The normal parser treats a bare empty line in an update hunk as an + // empty context line. Preserve that leniency in the streaming parser. + old_lines: vec![ + "context before".to_string(), + String::new(), + "context after".to_string(), + ], + new_lines: vec![ + "context before".to_string(), + String::new(), + "context after".to_string(), + ], + context_line_indices: vec![(0, 0), (1, 1), (2, 2)], + is_end_of_file: false, + }], + }]) + ); + } + + #[test] + fn test_streaming_patch_parser_ignores_empty_lines_after_end_of_file() { + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta( + "*** Begin Patch\n*** Update File: file.txt\n@@\n+quux\n*** End of File\n\n*** End Patch\n", + ), + Ok(vec![UpdateFile { + path: PathBuf::from("file.txt"), + move_path: None, + chunks: vec![UpdateFileChunk { + change_context: None, + old_lines: Vec::new(), + new_lines: vec!["quux".to_string()], + context_line_indices: vec![], + is_end_of_file: true, + }], + }]) + ); + } + + #[test] + fn test_streaming_patch_parser_matches_line_ending_behavior() { + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta("*** Begin Patch\r\n*** Update File: file.txt\r\n@@\r\n-old\r\n+new\r\n*** End Patch\r\n"), + Ok(vec![UpdateFile { + path: PathBuf::from("file.txt"), + move_path: None, + chunks: vec![UpdateFileChunk { + change_context: None, + old_lines: vec!["old".to_string()], + new_lines: vec!["new".to_string()], + context_line_indices: vec![], + is_end_of_file: false, + }], + }]) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta("*** Begin Patch\r\n*** Update File: file.txt\r\n@@\r\n-old\r\r\n+new\r\n*** End Patch\r\n"), + Ok(vec![UpdateFile { + path: PathBuf::from("file.txt"), + move_path: None, + chunks: vec![UpdateFileChunk { + change_context: None, + old_lines: vec!["old\r".to_string()], + new_lines: vec!["new".to_string()], + context_line_indices: vec![], + is_end_of_file: false, + }], + }]) + ); + } + + #[test] + fn test_streaming_patch_parser_finish_processes_final_line_without_newline() { + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta("*** Begin Patch\n*** Add File: file.txt\n+hello\n*** End Patch"), + Ok(vec![AddFile { + path: PathBuf::from("file.txt"), + contents: "hello\n".to_string(), + }]) + ); + assert_eq!( + parser.finish(), + Ok(vec![AddFile { + path: PathBuf::from("file.txt"), + contents: "hello\n".to_string(), + }]) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta( + "*** Begin Patch\n*** Update File: file.txt\n@@\n-old\n+new\n *** End Patch", + ), + Ok(vec![UpdateFile { + path: PathBuf::from("file.txt"), + move_path: None, + chunks: vec![UpdateFileChunk { + change_context: None, + old_lines: vec!["old".to_string()], + new_lines: vec!["new".to_string()], + context_line_indices: vec![], + is_end_of_file: false, + }], + }]) + ); + assert_eq!( + parser.finish(), + Ok(vec![UpdateFile { + path: PathBuf::from("file.txt"), + move_path: None, + chunks: vec![UpdateFileChunk { + change_context: None, + old_lines: vec!["old".to_string()], + new_lines: vec!["new".to_string()], + context_line_indices: vec![], + is_end_of_file: false, + }], + }]) + ); + } + + #[test] + fn test_streaming_patch_parser_finish_requires_end_patch() { + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta("*** Begin Patch\n*** Add File: file.txt\n+hello\n"), + Ok(vec![AddFile { + path: PathBuf::from("file.txt"), + contents: "hello\n".to_string(), + }]) + ); + assert_eq!( + parser.finish(), + Err(InvalidPatchError( + "The last line of the patch must be '*** End Patch'".to_string(), + )) + ); + } + + #[test] + fn test_streaming_patch_parser_rejects_content_after_end_patch() { + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta( + "*** Begin Patch\n*** Add File: file.txt\n+hello\n*** End Patch\nextra\n", + ), + Err(InvalidPatchError( + "The last line of the patch must be '*** End Patch'".to_string(), + )) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta( + "*** Begin Patch\n*** Add File: file.txt\n+hello\n*** End Patch\n \t\n", + ), + Ok(vec![AddFile { + path: PathBuf::from("file.txt"), + contents: "hello\n".to_string(), + }]) + ); + } + + #[test] + fn test_streaming_patch_parser_returns_errors() { + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta("bad\n"), + Err(InvalidPatchError( + "The first line of the patch must be '*** Begin Patch'".to_string(), + )) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!(parser.push_delta("*** Begin Patch\n"), Ok(Vec::new())); + assert_eq!( + parser.push_delta("bad\n"), + Err(InvalidHunkError { + message: "'bad' is not a valid hunk header. Valid hunk headers: '*** Add File: {path}', '*** Delete File: {path}', '*** Update File: {path}'" + .to_string(), + line_number: 2, + }) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta("*** Begin Patch\n*** Add File: file.txt\nbad\n"), + Err(InvalidHunkError { + message: "'bad' is not a valid hunk header. Valid hunk headers: '*** Add File: {path}', '*** Delete File: {path}', '*** Update File: {path}'" + .to_string(), + line_number: 3, + }) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta("*** Begin Patch\n*** Delete File: file.txt\nbad\n"), + Err(InvalidHunkError { + message: "'bad' is not a valid hunk header. Valid hunk headers: '*** Add File: {path}', '*** Delete File: {path}', '*** Update File: {path}'" + .to_string(), + line_number: 3, + }) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta("*** Begin Patch\n*** Update File: file.txt\n*** End Patch\n"), + Err(InvalidHunkError { + message: "Update file hunk for path 'file.txt' is empty".to_string(), + line_number: 2, + }) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta( + "*** Begin Patch\n*** Update File: old.txt\n*** Move to: new.txt\n*** Delete File: other.txt\n", + ), + Err(InvalidHunkError { + message: "Update file hunk for path 'old.txt' is empty".to_string(), + line_number: 2, + }) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta("*** Begin Patch\n*** Update File: file.txt\n@@\n*** End Patch\n"), + Err(InvalidHunkError { + message: "Update hunk does not contain any lines".to_string(), + line_number: 4, + }) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta("*** Begin Patch\n*** Update File: file.txt\n@@\n*** End of File\n"), + Err(InvalidHunkError { + message: "Update hunk does not contain any lines".to_string(), + line_number: 4, + }) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta("*** Begin Patch\n*** Update File: file.txt\n@@\n@@\n"), + Err(InvalidHunkError { + message: "Unexpected line found in update hunk: '@@'. Every line should start with ' ' (context line), '+' (added line), or '-' (removed line)" + .to_string(), + line_number: 4, + }) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta("*** Begin Patch\n*** Update File: file.txt\n@@\n-old\nbad\n"), + Err(InvalidHunkError { + message: "Expected update hunk to start with a @@ context marker, got: 'bad'" + .to_string(), + line_number: 5, + }) + ); + + let mut parser = StreamingPatchParser::default(); + assert_eq!( + parser.push_delta( + "*** Begin Patch\n*** Update File: file.txt\n@@\n*** Update File: other.txt\n", + ), + Err(InvalidHunkError { + message: "Unexpected line found in update hunk: '*** Update File: other.txt'. Every line should start with ' ' (context line), '+' (added line), or '-' (removed line)" + .to_string(), + line_number: 4, + }) + ); + } +} diff --git a/codex-rs/apply-patch/src/text_file.rs b/codex-rs/apply-patch/src/text_file.rs new file mode 100644 index 0000000000000000000000000000000000000000..b7d598ca973ec73f0ed064d2360f4d218e2a792d --- /dev/null +++ b/codex-rs/apply-patch/src/text_file.rs @@ -0,0 +1,121 @@ +pub(super) type Replacement = (usize, usize, Vec); + +#[derive(Clone, Copy)] +enum LineEnding { + Lf, + CrLf, + Cr, +} + +impl LineEnding { + fn as_str(self) -> &'static str { + match self { + Self::Lf => "\n", + Self::CrLf => "\r\n", + Self::Cr => "\r", + } + } +} + +struct SourceLine { + text: String, + ending: Option, +} + +pub(super) struct SourceFile { + lines: Vec, + preferred_ending: LineEnding, +} + +impl SourceFile { + /// Splits contents into logical lines while retaining each line ending. + /// + /// The first existing ending becomes the preferred style for inserted + /// lines; files without an ending default to LF. + pub(super) fn parse(contents: &str) -> Self { + let mut lines = Vec::new(); + let mut preferred_ending = None; + let mut line_start = 0; + let mut cursor = 0; + + while cursor < contents.len() { + let (ending, ending_len) = match contents.as_bytes()[cursor] { + b'\r' if contents.as_bytes().get(cursor + 1) == Some(&b'\n') => { + (LineEnding::CrLf, 2) + } + b'\r' => (LineEnding::Cr, 1), + b'\n' => (LineEnding::Lf, 1), + _ => { + cursor += 1; + continue; + } + }; + preferred_ending.get_or_insert(ending); + lines.push(SourceLine { + text: contents[line_start..cursor].to_string(), + ending: Some(ending), + }); + cursor += ending_len; + line_start = cursor; + } + + if line_start < contents.len() { + lines.push(SourceLine { + text: contents[line_start..].to_string(), + ending: None, + }); + } + + Self { + lines, + preferred_ending: preferred_ending.unwrap_or(LineEnding::Lf), + } + } + + pub(super) fn line_texts(&self) -> Vec { + self.lines.iter().map(|line| line.text.clone()).collect() + } + + /// Rebuilds the file from source-ordered, non-overlapping replacements. + /// + /// Unchanged lines retain their original endings, inserted lines use the + /// preferred ending, and every resulting line receives an ending to match + /// apply-patch's historical trailing-newline behavior. + pub(super) fn apply_replacements(&mut self, replacements: &[Replacement]) { + let mut source_lines = std::mem::take(&mut self.lines).into_iter(); + let mut new_lines = Vec::new(); + let mut source_index = 0; + + for (start_idx, old_len, new_segment) in replacements { + debug_assert!(*start_idx >= source_index); + for line in source_lines.by_ref().take(*start_idx - source_index) { + new_lines.push(line); + } + for _ in source_lines.by_ref().take(*old_len) {} + new_lines.extend(new_segment.iter().map(|text| SourceLine { + text: text.clone(), + ending: Some(self.preferred_ending), + })); + source_index = start_idx + old_len; + } + new_lines.extend(source_lines); + self.lines = new_lines; + + // Updates have historically added a trailing newline. This also gives + // an unterminated last line an ending if an insertion moved it inward. + for line in &mut self.lines { + line.ending.get_or_insert(self.preferred_ending); + } + } + + pub(super) fn into_contents(self) -> String { + let mut contents = String::new(); + for line in self.lines { + contents.push_str(&line.text); + if let Some(ending) = line.ending { + contents.push_str(ending.as_str()); + } + } + contents + } +} diff --git a/codex-rs/apply-patch/tests/all.rs b/codex-rs/apply-patch/tests/all.rs new file mode 100644 index 0000000000000000000000000000000000000000..7e136e4cce2a99915ff1b487d1f2b61c500f133a --- /dev/null +++ b/codex-rs/apply-patch/tests/all.rs @@ -0,0 +1,3 @@ +// Single integration test binary that aggregates all test modules. +// The submodules live in `tests/suite/`. +mod suite; diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/.gitattributes b/codex-rs/apply-patch/tests/fixtures/scenarios/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..3961a5e66a7496e08ef82295e353acca0d0b5a8d --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/.gitattributes @@ -0,0 +1,5 @@ +** text eol=lf +023_preserves_crlf_line_endings/input/*.txt -diff -text +023_preserves_crlf_line_endings/expected/*.txt -diff -text +024_preserves_mixed_line_endings/input/*.txt -diff -text +024_preserves_mixed_line_endings/expected/*.txt -diff -text diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/001_add_file/expected/bar.md b/codex-rs/apply-patch/tests/fixtures/scenarios/001_add_file/expected/bar.md new file mode 100644 index 0000000000000000000000000000000000000000..6dfa057f0d4a43d5a3025a9c14dea607de9e1dbb --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/001_add_file/expected/bar.md @@ -0,0 +1 @@ +This is a new file diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/001_add_file/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/001_add_file/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..37735b2a465d5534d403d6e0eba62d4e857973f7 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/001_add_file/patch.txt @@ -0,0 +1,4 @@ +*** Begin Patch +*** Add File: bar.md ++This is a new file +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/002_multiple_operations/expected/modify.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/002_multiple_operations/expected/modify.txt new file mode 100644 index 0000000000000000000000000000000000000000..1b2ee3e56691c2207698ae92e6d33129e40cbbfe --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/002_multiple_operations/expected/modify.txt @@ -0,0 +1,2 @@ +line1 +changed diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/002_multiple_operations/expected/nested/new.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/002_multiple_operations/expected/nested/new.txt new file mode 100644 index 0000000000000000000000000000000000000000..3151666398e7951f81d58893dc687970599625ab --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/002_multiple_operations/expected/nested/new.txt @@ -0,0 +1 @@ +created diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/002_multiple_operations/input/delete.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/002_multiple_operations/input/delete.txt new file mode 100644 index 0000000000000000000000000000000000000000..6e263abce10f69a67055f355ff61e4d51f66962d --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/002_multiple_operations/input/delete.txt @@ -0,0 +1 @@ +obsolete diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/002_multiple_operations/input/modify.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/002_multiple_operations/input/modify.txt new file mode 100644 index 0000000000000000000000000000000000000000..c0d0fb45c382919737f8d0c20aaf57cf89b74af8 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/002_multiple_operations/input/modify.txt @@ -0,0 +1,2 @@ +line1 +line2 diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/002_multiple_operations/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/002_multiple_operations/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..673dec2f78dd6ba4d1dd73824d72aa1ba8294aef --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/002_multiple_operations/patch.txt @@ -0,0 +1,9 @@ +*** Begin Patch +*** Add File: nested/new.txt ++created +*** Delete File: delete.txt +*** Update File: modify.txt +@@ +-line2 ++changed +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/003_multiple_chunks/expected/multi.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/003_multiple_chunks/expected/multi.txt new file mode 100644 index 0000000000000000000000000000000000000000..9054a72916ae61725d3cdb02d8eab81e14cb6e4f --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/003_multiple_chunks/expected/multi.txt @@ -0,0 +1,4 @@ +line1 +changed2 +line3 +changed4 diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/003_multiple_chunks/input/multi.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/003_multiple_chunks/input/multi.txt new file mode 100644 index 0000000000000000000000000000000000000000..84275f9939456e87efd6932bdf7fe01d52a53116 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/003_multiple_chunks/input/multi.txt @@ -0,0 +1,4 @@ +line1 +line2 +line3 +line4 diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/003_multiple_chunks/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/003_multiple_chunks/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..45733c714bb6bbbc3d455798cf1f16ceba947d5f --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/003_multiple_chunks/patch.txt @@ -0,0 +1,9 @@ +*** Begin Patch +*** Update File: multi.txt +@@ +-line2 ++changed2 +@@ +-line4 ++changed4 +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/004_move_to_new_directory/expected/old/other.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/004_move_to_new_directory/expected/old/other.txt new file mode 100644 index 0000000000000000000000000000000000000000..b61039d3df8a5590e91871eb6f131eaa5b0e1dae --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/004_move_to_new_directory/expected/old/other.txt @@ -0,0 +1 @@ +unrelated file diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/004_move_to_new_directory/expected/renamed/dir/name.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/004_move_to_new_directory/expected/renamed/dir/name.txt new file mode 100644 index 0000000000000000000000000000000000000000..b66ba06d315d46280bb09d54614cc52d1677809f --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/004_move_to_new_directory/expected/renamed/dir/name.txt @@ -0,0 +1 @@ +new content diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/004_move_to_new_directory/input/old/name.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/004_move_to_new_directory/input/old/name.txt new file mode 100644 index 0000000000000000000000000000000000000000..33194a0a6f3f99e366d606c24d9b1ab0e0086e69 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/004_move_to_new_directory/input/old/name.txt @@ -0,0 +1 @@ +old content diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/004_move_to_new_directory/input/old/other.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/004_move_to_new_directory/input/old/other.txt new file mode 100644 index 0000000000000000000000000000000000000000..b61039d3df8a5590e91871eb6f131eaa5b0e1dae --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/004_move_to_new_directory/input/old/other.txt @@ -0,0 +1 @@ +unrelated file diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/004_move_to_new_directory/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/004_move_to_new_directory/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..5e2d723a2b429040b300283402bff3491e97dc88 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/004_move_to_new_directory/patch.txt @@ -0,0 +1,7 @@ +*** Begin Patch +*** Update File: old/name.txt +*** Move to: renamed/dir/name.txt +@@ +-old content ++new content +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/005_rejects_empty_patch/expected/foo.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/005_rejects_empty_patch/expected/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..2bf5ad0447d3370461c6f32a0a5bc8a3177376aa --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/005_rejects_empty_patch/expected/foo.txt @@ -0,0 +1 @@ +stable diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/005_rejects_empty_patch/input/foo.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/005_rejects_empty_patch/input/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..2bf5ad0447d3370461c6f32a0a5bc8a3177376aa --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/005_rejects_empty_patch/input/foo.txt @@ -0,0 +1 @@ +stable diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/005_rejects_empty_patch/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/005_rejects_empty_patch/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..4fcfecbbc779ed8f085bdd8247c64fd66ffcbdca --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/005_rejects_empty_patch/patch.txt @@ -0,0 +1,2 @@ +*** Begin Patch +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/006_rejects_missing_context/expected/modify.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/006_rejects_missing_context/expected/modify.txt new file mode 100644 index 0000000000000000000000000000000000000000..c0d0fb45c382919737f8d0c20aaf57cf89b74af8 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/006_rejects_missing_context/expected/modify.txt @@ -0,0 +1,2 @@ +line1 +line2 diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/006_rejects_missing_context/input/modify.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/006_rejects_missing_context/input/modify.txt new file mode 100644 index 0000000000000000000000000000000000000000..c0d0fb45c382919737f8d0c20aaf57cf89b74af8 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/006_rejects_missing_context/input/modify.txt @@ -0,0 +1,2 @@ +line1 +line2 diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/006_rejects_missing_context/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/006_rejects_missing_context/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..488438b12b0652752cdf69e071405756d4cb010a --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/006_rejects_missing_context/patch.txt @@ -0,0 +1,6 @@ +*** Begin Patch +*** Update File: modify.txt +@@ +-missing ++changed +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/007_rejects_missing_file_delete/expected/foo.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/007_rejects_missing_file_delete/expected/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..2bf5ad0447d3370461c6f32a0a5bc8a3177376aa --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/007_rejects_missing_file_delete/expected/foo.txt @@ -0,0 +1 @@ +stable diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/007_rejects_missing_file_delete/input/foo.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/007_rejects_missing_file_delete/input/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..2bf5ad0447d3370461c6f32a0a5bc8a3177376aa --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/007_rejects_missing_file_delete/input/foo.txt @@ -0,0 +1 @@ +stable diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/007_rejects_missing_file_delete/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/007_rejects_missing_file_delete/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..6f95531db3612e28cf116e6cdc8942c110eb7039 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/007_rejects_missing_file_delete/patch.txt @@ -0,0 +1,3 @@ +*** Begin Patch +*** Delete File: missing.txt +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/008_rejects_empty_update_hunk/expected/foo.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/008_rejects_empty_update_hunk/expected/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..2bf5ad0447d3370461c6f32a0a5bc8a3177376aa --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/008_rejects_empty_update_hunk/expected/foo.txt @@ -0,0 +1 @@ +stable diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/008_rejects_empty_update_hunk/input/foo.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/008_rejects_empty_update_hunk/input/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..2bf5ad0447d3370461c6f32a0a5bc8a3177376aa --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/008_rejects_empty_update_hunk/input/foo.txt @@ -0,0 +1 @@ +stable diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/008_rejects_empty_update_hunk/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/008_rejects_empty_update_hunk/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..d7596a362be7e943804d8e9a9ff4045a5104ddfc --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/008_rejects_empty_update_hunk/patch.txt @@ -0,0 +1,3 @@ +*** Begin Patch +*** Update File: foo.txt +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/009_requires_existing_file_for_update/expected/foo.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/009_requires_existing_file_for_update/expected/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..2bf5ad0447d3370461c6f32a0a5bc8a3177376aa --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/009_requires_existing_file_for_update/expected/foo.txt @@ -0,0 +1 @@ +stable diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/009_requires_existing_file_for_update/input/foo.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/009_requires_existing_file_for_update/input/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..2bf5ad0447d3370461c6f32a0a5bc8a3177376aa --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/009_requires_existing_file_for_update/input/foo.txt @@ -0,0 +1 @@ +stable diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/009_requires_existing_file_for_update/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/009_requires_existing_file_for_update/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..a7de4f24c57886b3d43f5f50620e84359a849940 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/009_requires_existing_file_for_update/patch.txt @@ -0,0 +1,6 @@ +*** Begin Patch +*** Update File: missing.txt +@@ +-old ++new +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/expected/old/other.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/expected/old/other.txt new file mode 100644 index 0000000000000000000000000000000000000000..b61039d3df8a5590e91871eb6f131eaa5b0e1dae --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/expected/old/other.txt @@ -0,0 +1 @@ +unrelated file diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/expected/renamed/dir/name.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/expected/renamed/dir/name.txt new file mode 100644 index 0000000000000000000000000000000000000000..3e757656cf36eca53338e520d134963a44f793f8 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/expected/renamed/dir/name.txt @@ -0,0 +1 @@ +new diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/input/old/name.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/input/old/name.txt new file mode 100644 index 0000000000000000000000000000000000000000..3940df7cd8791c48c101a7543b3a9fd7ef930e78 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/input/old/name.txt @@ -0,0 +1 @@ +from diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/input/old/other.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/input/old/other.txt new file mode 100644 index 0000000000000000000000000000000000000000..b61039d3df8a5590e91871eb6f131eaa5b0e1dae --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/input/old/other.txt @@ -0,0 +1 @@ +unrelated file diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/input/renamed/dir/name.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/input/renamed/dir/name.txt new file mode 100644 index 0000000000000000000000000000000000000000..cbaf024e5e7fa87bdd8ad76b1393d9973585bc97 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/input/renamed/dir/name.txt @@ -0,0 +1 @@ +existing diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..c45ce6d7826992507b6e612b49b1724749876c80 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/010_move_overwrites_existing_destination/patch.txt @@ -0,0 +1,7 @@ +*** Begin Patch +*** Update File: old/name.txt +*** Move to: renamed/dir/name.txt +@@ +-from ++new +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/011_add_overwrites_existing_file/expected/duplicate.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/011_add_overwrites_existing_file/expected/duplicate.txt new file mode 100644 index 0000000000000000000000000000000000000000..b66ba06d315d46280bb09d54614cc52d1677809f --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/011_add_overwrites_existing_file/expected/duplicate.txt @@ -0,0 +1 @@ +new content diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/011_add_overwrites_existing_file/input/duplicate.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/011_add_overwrites_existing_file/input/duplicate.txt new file mode 100644 index 0000000000000000000000000000000000000000..33194a0a6f3f99e366d606c24d9b1ab0e0086e69 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/011_add_overwrites_existing_file/input/duplicate.txt @@ -0,0 +1 @@ +old content diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/011_add_overwrites_existing_file/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/011_add_overwrites_existing_file/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..bad9cf3fded135bc7e758b5f15bbd018af114efd --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/011_add_overwrites_existing_file/patch.txt @@ -0,0 +1,4 @@ +*** Begin Patch +*** Add File: duplicate.txt ++new content +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/012_delete_directory_fails/expected/dir/foo.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/012_delete_directory_fails/expected/dir/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..2bf5ad0447d3370461c6f32a0a5bc8a3177376aa --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/012_delete_directory_fails/expected/dir/foo.txt @@ -0,0 +1 @@ +stable diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/012_delete_directory_fails/input/dir/foo.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/012_delete_directory_fails/input/dir/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..2bf5ad0447d3370461c6f32a0a5bc8a3177376aa --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/012_delete_directory_fails/input/dir/foo.txt @@ -0,0 +1 @@ +stable diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/012_delete_directory_fails/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/012_delete_directory_fails/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..a10bcd9ea9c7d89e5f7eafb9b8595b2cd7dd5a33 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/012_delete_directory_fails/patch.txt @@ -0,0 +1,3 @@ +*** Begin Patch +*** Delete File: dir +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/013_rejects_invalid_hunk_header/expected/foo.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/013_rejects_invalid_hunk_header/expected/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..2bf5ad0447d3370461c6f32a0a5bc8a3177376aa --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/013_rejects_invalid_hunk_header/expected/foo.txt @@ -0,0 +1 @@ +stable diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/013_rejects_invalid_hunk_header/input/foo.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/013_rejects_invalid_hunk_header/input/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..2bf5ad0447d3370461c6f32a0a5bc8a3177376aa --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/013_rejects_invalid_hunk_header/input/foo.txt @@ -0,0 +1 @@ +stable diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/013_rejects_invalid_hunk_header/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/013_rejects_invalid_hunk_header/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..b35d7207d758490ff7495c4d7aa63375b6562b5a --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/013_rejects_invalid_hunk_header/patch.txt @@ -0,0 +1,3 @@ +*** Begin Patch +*** Frobnicate File: foo +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/014_update_file_appends_trailing_newline/expected/no_newline.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/014_update_file_appends_trailing_newline/expected/no_newline.txt new file mode 100644 index 0000000000000000000000000000000000000000..06fcdd77c9348567c50638b30d406500f521c304 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/014_update_file_appends_trailing_newline/expected/no_newline.txt @@ -0,0 +1,2 @@ +first line +second line diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/014_update_file_appends_trailing_newline/input/no_newline.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/014_update_file_appends_trailing_newline/input/no_newline.txt new file mode 100644 index 0000000000000000000000000000000000000000..a6e09874b5f93972808f3494a8ee14d4d948fc35 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/014_update_file_appends_trailing_newline/input/no_newline.txt @@ -0,0 +1 @@ +no newline at end diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/014_update_file_appends_trailing_newline/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/014_update_file_appends_trailing_newline/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..4ed5818eb1cb134ce7539c9ecfcda7a6346698d6 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/014_update_file_appends_trailing_newline/patch.txt @@ -0,0 +1,7 @@ +*** Begin Patch +*** Update File: no_newline.txt +@@ +-no newline at end ++first line ++second line +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/015_failure_after_partial_success_leaves_changes/expected/created.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/015_failure_after_partial_success_leaves_changes/expected/created.txt new file mode 100644 index 0000000000000000000000000000000000000000..ce013625030ba8dba906f756967f9e9ca394464a --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/015_failure_after_partial_success_leaves_changes/expected/created.txt @@ -0,0 +1 @@ +hello diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/015_failure_after_partial_success_leaves_changes/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/015_failure_after_partial_success_leaves_changes/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..a6e9709d1f6af0058a709aa758c12e4b94f3cc3b --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/015_failure_after_partial_success_leaves_changes/patch.txt @@ -0,0 +1,8 @@ +*** Begin Patch +*** Add File: created.txt ++hello +*** Update File: missing.txt +@@ +-old ++new +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/016_pure_addition_update_chunk/expected/input.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/016_pure_addition_update_chunk/expected/input.txt new file mode 100644 index 0000000000000000000000000000000000000000..f6d6f0bef856b102699aa7acd595fbff9aed34ee --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/016_pure_addition_update_chunk/expected/input.txt @@ -0,0 +1,4 @@ +line1 +line2 +added line 1 +added line 2 diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/016_pure_addition_update_chunk/input/input.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/016_pure_addition_update_chunk/input/input.txt new file mode 100644 index 0000000000000000000000000000000000000000..c0d0fb45c382919737f8d0c20aaf57cf89b74af8 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/016_pure_addition_update_chunk/input/input.txt @@ -0,0 +1,2 @@ +line1 +line2 diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/016_pure_addition_update_chunk/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/016_pure_addition_update_chunk/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..56337549f9ffe1115f00b3fcf5037dee31cdeda9 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/016_pure_addition_update_chunk/patch.txt @@ -0,0 +1,6 @@ +*** Begin Patch +*** Update File: input.txt +@@ ++added line 1 ++added line 2 +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/017_whitespace_padded_hunk_header/expected/foo.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/017_whitespace_padded_hunk_header/expected/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..3e757656cf36eca53338e520d134963a44f793f8 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/017_whitespace_padded_hunk_header/expected/foo.txt @@ -0,0 +1 @@ +new diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/017_whitespace_padded_hunk_header/input/foo.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/017_whitespace_padded_hunk_header/input/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..3367afdbbf91e638efe983616377c60477cc6612 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/017_whitespace_padded_hunk_header/input/foo.txt @@ -0,0 +1 @@ +old diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/017_whitespace_padded_hunk_header/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/017_whitespace_padded_hunk_header/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..21e6c1958d8c834b899243c6309b862b78aa91b1 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/017_whitespace_padded_hunk_header/patch.txt @@ -0,0 +1,6 @@ +*** Begin Patch + *** Update File: foo.txt +@@ +-old ++new +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/018_whitespace_padded_patch_markers/expected/file.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/018_whitespace_padded_patch_markers/expected/file.txt new file mode 100644 index 0000000000000000000000000000000000000000..f719efd430d52bcfc8566a43b2eb655688d38871 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/018_whitespace_padded_patch_markers/expected/file.txt @@ -0,0 +1 @@ +two diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/018_whitespace_padded_patch_markers/input/file.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/018_whitespace_padded_patch_markers/input/file.txt new file mode 100644 index 0000000000000000000000000000000000000000..5626abf0f72e58d7a153368ba57db4c673c0e171 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/018_whitespace_padded_patch_markers/input/file.txt @@ -0,0 +1 @@ +one diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/018_whitespace_padded_patch_markers/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/018_whitespace_padded_patch_markers/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..2648721797e8017ab4ab561d082e93c08c571e33 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/018_whitespace_padded_patch_markers/patch.txt @@ -0,0 +1,6 @@ + *** Begin Patch +*** Update File: file.txt +@@ +-one ++two +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/019_unicode_simple/expected/foo.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/019_unicode_simple/expected/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..99d5a6e9c2f45d932e02f40796d655bf13fd513b --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/019_unicode_simple/expected/foo.txt @@ -0,0 +1,3 @@ +line1 +naïve café ✅ +line3 diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/019_unicode_simple/input/foo.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/019_unicode_simple/input/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..b17094871378284e284c29384e0384a7f96268fe --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/019_unicode_simple/input/foo.txt @@ -0,0 +1,3 @@ +line1 +naïve café +line3 diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/019_unicode_simple/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/019_unicode_simple/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..9514207fdeacc7fad3b12192ecedb521a5458128 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/019_unicode_simple/patch.txt @@ -0,0 +1,7 @@ +*** Begin Patch +*** Update File: foo.txt +@@ + line1 +-naïve café ++naïve café ✅ +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/020_delete_file_success/expected/keep.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/020_delete_file_success/expected/keep.txt new file mode 100644 index 0000000000000000000000000000000000000000..2fa992c0b8b5c6acd2bdd4fa31de29d29799bdd5 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/020_delete_file_success/expected/keep.txt @@ -0,0 +1 @@ +keep diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/020_delete_file_success/input/keep.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/020_delete_file_success/input/keep.txt new file mode 100644 index 0000000000000000000000000000000000000000..2fa992c0b8b5c6acd2bdd4fa31de29d29799bdd5 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/020_delete_file_success/input/keep.txt @@ -0,0 +1 @@ +keep diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/020_delete_file_success/input/obsolete.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/020_delete_file_success/input/obsolete.txt new file mode 100644 index 0000000000000000000000000000000000000000..6e263abce10f69a67055f355ff61e4d51f66962d --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/020_delete_file_success/input/obsolete.txt @@ -0,0 +1 @@ +obsolete diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/020_delete_file_success/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/020_delete_file_success/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..5978f73889453c2659941c191713c4c440633c6b --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/020_delete_file_success/patch.txt @@ -0,0 +1,3 @@ +*** Begin Patch +*** Delete File: obsolete.txt +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/020_whitespace_padded_patch_marker_lines/expected/file.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/020_whitespace_padded_patch_marker_lines/expected/file.txt new file mode 100644 index 0000000000000000000000000000000000000000..f719efd430d52bcfc8566a43b2eb655688d38871 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/020_whitespace_padded_patch_marker_lines/expected/file.txt @@ -0,0 +1 @@ +two diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/020_whitespace_padded_patch_marker_lines/input/file.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/020_whitespace_padded_patch_marker_lines/input/file.txt new file mode 100644 index 0000000000000000000000000000000000000000..5626abf0f72e58d7a153368ba57db4c673c0e171 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/020_whitespace_padded_patch_marker_lines/input/file.txt @@ -0,0 +1 @@ +one diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/020_whitespace_padded_patch_marker_lines/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/020_whitespace_padded_patch_marker_lines/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..3d2a1dbe5ec3c5860a614ad8f6174d2f28c8e312 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/020_whitespace_padded_patch_marker_lines/patch.txt @@ -0,0 +1,6 @@ +*** Begin Patch +*** Update File: file.txt +@@ +-one ++two + *** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/021_update_file_deletion_only/expected/lines.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/021_update_file_deletion_only/expected/lines.txt new file mode 100644 index 0000000000000000000000000000000000000000..8129d305c8e06efd4a0e742b2aab8d0958a2bc0f --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/021_update_file_deletion_only/expected/lines.txt @@ -0,0 +1,2 @@ +line1 +line3 diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/021_update_file_deletion_only/input/lines.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/021_update_file_deletion_only/input/lines.txt new file mode 100644 index 0000000000000000000000000000000000000000..83db48f84ec878fbfb30b46d16630e944e34f205 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/021_update_file_deletion_only/input/lines.txt @@ -0,0 +1,3 @@ +line1 +line2 +line3 diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/021_update_file_deletion_only/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/021_update_file_deletion_only/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..860c6c9a9902885b5ab3c5a4a403a0827b0e793a --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/021_update_file_deletion_only/patch.txt @@ -0,0 +1,7 @@ +*** Begin Patch +*** Update File: lines.txt +@@ + line1 +-line2 + line3 +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/022_update_file_end_of_file_marker/expected/tail.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/022_update_file_end_of_file_marker/expected/tail.txt new file mode 100644 index 0000000000000000000000000000000000000000..87463f92d717a4901543a4058f9c930203e51997 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/022_update_file_end_of_file_marker/expected/tail.txt @@ -0,0 +1,2 @@ +first +second updated diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/022_update_file_end_of_file_marker/input/tail.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/022_update_file_end_of_file_marker/input/tail.txt new file mode 100644 index 0000000000000000000000000000000000000000..66a52ee7a1d803dc57859c3e95ac9dcdc87c0164 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/022_update_file_end_of_file_marker/input/tail.txt @@ -0,0 +1,2 @@ +first +second diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/022_update_file_end_of_file_marker/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/022_update_file_end_of_file_marker/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..8b16b5bd9ee31912d1e91dec35b38678f508c87e --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/022_update_file_end_of_file_marker/patch.txt @@ -0,0 +1,8 @@ +*** Begin Patch +*** Update File: tail.txt +@@ + first +-second ++second updated +*** End of File +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/023_preserves_crlf_line_endings/expected/lines.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/023_preserves_crlf_line_endings/expected/lines.txt new file mode 100644 index 0000000000000000000000000000000000000000..43321636e611bea29569e72f60494ae3c6226ae2 Binary files /dev/null and b/codex-rs/apply-patch/tests/fixtures/scenarios/023_preserves_crlf_line_endings/expected/lines.txt differ diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/023_preserves_crlf_line_endings/input/lines.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/023_preserves_crlf_line_endings/input/lines.txt new file mode 100644 index 0000000000000000000000000000000000000000..e1587ff9ea6ced3af32e07e4b6b68d21374907ec Binary files /dev/null and b/codex-rs/apply-patch/tests/fixtures/scenarios/023_preserves_crlf_line_endings/input/lines.txt differ diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/023_preserves_crlf_line_endings/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/023_preserves_crlf_line_endings/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..6df72a09cfead674e32edbf76a1509e656031c22 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/023_preserves_crlf_line_endings/patch.txt @@ -0,0 +1,9 @@ +*** Begin Patch +*** Update File: lines.txt +@@ +-one ++ONE + two ++between + three +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/024_preserves_mixed_line_endings/expected/lines.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/024_preserves_mixed_line_endings/expected/lines.txt new file mode 100644 index 0000000000000000000000000000000000000000..967a20a3a7f3da2cbe11f49a82d07de06c7dae58 Binary files /dev/null and b/codex-rs/apply-patch/tests/fixtures/scenarios/024_preserves_mixed_line_endings/expected/lines.txt differ diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/024_preserves_mixed_line_endings/input/lines.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/024_preserves_mixed_line_endings/input/lines.txt new file mode 100644 index 0000000000000000000000000000000000000000..92a67d3b1cf355f3828c2cbf19a35b73e1a575ef Binary files /dev/null and b/codex-rs/apply-patch/tests/fixtures/scenarios/024_preserves_mixed_line_endings/input/lines.txt differ diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/024_preserves_mixed_line_endings/patch.txt b/codex-rs/apply-patch/tests/fixtures/scenarios/024_preserves_mixed_line_endings/patch.txt new file mode 100644 index 0000000000000000000000000000000000000000..a97b2f35a4b0de0ccbf352dcc548cace7b088bbf --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/024_preserves_mixed_line_endings/patch.txt @@ -0,0 +1,9 @@ +*** Begin Patch +*** Update File: lines.txt +@@ + one + two +-three ++THREE + four +*** End Patch diff --git a/codex-rs/apply-patch/tests/fixtures/scenarios/README.md b/codex-rs/apply-patch/tests/fixtures/scenarios/README.md new file mode 100644 index 0000000000000000000000000000000000000000..65d1fbe2e4bbf48e8797deda2974a33b9460a7a9 --- /dev/null +++ b/codex-rs/apply-patch/tests/fixtures/scenarios/README.md @@ -0,0 +1,18 @@ +# Overview +This directory is a collection of end to end tests for the apply-patch specification, meant to be easily portable to other languages or platforms. + + +# Specification +Each test case is one directory, composed of input state (input/), the patch operation (patch.txt), and the expected final state (expected/). This structure is designed to keep tests simple (i.e. test exactly one patch at a time) while still providing enough flexibility to test any given operation across files. + +Here's what this would look like for a simple test apply-patch test case to create a new file: + +``` +001_add/ + input/ + foo.md + expected/ + foo.md + bar.md + patch.txt +``` diff --git a/codex-rs/apply-patch/tests/suite/cli.rs b/codex-rs/apply-patch/tests/suite/cli.rs new file mode 100644 index 0000000000000000000000000000000000000000..c982c7aa864ffe80ba8f797439379ada50a20ad2 --- /dev/null +++ b/codex-rs/apply-patch/tests/suite/cli.rs @@ -0,0 +1,91 @@ +use assert_cmd::Command; +use std::fs; +use tempfile::tempdir; + +fn apply_patch_command() -> anyhow::Result { + Ok(Command::new(codex_utils_cargo_bin::cargo_bin( + "apply_patch", + )?)) +} + +#[test] +fn test_apply_patch_cli_add_and_update() -> anyhow::Result<()> { + let tmp = tempdir()?; + let file = "cli_test.txt"; + let absolute_path = tmp.path().join(file); + + // 1) Add a file + let add_patch = format!( + r#"*** Begin Patch +*** Add File: {file} ++hello +*** End Patch"# + ); + apply_patch_command()? + .arg(add_patch) + .current_dir(tmp.path()) + .assert() + .success() + .stdout(format!("Success. Updated the following files:\nA {file}\n")); + assert_eq!(fs::read_to_string(&absolute_path)?, "hello\n"); + + // 2) Update the file + let update_patch = format!( + r#"*** Begin Patch +*** Update File: {file} +@@ +-hello ++world +*** End Patch"# + ); + apply_patch_command()? + .arg(update_patch) + .current_dir(tmp.path()) + .assert() + .success() + .stdout(format!("Success. Updated the following files:\nM {file}\n")); + assert_eq!(fs::read_to_string(&absolute_path)?, "world\n"); + + Ok(()) +} + +#[test] +fn test_apply_patch_cli_stdin_add_and_update() -> anyhow::Result<()> { + let tmp = tempdir()?; + let file = "cli_test_stdin.txt"; + let absolute_path = tmp.path().join(file); + + // 1) Add a file via stdin + let add_patch = format!( + r#"*** Begin Patch +*** Add File: {file} ++hello +*** End Patch"# + ); + apply_patch_command()? + .current_dir(tmp.path()) + .write_stdin(add_patch) + .assert() + .success() + .stdout(format!("Success. Updated the following files:\nA {file}\n")); + assert_eq!(fs::read_to_string(&absolute_path)?, "hello\n"); + + // 2) Update the file via stdin + let update_patch = format!( + r#"*** Begin Patch +*** Update File: {file} +@@ +-hello ++world +*** End Patch"# + ); + apply_patch_command()? + .current_dir(tmp.path()) + .write_stdin(update_patch) + .assert() + .success() + .stdout(format!("Success. Updated the following files:\nM {file}\n")); + assert_eq!(fs::read_to_string(&absolute_path)?, "world\n"); + + Ok(()) +} diff --git a/codex-rs/apply-patch/tests/suite/mod.rs b/codex-rs/apply-patch/tests/suite/mod.rs new file mode 100644 index 0000000000000000000000000000000000000000..6c827e70b388df45eb0994d67aff6743db207e33 --- /dev/null +++ b/codex-rs/apply-patch/tests/suite/mod.rs @@ -0,0 +1,6 @@ +mod cli; +#[cfg(unix)] +mod no_follow; +mod scenarios; +#[cfg(not(target_os = "windows"))] +mod tool; diff --git a/codex-rs/apply-patch/tests/suite/no_follow.rs b/codex-rs/apply-patch/tests/suite/no_follow.rs new file mode 100644 index 0000000000000000000000000000000000000000..d2b05bfe1ee87529a36d9fb591f5cfd3e95f0a19 --- /dev/null +++ b/codex-rs/apply-patch/tests/suite/no_follow.rs @@ -0,0 +1,153 @@ +use codex_apply_patch::ApplyPatchOptions; +use codex_apply_patch::MaybeApplyPatchVerified; +use codex_apply_patch::apply_patch_with_options; +use codex_apply_patch::parse_patch; +use codex_apply_patch::verify_apply_patch_args; +use codex_exec_server::LOCAL_FS; +use codex_utils_path_uri::PathUri; +use pretty_assertions::assert_eq; +use std::fs; +use std::os::unix::fs::symlink; + +#[tokio::test] +async fn no_follow_rejects_leaf_and_ancestor_links_for_every_patch_operation() -> anyhow::Result<()> +{ + for path in ["link.txt", "linked/victim.txt"] { + for body in [ + format!("*** Add File: {path}\n+changed"), + format!("*** Update File: {path}\n@@\n-original\n+changed"), + format!("*** Delete File: {path}"), + format!("*** Update File: {path}\n*** Move to: moved.txt\n@@\n-original\n+changed"), + format!("*** Update File: source.txt\n*** Move to: {path}\n@@\n-original\n+changed"), + "*** Add File: linked/new/nested.txt\n+changed".to_string(), + "*** Update File: source.txt\n*** Move to: linked/new/nested.txt\n@@\n-original\n+changed".to_string(), + ] { + let temp = tempfile::tempdir()?; + // macOS's temporary directory can itself contain a /var symlink. + let root = temp.path().canonicalize()?; + let work = root.join("work"); + let outside = root.join("outside"); + fs::create_dir(&work)?; + fs::create_dir(&outside)?; + fs::write(outside.join("victim.txt"), "original\n")?; + fs::write(work.join("source.txt"), "original\n")?; + symlink(outside.join("victim.txt"), work.join("link.txt"))?; + symlink(&outside, work.join("linked"))?; + + let patch = format!("*** Begin Patch\n{body}\n*** End Patch"); + let failure = apply_patch_with_options( + &patch, + ApplyPatchOptions { + follow_symlinks: false, + ..Default::default() + }, + &PathUri::from_host_native_path(&work)?, + &mut Vec::new(), + &mut Vec::new(), + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .expect_err(&body); + + assert!(failure.delta().is_empty(), "{body}"); + assert_eq!(fs::read_to_string(outside.join("victim.txt"))?, "original\n"); + assert_eq!(fs::read_to_string(work.join("source.txt"))?, "original\n"); + assert!(!outside.join("new").exists(), "{body}"); + assert!(!work.join("moved.txt").exists(), "{body}"); + assert!(fs::symlink_metadata(work.join("link.txt"))?.is_symlink()); + } + } + Ok(()) +} + +#[tokio::test] +async fn no_follow_rechecks_paths_after_verification() -> anyhow::Result<()> { + let temp = tempfile::tempdir()?; + let root = temp.path().canonicalize()?; + fs::create_dir(root.join("approved"))?; + fs::create_dir(root.join("outside"))?; + fs::write(root.join("approved/file.txt"), "original\n")?; + fs::write(root.join("outside/file.txt"), "original\n")?; + let cwd = PathUri::from_host_native_path(&root)?; + let patch = "*** Begin Patch\n*** Update File: approved/file.txt\n@@\n-original\n+changed\n*** End Patch"; + let MaybeApplyPatchVerified::Body(action) = verify_apply_patch_args( + parse_patch(patch)?, + &cwd, + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + else { + panic!("regular-file patch should verify"); + }; + fs::rename(root.join("approved"), root.join("original"))?; + symlink(root.join("outside"), root.join("approved"))?; + + apply_patch_with_options( + &action.patch, + ApplyPatchOptions { + follow_symlinks: false, + ..Default::default() + }, + &cwd, + &mut Vec::new(), + &mut Vec::new(), + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await + .expect_err("swapped ancestor must not be followed"); + assert_eq!( + fs::read_to_string(root.join("outside/file.txt"))?, + "original\n" + ); + assert_eq!( + fs::read_to_string(root.join("original/file.txt"))?, + "original\n" + ); + Ok(()) +} + +#[tokio::test] +async fn no_follow_applies_regular_files_and_default_still_follows_links() -> anyhow::Result<()> { + let temp = tempfile::tempdir()?; + let root = temp.path().canonicalize()?; + let cwd = PathUri::from_host_native_path(&root)?; + let patch = "*** Begin Patch\n*** Add File: new/file.txt\n+original\n*** Update File: new/file.txt\n*** Move to: moved/file.txt\n@@\n-original\n+changed\n*** Add File: remove.txt\n+remove\n*** Delete File: remove.txt\n*** End Patch"; + apply_patch_with_options( + patch, + ApplyPatchOptions { + follow_symlinks: false, + ..Default::default() + }, + &cwd, + &mut Vec::new(), + &mut Vec::new(), + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await?; + assert_eq!( + fs::read_to_string(root.join("moved/file.txt"))?, + "changed\n" + ); + assert!(!root.join("new/file.txt").exists()); + assert!(!root.join("remove.txt").exists()); + + symlink(root.join("moved/file.txt"), root.join("link.txt"))?; + codex_apply_patch::apply_patch( + "*** Begin Patch\n*** Update File: link.txt\n@@\n-changed\n+followed\n*** End Patch", + &cwd, + &mut Vec::new(), + &mut Vec::new(), + LOCAL_FS.as_ref(), + /*sandbox*/ None, + ) + .await?; + assert_eq!( + fs::read_to_string(root.join("moved/file.txt"))?, + "followed\n" + ); + Ok(()) +} diff --git a/codex-rs/apply-patch/tests/suite/scenarios.rs b/codex-rs/apply-patch/tests/suite/scenarios.rs new file mode 100644 index 0000000000000000000000000000000000000000..ed8d6a37a44db9952aab2f9bd2ff72bbd13b1f63 --- /dev/null +++ b/codex-rs/apply-patch/tests/suite/scenarios.rs @@ -0,0 +1,133 @@ +use anyhow::Context; +use codex_apply_patch::CODEX_APPLY_PATCH_PRESERVE_LINE_ENDINGS_ENV_VAR; +use codex_utils_cargo_bin::find_resource; +use pretty_assertions::assert_eq; +use std::collections::BTreeMap; +use std::fs; +use std::path::Path; +use std::path::PathBuf; +use std::process::Command; +use tempfile::tempdir; + +#[test] +fn test_apply_patch_scenarios() -> anyhow::Result<()> { + let scenarios_marker = find_resource!("tests/fixtures/scenarios/.gitattributes")?; + let scenarios_dir = scenarios_marker + .parent() + .context("scenario marker should have a parent directory")?; + for scenario in fs::read_dir(scenarios_dir) + .with_context(|| format!("failed to read {}", scenarios_dir.display()))? + { + let scenario = scenario?; + let path = scenario.path(); + if path.is_dir() { + run_apply_patch_scenario(&path) + .with_context(|| format!("failed to run scenario {}", path.display()))?; + } + } + Ok(()) +} + +/// Reads a scenario directory, copies the input files to a temporary directory, runs apply-patch, +/// and asserts that the final state matches the expected state exactly. +fn run_apply_patch_scenario(dir: &Path) -> anyhow::Result<()> { + let tmp = tempdir()?; + + // Copy the input files to the temporary directory + let input_dir = dir.join("input"); + if input_dir.is_dir() { + copy_dir_recursive(&input_dir, tmp.path())?; + } + + // Read the patch.txt file + let patch_path = dir.join("patch.txt"); + let patch = fs::read_to_string(&patch_path) + .with_context(|| format!("failed to read {}", patch_path.display()))?; + + // Run apply_patch in the temporary directory. We intentionally do not assert + // on the exit status here; the scenarios are specified purely in terms of + // final filesystem state, which we compare below. + Command::new(codex_utils_cargo_bin::cargo_bin("apply_patch")?) + .arg(patch) + .env(CODEX_APPLY_PATCH_PRESERVE_LINE_ENDINGS_ENV_VAR, "1") + .current_dir(tmp.path()) + .output() + .with_context(|| format!("failed to run scenario {}", dir.display()))?; + + // Assert that the final state matches the expected state exactly + let expected_dir = dir.join("expected"); + let expected_snapshot = snapshot_dir(&expected_dir)?; + let actual_snapshot = snapshot_dir(tmp.path())?; + + assert_eq!( + actual_snapshot, + expected_snapshot, + "Scenario {} did not match expected final state", + dir.display() + ); + + Ok(()) +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum Entry { + File(Vec), + Dir, +} + +fn snapshot_dir(root: &Path) -> anyhow::Result> { + let mut entries = BTreeMap::new(); + if root.is_dir() { + snapshot_dir_recursive(root, root, &mut entries)?; + } + Ok(entries) +} + +fn snapshot_dir_recursive( + base: &Path, + dir: &Path, + entries: &mut BTreeMap, +) -> anyhow::Result<()> { + for entry in fs::read_dir(dir)? { + let entry = entry?; + let path = entry.path(); + let Some(stripped) = path.strip_prefix(base).ok() else { + continue; + }; + let rel = stripped.to_path_buf(); + + // Under Buck2, files in `__srcs` are often materialized as symlinks. + // Use `metadata()` (follows symlinks) so our fixture snapshots work + // under both Cargo and Buck2. + let metadata = fs::metadata(&path)?; + if metadata.is_dir() { + entries.insert(rel.clone(), Entry::Dir); + snapshot_dir_recursive(base, &path, entries)?; + } else if metadata.is_file() { + let contents = fs::read(&path)?; + entries.insert(rel, Entry::File(contents)); + } + } + Ok(()) +} + +fn copy_dir_recursive(src: &Path, dst: &Path) -> anyhow::Result<()> { + for entry in fs::read_dir(src)? { + let entry = entry?; + let path = entry.path(); + let dest_path = dst.join(entry.file_name()); + + // See note in `snapshot_dir_recursive` about Buck2 symlink trees. + let metadata = fs::metadata(&path)?; + if metadata.is_dir() { + fs::create_dir_all(&dest_path)?; + copy_dir_recursive(&path, &dest_path)?; + } else if metadata.is_file() { + if let Some(parent) = dest_path.parent() { + fs::create_dir_all(parent)?; + } + fs::copy(&path, &dest_path)?; + } + } + Ok(()) +} diff --git a/codex-rs/apply-patch/tests/suite/tool.rs b/codex-rs/apply-patch/tests/suite/tool.rs new file mode 100644 index 0000000000000000000000000000000000000000..36d37efe1c648854815cdbcd6531c0d9a524b064 --- /dev/null +++ b/codex-rs/apply-patch/tests/suite/tool.rs @@ -0,0 +1,437 @@ +use assert_cmd::Command; +use codex_apply_patch::CODEX_APPLY_PATCH_PRESERVE_LINE_ENDINGS_ENV_VAR; +use pretty_assertions::assert_eq; +use std::fs; +use std::path::Path; +use std::path::PathBuf; +use tempfile::tempdir; + +fn run_apply_patch_in_dir(dir: &Path, patch: &str) -> anyhow::Result { + let mut cmd = Command::new(codex_utils_cargo_bin::cargo_bin("apply_patch")?); + cmd.env(CODEX_APPLY_PATCH_PRESERVE_LINE_ENDINGS_ENV_VAR, "1"); + cmd.current_dir(dir); + Ok(cmd.arg(patch).assert()) +} + +fn assert_apply_patch_updates_file( + file_name: &str, + original: &[u8], + patch: &str, + expected: &[u8], +) -> anyhow::Result<()> { + let tmp = tempdir()?; + let target_path = tmp.path().join(file_name); + fs::write(&target_path, original)?; + + run_apply_patch_in_dir(tmp.path(), patch)? + .success() + .stdout(format!( + "Success. Updated the following files:\nM {file_name}\n" + )); + + assert_eq!(fs::read(target_path)?, expected); + Ok(()) +} + +fn apply_patch_command(dir: &Path) -> anyhow::Result { + let mut cmd = Command::new(codex_utils_cargo_bin::cargo_bin("apply_patch")?); + cmd.env(CODEX_APPLY_PATCH_PRESERVE_LINE_ENDINGS_ENV_VAR, "1"); + cmd.current_dir(dir); + Ok(cmd) +} + +fn resolved_under(root: &Path, path: &str) -> anyhow::Result { + Ok(root.canonicalize()?.join(path)) +} + +#[test] +fn test_apply_patch_cli_applies_multiple_operations() -> anyhow::Result<()> { + let tmp = tempdir()?; + let add_path = tmp.path().join("nested/new.txt"); + let modify_path = tmp.path().join("modify.txt"); + let delete_path = tmp.path().join("delete.txt"); + + fs::write(&modify_path, "line1\nline2\n")?; + fs::write(&delete_path, "obsolete\n")?; + + let patch = "*** Begin Patch\n*** Add File: nested/new.txt\n+created\n*** Delete File: delete.txt\n*** Update File: modify.txt\n@@\n-line2\n+changed\n*** End Patch"; + + run_apply_patch_in_dir(tmp.path(), patch)?.success().stdout( + "Success. Updated the following files:\nA nested/new.txt\nM modify.txt\nD delete.txt\n", + ); + + assert_eq!(fs::read_to_string(add_path)?, "created\n"); + assert_eq!(fs::read_to_string(&modify_path)?, "line1\nchanged\n"); + assert!(!delete_path.exists()); + + Ok(()) +} + +#[test] +fn test_apply_patch_cli_applies_multiple_chunks() -> anyhow::Result<()> { + let tmp = tempdir()?; + let target_path = tmp.path().join("multi.txt"); + fs::write(&target_path, "line1\nline2\nline3\nline4\n")?; + + let patch = "*** Begin Patch\n*** Update File: multi.txt\n@@\n-line2\n+changed2\n@@\n-line4\n+changed4\n*** End Patch"; + + run_apply_patch_in_dir(tmp.path(), patch)? + .success() + .stdout("Success. Updated the following files:\nM multi.txt\n"); + + assert_eq!( + fs::read_to_string(&target_path)?, + "line1\nchanged2\nline3\nchanged4\n" + ); + + Ok(()) +} + +#[test] +fn test_apply_patch_cli_rejects_overlapping_end_of_file_chunks() -> anyhow::Result<()> { + let tmp = tempdir()?; + let target_path = tmp.path().join("overlapping.txt"); + let expected_target_path = resolved_under(tmp.path(), "overlapping.txt")?; + fs::write(&target_path, "one\n")?; + + let patch = "*** Begin Patch\n*** Update File: overlapping.txt\n@@\n-one\n+first\n@@\n-one\n+second\n*** End of File\n*** End Patch"; + + run_apply_patch_in_dir(tmp.path(), patch)? + .failure() + .stderr(format!( + "Failed to find expected lines in {}:\none\n", + expected_target_path.display() + )); + + assert_eq!(fs::read_to_string(target_path)?, "one\n"); + Ok(()) +} + +#[test] +fn test_apply_patch_cli_allows_overlapping_eof_chunks_in_legacy_mode() -> anyhow::Result<()> { + let tmp = tempdir()?; + let target_path = tmp.path().join("overlapping.txt"); + fs::write(&target_path, "one\n")?; + + let patch = "*** Begin Patch\n*** Update File: overlapping.txt\n@@\n-one\n+first\n@@\n-one\n+second\n*** End of File\n*** End Patch"; + + Command::new(codex_utils_cargo_bin::cargo_bin("apply_patch")?) + .env_remove(CODEX_APPLY_PATCH_PRESERVE_LINE_ENDINGS_ENV_VAR) + .arg(patch) + .current_dir(tmp.path()) + .assert() + .success() + .stdout("Success. Updated the following files:\nM overlapping.txt\n"); + + assert_eq!(fs::read_to_string(target_path)?, "first\n"); + Ok(()) +} + +#[test] +fn test_apply_patch_cli_preserves_crlf_from_target_file() -> anyhow::Result<()> { + let patch = "*** Begin Patch\n*** Update File: crlf.txt\n@@\n-one\n+uno\n@@\n two\n+\n+between\n three\n*** End Patch"; + + assert_apply_patch_updates_file( + "crlf.txt", + b"one\r\ntwo\r\nthree\r\n", + patch, + b"uno\r\ntwo\r\n\r\nbetween\r\nthree\r\n", + ) +} + +#[test] +fn test_apply_patch_cli_appends_after_trailing_blank_crlf_line() -> anyhow::Result<()> { + let patch = "*** Begin Patch\n*** Update File: trailing_blank.txt\n@@\n+new\n*** End Patch"; + + assert_apply_patch_updates_file( + "trailing_blank.txt", + b"a\r\n\r\n", + patch, + b"a\r\n\r\nnew\r\n", + ) +} + +#[test] +fn test_apply_patch_cli_uses_legacy_line_handling_without_rollout_env() -> anyhow::Result<()> { + let tmp = tempdir()?; + let target_path = tmp.path().join("crlf.txt"); + fs::write(&target_path, b"one\r\n")?; + let patch = "*** Begin Patch\n*** Update File: crlf.txt\n@@\n-one\n+uno\n*** End Patch"; + + Command::new(codex_utils_cargo_bin::cargo_bin("apply_patch")?) + .env_remove(CODEX_APPLY_PATCH_PRESERVE_LINE_ENDINGS_ENV_VAR) + .arg(patch) + .current_dir(tmp.path()) + .assert() + .success(); + + assert_eq!(fs::read(target_path)?, b"uno\n"); + Ok(()) +} + +#[test] +fn test_apply_patch_cli_preserves_cr_from_target_file() -> anyhow::Result<()> { + let patch = "*** Begin Patch\n*** Update File: cr.txt\n@@\n-one\n+uno\n@@\n two\n+\n+between\n three\n*** End Patch"; + + assert_apply_patch_updates_file( + "cr.txt", + b"one\rtwo\rthree\r", + patch, + b"uno\rtwo\r\rbetween\rthree\r", + ) +} + +#[test] +fn test_apply_patch_cli_preserves_change_order_with_repeated_lines() -> anyhow::Result<()> { + let patch = + "*** Begin Patch\n*** Update File: repeated.txt\n@@\n-a\n-b\n+b\n+b\n+a\n*** End Patch"; + + assert_apply_patch_updates_file("repeated.txt", b"a\nb\n", patch, b"b\nb\na\n") +} + +#[test] +fn test_apply_patch_cli_preserves_repeated_context_line_ending() -> anyhow::Result<()> { + let patch = + "*** Begin Patch\n*** Update File: repeated_context.txt\n@@\n-same\n same\n*** End Patch"; + + assert_apply_patch_updates_file("repeated_context.txt", b"same\r\nsame\n", patch, b"same\n") +} + +#[test] +fn test_apply_patch_cli_preserves_untouched_mixed_line_endings() -> anyhow::Result<()> { + let patch = "*** Begin Patch\n*** Update File: mixed.txt\n@@\n one\n two\n-three\n+THREE\n four\n*** End Patch"; + + assert_apply_patch_updates_file( + "mixed.txt", + b"one\r\ntwo\rthree\nfour\r\n", + patch, + b"one\r\ntwo\rTHREE\r\nfour\r\n", + ) +} + +#[test] +fn test_apply_patch_cli_uses_crlf_for_new_trailing_newline() -> anyhow::Result<()> { + let patch = + "*** Begin Patch\n*** Update File: no_trailing_newline.txt\n@@\n-one\n+ONE\n*** End Patch"; + + assert_apply_patch_updates_file( + "no_trailing_newline.txt", + b"one\r\ntwo", + patch, + b"ONE\r\ntwo\r\n", + ) +} + +#[test] +fn test_apply_patch_cli_moves_file_to_new_directory() -> anyhow::Result<()> { + let tmp = tempdir()?; + let original_path = tmp.path().join("old/name.txt"); + let new_path = tmp.path().join("renamed/dir/name.txt"); + fs::create_dir_all(original_path.parent().expect("parent should exist"))?; + fs::write(&original_path, "old content\n")?; + + let patch = "*** Begin Patch\n*** Update File: old/name.txt\n*** Move to: renamed/dir/name.txt\n@@\n-old content\n+new content\n*** End Patch"; + + run_apply_patch_in_dir(tmp.path(), patch)? + .success() + .stdout("Success. Updated the following files:\nM renamed/dir/name.txt\n"); + + assert!(!original_path.exists()); + assert_eq!(fs::read_to_string(&new_path)?, "new content\n"); + + Ok(()) +} + +#[test] +fn test_apply_patch_cli_rejects_empty_patch() -> anyhow::Result<()> { + let tmp = tempdir()?; + + apply_patch_command(tmp.path())? + .arg("*** Begin Patch\n*** End Patch") + .assert() + .failure() + .stderr("No files were modified.\n"); + + Ok(()) +} + +#[test] +fn test_apply_patch_cli_reports_missing_context() -> anyhow::Result<()> { + let tmp = tempdir()?; + let target_path = tmp.path().join("modify.txt"); + let expected_target_path = resolved_under(tmp.path(), "modify.txt")?; + fs::write(&target_path, "line1\nline2\n")?; + + apply_patch_command(tmp.path())? + .arg("*** Begin Patch\n*** Update File: modify.txt\n@@\n-missing\n+changed\n*** End Patch") + .assert() + .failure() + .stderr(format!( + "Failed to find expected lines in {}:\nmissing\n", + expected_target_path.display() + )); + assert_eq!(fs::read_to_string(&target_path)?, "line1\nline2\n"); + + Ok(()) +} + +#[test] +fn test_apply_patch_cli_rejects_missing_file_delete() -> anyhow::Result<()> { + let tmp = tempdir()?; + let missing_path = resolved_under(tmp.path(), "missing.txt")?; + + apply_patch_command(tmp.path())? + .arg("*** Begin Patch\n*** Delete File: missing.txt\n*** End Patch") + .assert() + .failure() + .stderr(format!( + "Failed to delete file {}\n", + missing_path.display() + )); + + Ok(()) +} + +#[test] +fn test_apply_patch_cli_rejects_empty_update_hunk() -> anyhow::Result<()> { + let tmp = tempdir()?; + + apply_patch_command(tmp.path())? + .arg("*** Begin Patch\n*** Update File: foo.txt\n*** End Patch") + .assert() + .failure() + .stderr("Invalid patch hunk on line 2: Update file hunk for path 'foo.txt' is empty\n"); + + Ok(()) +} + +#[test] +fn test_apply_patch_cli_requires_existing_file_for_update() -> anyhow::Result<()> { + let tmp = tempdir()?; + let missing_path = resolved_under(tmp.path(), "missing.txt")?; + + apply_patch_command(tmp.path())? + .arg("*** Begin Patch\n*** Update File: missing.txt\n@@\n-old\n+new\n*** End Patch") + .assert() + .failure() + .stderr(format!( + "Failed to read file to update {}: No such file or directory (os error 2)\n", + missing_path.display() + )); + + Ok(()) +} + +#[test] +fn test_apply_patch_cli_move_overwrites_existing_destination() -> anyhow::Result<()> { + let tmp = tempdir()?; + let original_path = tmp.path().join("old/name.txt"); + let destination = tmp.path().join("renamed/dir/name.txt"); + fs::create_dir_all(original_path.parent().expect("parent should exist"))?; + fs::create_dir_all(destination.parent().expect("parent should exist"))?; + fs::write(&original_path, "from\n")?; + fs::write(&destination, "existing\n")?; + + run_apply_patch_in_dir( + tmp.path(), + "*** Begin Patch\n*** Update File: old/name.txt\n*** Move to: renamed/dir/name.txt\n@@\n-from\n+new\n*** End Patch", + )? + .success() + .stdout("Success. Updated the following files:\nM renamed/dir/name.txt\n"); + + assert!(!original_path.exists()); + assert_eq!(fs::read_to_string(&destination)?, "new\n"); + + Ok(()) +} + +#[test] +fn test_apply_patch_cli_add_overwrites_existing_file() -> anyhow::Result<()> { + let tmp = tempdir()?; + let path = tmp.path().join("duplicate.txt"); + fs::write(&path, "old content\n")?; + + run_apply_patch_in_dir( + tmp.path(), + "*** Begin Patch\n*** Add File: duplicate.txt\n+new content\n*** End Patch", + )? + .success() + .stdout("Success. Updated the following files:\nA duplicate.txt\n"); + + assert_eq!(fs::read_to_string(&path)?, "new content\n"); + + Ok(()) +} + +#[test] +fn test_apply_patch_cli_delete_directory_fails() -> anyhow::Result<()> { + let tmp = tempdir()?; + let dir = tmp.path().join("dir"); + let expected_dir = resolved_under(tmp.path(), "dir")?; + fs::create_dir(&dir)?; + + apply_patch_command(tmp.path())? + .arg("*** Begin Patch\n*** Delete File: dir\n*** End Patch") + .assert() + .failure() + .stderr(format!( + "Failed to delete file {}\n", + expected_dir.display() + )); + + Ok(()) +} + +#[test] +fn test_apply_patch_cli_rejects_invalid_hunk_header() -> anyhow::Result<()> { + let tmp = tempdir()?; + + apply_patch_command(tmp.path())? + .arg("*** Begin Patch\n*** Frobnicate File: foo\n*** End Patch") + .assert() + .failure() + .stderr("Invalid patch hunk on line 2: '*** Frobnicate File: foo' is not a valid hunk header. Valid hunk headers: '*** Add File: {path}', '*** Delete File: {path}', '*** Update File: {path}'\n"); + + Ok(()) +} + +#[test] +fn test_apply_patch_cli_updates_file_appends_trailing_newline() -> anyhow::Result<()> { + let tmp = tempdir()?; + let target_path = tmp.path().join("no_newline.txt"); + fs::write(&target_path, "no newline at end")?; + + run_apply_patch_in_dir( + tmp.path(), + "*** Begin Patch\n*** Update File: no_newline.txt\n@@\n-no newline at end\n+first line\n+second line\n*** End Patch", + )? + .success() + .stdout("Success. Updated the following files:\nM no_newline.txt\n"); + + let contents = fs::read_to_string(&target_path)?; + assert!(contents.ends_with('\n')); + assert_eq!(contents, "first line\nsecond line\n"); + + Ok(()) +} + +#[test] +fn test_apply_patch_cli_failure_after_partial_success_leaves_changes() -> anyhow::Result<()> { + let tmp = tempdir()?; + let new_file = tmp.path().join("created.txt"); + let missing_file = resolved_under(tmp.path(), "missing.txt")?; + + apply_patch_command(tmp.path())? + .arg("*** Begin Patch\n*** Add File: created.txt\n+hello\n*** Update File: missing.txt\n@@\n-old\n+new\n*** End Patch") + .assert() + .failure() + .stdout("") + .stderr(format!( + "Failed to read file to update {}: No such file or directory (os error 2)\n", + missing_file.display() + )); + + assert_eq!(fs::read_to_string(&new_file)?, "hello\n"); + + Ok(()) +} diff --git a/codex-rs/arg0/src/lib.rs b/codex-rs/arg0/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..6b0d99a4372414409e2cdfd220c944b8190e51cc --- /dev/null +++ b/codex-rs/arg0/src/lib.rs @@ -0,0 +1,813 @@ +use std::ffi::OsString; +use std::fs::File; +use std::future::Future; +use std::path::Path; +use std::path::PathBuf; + +use codex_apply_patch::CODEX_CORE_APPLY_PATCH_ARG1; +use codex_async_utils::THREAD_STACK_SIZE_BYTES; +#[cfg(unix)] +use codex_exec_server::CODEX_ARG0_EXEC_HELPER_ARG1; +use codex_exec_server::CODEX_FS_HELPER_ARG1; +use codex_install_context::InstallContext; +use codex_sandboxing::landlock::CODEX_LINUX_SANDBOX_ARG0; +use codex_utils_home_dir::find_codex_home; +#[cfg(target_os = "windows")] +use codex_windows_sandbox::CODEX_WINDOWS_SANDBOX_ARG1; +#[cfg(unix)] +use std::os::unix::fs::symlink; +use tempfile::TempDir; + +const APPLY_PATCH_ARG0: &str = "apply_patch"; +const MISSPELLED_APPLY_PATCH_ARG0: &str = "applypatch"; +#[cfg(unix)] +const EXECVE_WRAPPER_ARG0: &str = "codex-execve-wrapper"; +const LOCK_FILENAME: &str = ".lock"; + +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct Arg0DispatchPaths { + /// Stable path to the current Codex executable for child re-execs. + /// + /// Prefer this over [`std::env::current_exe()`] in code that may run under + /// a test harness, where `current_exe()` can point at the harness binary + /// instead of the real Codex CLI. + pub codex_self_exe: Option, + pub codex_linux_sandbox_exe: Option, + pub main_execve_wrapper_exe: Option, +} + +/// Keeps the per-session PATH entry alive and locked for the process lifetime. +pub struct Arg0PathEntryGuard { + _temp_dir: TempDir, + _lock_file: File, + paths: Arg0DispatchPaths, +} + +impl Arg0PathEntryGuard { + fn new(temp_dir: TempDir, lock_file: File, paths: Arg0DispatchPaths) -> Self { + Self { + _temp_dir: temp_dir, + _lock_file: lock_file, + paths, + } + } + + pub fn paths(&self) -> &Arg0DispatchPaths { + &self.paths + } +} + +pub fn arg0_dispatch() -> Option { + // Determine if we were invoked via the special alias. + let mut args = std::env::args_os(); + let argv0 = args.next().unwrap_or_default(); + let exe_name = Path::new(&argv0) + .file_name() + .and_then(|s| s.to_str()) + .unwrap_or(""); + + #[cfg(unix)] + if exe_name == EXECVE_WRAPPER_ARG0 { + let mut args = std::env::args(); + let _ = args.next(); + let file = match args.next() { + Some(file) => file, + None => std::process::exit(1), + }; + let argv = args.collect::>(); + + let runtime = match tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + { + Ok(runtime) => runtime, + Err(_) => std::process::exit(1), + }; + let exit_code = runtime.block_on( + codex_shell_escalation::run_shell_escalation_execve_wrapper(file, argv), + ); + match exit_code { + Ok(exit_code) => std::process::exit(exit_code), + Err(_) => std::process::exit(1), + } + } + + if exe_name == CODEX_LINUX_SANDBOX_ARG0 { + // Safety: [`run_main`] never returns. + codex_linux_sandbox::run_main(); + } else if exe_name == APPLY_PATCH_ARG0 || exe_name == MISSPELLED_APPLY_PATCH_ARG0 { + codex_apply_patch::main(); + } + + let argv1 = args.next().unwrap_or_default(); + if argv1 == codex_sandboxing::CODEX_WINDOWS_MXC_ARG1 { + codex_sandboxing::run_windows_mxc_main(); + } + #[cfg(unix)] + if argv1 == CODEX_ARG0_EXEC_HELPER_ARG1 { + codex_exec_server::run_arg0_exec_helper_main(); + } + if argv1 == CODEX_FS_HELPER_ARG1 { + codex_exec_server::run_fs_helper_main(); + } + #[cfg(target_os = "windows")] + if argv1 == CODEX_WINDOWS_SANDBOX_ARG1 { + codex_windows_sandbox::run_windows_sandbox_wrapper_main(); + } + if argv1 == CODEX_CORE_APPLY_PATCH_ARG1 { + let patch_arg = args.next().and_then(|s| s.to_str().map(str::to_owned)); + let exit_code = match patch_arg { + Some(patch_arg) => { + let mut stdout = std::io::stdout(); + let mut stderr = std::io::stderr(); + let cwd = match codex_utils_absolute_path::AbsolutePathBuf::current_dir() { + Ok(cwd) => cwd, + Err(_) => std::process::exit(1), + }; + let runtime = match tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + { + Ok(runtime) => runtime, + Err(_) => std::process::exit(1), + }; + let cwd = cwd.into(); + let update_file_mode = codex_apply_patch::apply_patch_file_update_mode_from_env(); + match runtime.block_on(codex_apply_patch::apply_patch_with_options( + &patch_arg, + codex_apply_patch::ApplyPatchOptions { + update_file_mode, + ..Default::default() + }, + &cwd, + &mut stdout, + &mut stderr, + codex_exec_server::LOCAL_FS.as_ref(), + /*sandbox*/ None, + )) { + Ok(_) => 0, + Err(_) => 1, + } + } + None => { + eprintln!("Error: {CODEX_CORE_APPLY_PATCH_ARG1} requires a UTF-8 PATCH argument."); + 1 + } + }; + std::process::exit(exit_code); + } + + // This modifies the environment, which is not thread-safe, so do this + // before creating any threads/the Tokio runtime. + load_dotenv(); + + let (path_entry_guard, updated_path_env_var) = prepare_path_env_var_with_aliases( + InstallContext::current(), + std::env::var_os("PATH"), + prepare_path_entry_for_codex_aliases, + ); + if let Some(updated_path_env_var) = updated_path_env_var { + // It is safe to call set_var() because our process is single-threaded at + // this point in its execution. + unsafe { + std::env::set_var("PATH", updated_path_env_var); + } + } + path_entry_guard +} + +fn prepare_path_env_var_with_aliases( + install_context: &InstallContext, + existing_path: Option, + prepare_aliases: impl FnOnce(Option) -> std::io::Result<(Arg0PathEntryGuard, OsString)>, +) -> (Option, Option) { + let package_path = path_env_with_package_path_dir(install_context, existing_path.clone()); + let path_for_aliases = package_path.clone().or(existing_path); + + match prepare_aliases(path_for_aliases) { + Ok((path_entry, updated_path_env_var)) => (Some(path_entry), Some(updated_path_env_var)), + Err(err) => { + // It is possible that Codex will proceed successfully even if + // creating helper aliases fails, so warn the user and move on. + eprintln!("WARNING: proceeding, even though we could not create PATH aliases: {err}"); + (None, package_path) + } + } +} + +/// While we want to deploy the Codex CLI as a single executable for simplicity, +/// we also want to expose some of its functionality as distinct CLIs, so we use +/// the "arg0 trick" to determine which CLI to dispatch. This effectively allows +/// us to simulate deploying multiple executables as a single binary on Mac and +/// Linux (but not Windows). +/// +/// When the current executable is invoked through the hard-link or alias named +/// `codex-linux-sandbox` we *directly* execute +/// [`codex_linux_sandbox::run_main`] (which never returns). Otherwise we: +/// +/// 1. Load `.env` values from `~/.codex/.env` before creating any threads. +/// 2. Spawn a main runtime thread with a controlled stack size. +/// 3. Construct a Tokio multi-thread runtime. +/// 4. Capture the current executable path and derive the +/// `codex-linux-sandbox` helper path (falling back to the current +/// executable if needed) so children can re-invoke the sandbox when running +/// on Linux. +/// 5. Execute the provided async `main_fn` inside that runtime, forwarding any +/// error. Note that `main_fn` receives [`Arg0DispatchPaths`], which +/// contains the helper executable paths needed to construct +/// [`codex_core::config::Config`]. +/// +/// This function should be used to wrap any `main()` function in binary crates +/// in this workspace that depends on these helper CLIs. +pub fn arg0_dispatch_or_else(main_fn: F) -> anyhow::Result<()> +where + F: FnOnce(Arg0DispatchPaths) -> Fut + Send + 'static, + Fut: Future>, +{ + // Retain the TempDir so it exists for the lifetime of the invocation of + // this executable. Admittedly, we could invoke `keep()` on it, but it + // would be nice to avoid leaving temporary directories behind, if possible. + let path_entry_guard = arg0_dispatch(); + let current_exe = std::env::current_exe().ok(); + + // Regular invocation. Run the async entry point on a thread with the same + // stack budget as Tokio workers; `Runtime::block_on` otherwise runs the + // top-level future on the caller's OS stack. + let handle = std::thread::Builder::new() + .name("codex-main".to_string()) + .stack_size(THREAD_STACK_SIZE_BYTES) + .spawn(move || { + let runtime = build_runtime()?; + runtime.block_on(run_main_with_arg0_guard( + path_entry_guard, + current_exe, + main_fn, + )) + })?; + match handle.join() { + Ok(result) => result, + Err(payload) => std::panic::resume_unwind(payload), + } +} + +async fn run_main_with_arg0_guard( + path_entry_guard: Option, + current_exe: Option, + main_fn: F, +) -> anyhow::Result<()> +where + F: FnOnce(Arg0DispatchPaths) -> Fut, + Fut: Future>, +{ + let paths = Arg0DispatchPaths { + codex_self_exe: current_exe.clone(), + codex_linux_sandbox_exe: if cfg!(target_os = "linux") { + linux_sandbox_exe_path(path_entry_guard.as_ref(), current_exe) + } else { + None + }, + main_execve_wrapper_exe: path_entry_guard + .as_ref() + .and_then(|path_entry| path_entry.paths().main_execve_wrapper_exe.clone()), + }; + + let result = main_fn(paths).await; + // Keep the arg0 tempdir guard alive until the async entry point finishes; + // runtime paths above can point at aliases inside that directory. + drop(path_entry_guard); + result +} + +fn linux_sandbox_exe_path( + path_entry_guard: Option<&Arg0PathEntryGuard>, + current_exe: Option, +) -> Option { + // Prefer the `codex-linux-sandbox` alias when available so callers can + // re-exec through a path whose basename still triggers arg0 dispatch on + // bubblewrap builds that do not support `--argv0`. + path_entry_guard + .and_then(|path_entry| path_entry.paths().codex_linux_sandbox_exe.clone()) + .or(current_exe) +} + +fn build_runtime() -> anyhow::Result { + let mut builder = tokio::runtime::Builder::new_multi_thread(); + builder.enable_all(); + builder.thread_stack_size(THREAD_STACK_SIZE_BYTES); + Ok(builder.build()?) +} + +const ILLEGAL_ENV_VAR_PREFIX: &str = "CODEX_"; + +/// Load env vars from ~/.codex/.env. +/// +/// Security: Do not allow `.env` files to create or modify any variables +/// with names starting with `CODEX_`. +fn load_dotenv() { + if let Ok(codex_home) = find_codex_home() + && let Ok(iter) = dotenvy::from_path_iter(codex_home.join(".env")) + { + set_filtered(iter); + } +} + +/// Helper to set vars from a dotenvy iterator while filtering out `CODEX_` keys. +fn set_filtered(iter: I) +where + I: IntoIterator>, +{ + for (key, value) in iter.into_iter().flatten() { + if !key.to_ascii_uppercase().starts_with(ILLEGAL_ENV_VAR_PREFIX) { + // It is safe to call set_var() because our process is + // single-threaded at this point in its execution. + unsafe { std::env::set_var(&key, &value) }; + } + } +} + +/// Creates a temporary directory with either: +/// +/// - UNIX: `apply_patch` symlink to the current executable +/// - WINDOWS: `apply_patch.bat` batch script to invoke the current executable +/// with the hidden `--codex-run-as-apply-patch` flag. +/// +/// Returns the temporary directory guard and the PATH value that prepends the +/// temporary directory so `apply_patch` can be on the PATH without requiring the +/// user to install a separate executable, simplifying the deployment of Codex +/// CLI. +/// Note: In debug builds the temp-dir guard is disabled to ease local testing. +/// +/// IMPORTANT: Callers must update PATH before multiple threads are spawned. +fn prepare_path_entry_for_codex_aliases( + existing_path: Option, +) -> std::io::Result<(Arg0PathEntryGuard, OsString)> { + let codex_home = find_codex_home()?; + #[cfg(not(debug_assertions))] + { + // Guard against placing helpers in system temp directories outside debug builds. + let temp_root = std::env::temp_dir(); + if codex_home.starts_with(&temp_root) { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "Refusing to create helper binaries under temporary dir {temp_root:?} (codex_home: {codex_home:?})" + ), + )); + } + } + + std::fs::create_dir_all(&codex_home)?; + // Use a CODEX_HOME-scoped temp root to avoid cluttering the top-level directory. + let temp_root = codex_home.join("tmp").join("arg0"); + std::fs::create_dir_all(&temp_root)?; + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + + // Ensure only the current user can access the temp directory. + std::fs::set_permissions(&temp_root, std::fs::Permissions::from_mode(0o700))?; + } + + // Best-effort cleanup of stale per-session dirs. Ignore failures so startup proceeds. + if let Err(err) = janitor_cleanup(&temp_root) { + eprintln!("WARNING: failed to clean up stale arg0 temp dirs: {err}"); + } + + let temp_dir = tempfile::Builder::new() + .prefix("codex-arg0") + .tempdir_in(&temp_root)?; + let path = temp_dir.path(); + + let lock_path = path.join(LOCK_FILENAME); + let lock_file = File::options() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(&lock_path)?; + lock_file.try_lock()?; + + for filename in &[ + APPLY_PATCH_ARG0, + MISSPELLED_APPLY_PATCH_ARG0, + #[cfg(target_os = "linux")] + CODEX_LINUX_SANDBOX_ARG0, + #[cfg(unix)] + EXECVE_WRAPPER_ARG0, + ] { + let exe = std::env::current_exe()?; + + #[cfg(unix)] + { + let link = path.join(filename); + symlink(&exe, &link)?; + } + + #[cfg(windows)] + { + let batch_script = path.join(format!("{filename}.bat")); + let exe = windows_batch_executable_path(&exe, path); + std::fs::write( + &batch_script, + format!( + r#"@echo off +"{exe}" {CODEX_CORE_APPLY_PATCH_ARG1} %* +"#, + ), + )?; + } + } + + let updated_path_env_var = path_env_with_entry(path, existing_path); + + let paths = Arg0DispatchPaths { + codex_self_exe: std::env::current_exe().ok(), + codex_linux_sandbox_exe: { + #[cfg(target_os = "linux")] + { + Some(path.join(CODEX_LINUX_SANDBOX_ARG0)) + } + #[cfg(not(target_os = "linux"))] + { + None + } + }, + main_execve_wrapper_exe: { + #[cfg(unix)] + { + Some(path.join(EXECVE_WRAPPER_ARG0)) + } + #[cfg(not(unix))] + { + None + } + }, + }; + + Ok(( + Arg0PathEntryGuard::new(temp_dir, lock_file, paths), + updated_path_env_var, + )) +} + +#[cfg(windows)] +fn windows_batch_executable_path(executable: &Path, alias_directory: &Path) -> String { + pathdiff::diff_paths(executable, alias_directory) + .filter(|relative_path| relative_path.is_relative()) + .map(|relative_path| format!("%~dp0{}", relative_path.display())) + .unwrap_or_else(|| executable.display().to_string()) +} + +fn path_env_with_package_path_dir( + install_context: &InstallContext, + existing_path: Option, +) -> Option { + let path_dir = install_context + .package_layout + .as_ref() + .and_then(|package_layout| package_layout.path_dir.as_ref())?; + Some(path_env_with_entry(path_dir.as_path(), existing_path)) +} + +fn path_env_with_entry(path_entry: &Path, existing_path: Option) -> OsString { + #[cfg(unix)] + const PATH_SEPARATOR: &str = ":"; + + #[cfg(windows)] + const PATH_SEPARATOR: &str = ";"; + + let capacity = path_entry.as_os_str().len() + + existing_path + .as_ref() + .map_or(0, |existing_path| 1 + existing_path.len()); + let mut path_env_var = OsString::with_capacity(capacity); + path_env_var.push(path_entry); + if let Some(existing_path) = existing_path { + path_env_var.push(PATH_SEPARATOR); + path_env_var.push(existing_path); + } + path_env_var +} + +fn janitor_cleanup(temp_root: &Path) -> std::io::Result<()> { + let entries = match std::fs::read_dir(temp_root) { + Ok(entries) => entries, + Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(()), + Err(err) => return Err(err), + }; + + for entry in entries.flatten() { + let path = entry.path(); + if !path.is_dir() { + continue; + } + + // Skip the directory if locking fails or the lock is currently held. + let Some(_lock_file) = try_lock_dir(&path)? else { + continue; + }; + + match std::fs::remove_dir_all(&path) { + Ok(()) => {} + // Expected TOCTOU race: directory can disappear after read_dir/lock checks. + Err(err) if err.kind() == std::io::ErrorKind::NotFound => continue, + Err(err) => return Err(err), + } + } + + Ok(()) +} + +fn try_lock_dir(dir: &Path) -> std::io::Result> { + let lock_path = dir.join(LOCK_FILENAME); + let lock_file = match File::options().read(true).write(true).open(&lock_path) { + Ok(file) => file, + Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(err) => return Err(err), + }; + + match lock_file.try_lock() { + Ok(()) => Ok(Some(lock_file)), + Err(std::fs::TryLockError::WouldBlock) => Ok(None), + Err(err) => Err(err.into()), + } +} + +#[cfg(test)] +mod tests { + use super::Arg0DispatchPaths; + use super::Arg0PathEntryGuard; + use super::LOCK_FILENAME; + use super::janitor_cleanup; + use super::linux_sandbox_exe_path; + #[cfg(unix)] + use super::run_main_with_arg0_guard; + #[cfg(unix)] + use anyhow::ensure; + use codex_install_context::CodexPackageLayout; + use codex_install_context::InstallContext; + use codex_install_context::InstallMethod; + use codex_utils_absolute_path::AbsolutePathBuf; + use pretty_assertions::assert_eq; + use std::fs; + use std::fs::File; + #[cfg(windows)] + use std::os::windows::process::CommandExt; + use std::path::Path; + use std::path::PathBuf; + use tempfile::TempDir; + + struct PackagePathTestFixture { + _temp_dir: TempDir, + arg0_dir: PathBuf, + existing_dir: PathBuf, + install_context: InstallContext, + path_dir: AbsolutePathBuf, + } + + fn create_lock(dir: &Path) -> std::io::Result { + let lock_path = dir.join(LOCK_FILENAME); + File::options() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(lock_path) + } + + fn package_path_test_fixture() -> anyhow::Result { + let temp_dir = TempDir::new()?; + let arg0_dir = temp_dir.path().join("arg0"); + let package_dir = temp_dir.path().join("package"); + let bin_dir = package_dir.join("bin"); + let path_dir = package_dir.join("codex-path"); + let existing_dir = temp_dir.path().join("existing-bin"); + fs::create_dir_all(&arg0_dir)?; + fs::create_dir_all(&bin_dir)?; + fs::create_dir_all(&path_dir)?; + fs::create_dir_all(&existing_dir)?; + let path_dir = AbsolutePathBuf::from_absolute_path(path_dir.canonicalize()?)?; + let install_context = InstallContext { + method: InstallMethod::Other, + package_layout: Some(CodexPackageLayout { + package_dir: AbsolutePathBuf::from_absolute_path(package_dir.canonicalize()?)?, + bin_dir: AbsolutePathBuf::from_absolute_path(bin_dir.canonicalize()?)?, + resources_dir: None, + path_dir: Some(path_dir.clone()), + }), + }; + + Ok(PackagePathTestFixture { + _temp_dir: temp_dir, + arg0_dir, + existing_dir, + install_context, + path_dir, + }) + } + + #[cfg(windows)] + #[test] + fn windows_batch_alias_preserves_unicode_executable_paths() -> anyhow::Result<()> { + let root = TempDir::new()?; + let profile = root.path().join("用户"); + let alias_directory = profile.join(".codex").join("tmp").join("arg0"); + let executable_directory = profile.join("bin"); + fs::create_dir_all(&alias_directory)?; + fs::create_dir_all(&executable_directory)?; + + let system_root = std::env::var_os("SystemRoot") + .ok_or_else(|| anyhow::anyhow!("missing Windows system root"))?; + let command_shell = PathBuf::from(system_root).join("System32").join("cmd.exe"); + let executable = executable_directory.join("cmd.exe"); + fs::copy(&command_shell, &executable)?; + + let batch_path = alias_directory.join("apply_patch.bat"); + let executable_path = super::windows_batch_executable_path(&executable, &alias_directory); + fs::write( + &batch_path, + format!("@echo off\r\n\"{executable_path}\" /d /c exit 37\r\n"), + )?; + + let output = std::process::Command::new(command_shell) + .args(["/d", "/c"]) + .raw_arg(format!("chcp 437>nul & call \"{}\"", batch_path.display())) + .output()?; + + assert_eq!(output.status.code(), Some(37)); + Ok(()) + } + + #[cfg(windows)] + #[test] + fn windows_batch_alias_preserves_cross_volume_executable_paths() { + assert_eq!( + super::windows_batch_executable_path( + Path::new(r"D:\Tools\codex.exe"), + Path::new(r"C:\Users\person\.codex\tmp\arg0"), + ), + r"D:\Tools\codex.exe", + ); + } + + #[test] + fn linux_sandbox_exe_path_prefers_codex_linux_sandbox_alias() -> std::io::Result<()> { + let temp_dir = TempDir::new()?; + let lock_file = create_lock(temp_dir.path())?; + let alias_path = temp_dir.path().join("codex-linux-sandbox"); + let path_entry = Arg0PathEntryGuard::new( + temp_dir, + lock_file, + Arg0DispatchPaths { + codex_self_exe: Some(PathBuf::from("/usr/bin/codex")), + codex_linux_sandbox_exe: Some(alias_path.clone()), + main_execve_wrapper_exe: None, + }, + ); + + assert_eq!( + linux_sandbox_exe_path(Some(&path_entry), Some(PathBuf::from("/usr/bin/codex"))), + Some(alias_path), + ); + Ok(()) + } + + #[test] + fn path_env_can_prepend_package_path_before_arg0_alias_dir() -> anyhow::Result<()> { + let fixture = package_path_test_fixture()?; + + let package_path = super::path_env_with_package_path_dir( + &fixture.install_context, + Some(fixture.existing_dir.as_os_str().to_owned()), + ) + .expect("package path dir should update PATH"); + let updated_path = super::path_env_with_entry(&fixture.arg0_dir, Some(package_path)); + + assert_eq!( + std::env::split_paths(&updated_path).collect::>(), + vec![ + fixture.arg0_dir, + fixture.path_dir.as_path().to_path_buf(), + fixture.existing_dir + ], + ); + Ok(()) + } + + #[test] + fn package_path_survives_arg0_alias_setup_failure() -> anyhow::Result<()> { + let fixture = package_path_test_fixture()?; + + let (path_entry_guard, updated_path_env_var) = super::prepare_path_env_var_with_aliases( + &fixture.install_context, + Some(fixture.existing_dir.as_os_str().to_owned()), + |path_for_aliases| { + assert_eq!( + std::env::split_paths( + &path_for_aliases.expect("package PATH should be passed to alias setup") + ) + .collect::>(), + vec![ + fixture.path_dir.as_path().to_path_buf(), + fixture.existing_dir.clone() + ], + ); + Err(std::io::Error::other("alias setup failed")) + }, + ); + + assert!(path_entry_guard.is_none()); + let updated_path_env_var = + updated_path_env_var.expect("package PATH should survive alias setup failure"); + assert_eq!( + std::env::split_paths(&updated_path_env_var).collect::>(), + vec![ + fixture.path_dir.as_path().to_path_buf(), + fixture.existing_dir + ], + ); + Ok(()) + } + + #[cfg(unix)] + #[test] + fn run_main_with_arg0_guard_keeps_aliases_alive_until_main_returns() -> anyhow::Result<()> { + let temp_dir = TempDir::new()?; + let alias_path = temp_dir.path().join("codex-helper-alias"); + fs::write(&alias_path, b"")?; + let lock_file = create_lock(temp_dir.path())?; + let path_entry = Arg0PathEntryGuard::new( + temp_dir, + lock_file, + Arg0DispatchPaths { + codex_self_exe: Some(PathBuf::from("/usr/bin/codex")), + codex_linux_sandbox_exe: Some(alias_path.clone()), + main_execve_wrapper_exe: Some(alias_path), + }, + ); + + super::build_runtime()?.block_on(run_main_with_arg0_guard( + /*path_entry_guard*/ Some(path_entry), + Some(PathBuf::from("/usr/bin/codex")), + |paths| async move { + let alias_path = paths + .codex_linux_sandbox_exe + .or(paths.main_execve_wrapper_exe) + .expect("unix dispatch should create at least one alias path"); + ensure!( + alias_path.exists(), + "alias path disappeared before main future was polled: {}", + alias_path.display() + ); + + tokio::task::yield_now().await; + + ensure!( + alias_path.exists(), + "alias path disappeared while main future was running: {}", + alias_path.display() + ); + Ok(()) + }, + )) + } + + #[test] + fn janitor_skips_dirs_without_lock_file() -> std::io::Result<()> { + let root = tempfile::tempdir()?; + let dir = root.path().join("no-lock"); + fs::create_dir(&dir)?; + + janitor_cleanup(root.path())?; + + assert!(dir.exists()); + Ok(()) + } + + #[test] + fn janitor_skips_dirs_with_held_lock() -> std::io::Result<()> { + let root = tempfile::tempdir()?; + let dir = root.path().join("locked"); + fs::create_dir(&dir)?; + let lock_file = create_lock(&dir)?; + lock_file.try_lock()?; + + janitor_cleanup(root.path())?; + + assert!(dir.exists()); + Ok(()) + } + + #[test] + fn janitor_removes_dirs_with_unlocked_lock() -> std::io::Result<()> { + let root = tempfile::tempdir()?; + let dir = root.path().join("stale"); + fs::create_dir(&dir)?; + create_lock(&dir)?; + + janitor_cleanup(root.path())?; + + assert!(!dir.exists()); + Ok(()) + } +} diff --git a/codex-rs/async-utils/src/lib.rs b/codex-rs/async-utils/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..16a86229ee73b63f8fc0b74e0700c2f0104d57c9 --- /dev/null +++ b/codex-rs/async-utils/src/lib.rs @@ -0,0 +1,89 @@ +use std::future::Future; +use tokio_util::sync::CancellationToken; + +/// Stack budget for threads that poll Codex async work. +pub const THREAD_STACK_SIZE_BYTES: usize = 16 * 1024 * 1024; + +#[derive(Debug, PartialEq, Eq)] +pub enum CancelErr { + Cancelled, +} + +pub trait OrCancelExt: Sized { + type Output; + + fn or_cancel( + self, + token: &CancellationToken, + ) -> impl Future> + Send; +} + +impl OrCancelExt for F +where + F: Future + Send, + F::Output: Send, +{ + type Output = F::Output; + + async fn or_cancel(self, token: &CancellationToken) -> Result { + tokio::select! { + _ = token.cancelled() => Err(CancelErr::Cancelled), + res = self => Ok(res), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + use std::time::Duration; + use tokio::task; + use tokio::time::sleep; + + #[tokio::test] + async fn returns_ok_when_future_completes_first() { + let token = CancellationToken::new(); + let value = async { 42 }; + + let result = value.or_cancel(&token).await; + + assert_eq!(Ok(42), result); + } + + #[tokio::test] + async fn returns_err_when_token_cancelled_first() { + let token = CancellationToken::new(); + let token_clone = token.clone(); + + let cancel_handle = task::spawn(async move { + sleep(Duration::from_millis(10)).await; + token_clone.cancel(); + }); + + let result = async { + sleep(Duration::from_millis(100)).await; + 7 + } + .or_cancel(&token) + .await; + + cancel_handle.await.expect("cancel task panicked"); + assert_eq!(Err(CancelErr::Cancelled), result); + } + + #[tokio::test] + async fn returns_err_when_token_already_cancelled() { + let token = CancellationToken::new(); + token.cancel(); + + let result = async { + sleep(Duration::from_millis(50)).await; + 5 + } + .or_cancel(&token) + .await; + + assert_eq!(Err(CancelErr::Cancelled), result); + } +} diff --git a/codex-rs/chatgpt/src/apply_command.rs b/codex-rs/chatgpt/src/apply_command.rs new file mode 100644 index 0000000000000000000000000000000000000000..70fe4481db70c4bb0f71a89daff499ca5c7bb27b --- /dev/null +++ b/codex-rs/chatgpt/src/apply_command.rs @@ -0,0 +1,77 @@ +use std::path::PathBuf; + +use clap::Parser; +use codex_core::config::Config; +use codex_git_utils::ApplyGitRequest; +use codex_git_utils::apply_git_patch; +use codex_utils_cli::CliConfigOverrides; + +use crate::get_task::GetTaskResponse; +use crate::get_task::OutputItem; +use crate::get_task::PrOutputItem; +use crate::get_task::get_task; + +/// Applies the latest diff from a Codex agent task. +#[derive(Debug, Parser)] +pub struct ApplyCommand { + pub task_id: String, + + #[clap(flatten)] + pub config_overrides: CliConfigOverrides, +} +pub async fn run_apply_command( + apply_cli: ApplyCommand, + cwd: Option, +) -> anyhow::Result<()> { + let config = Config::load_with_cli_overrides( + apply_cli + .config_overrides + .parse_overrides() + .map_err(anyhow::Error::msg)?, + ) + .await?; + + let task_response = get_task(&config, apply_cli.task_id).await?; + apply_diff_from_task(task_response, cwd).await +} + +pub async fn apply_diff_from_task( + task_response: GetTaskResponse, + cwd: Option, +) -> anyhow::Result<()> { + let diff_turn = match task_response.current_diff_task_turn { + Some(turn) => turn, + None => anyhow::bail!("No diff turn found"), + }; + let output_diff = diff_turn.output_items.iter().find_map(|item| match item { + OutputItem::Pr(PrOutputItem { output_diff }) => Some(output_diff), + _ => None, + }); + match output_diff { + Some(output_diff) => apply_diff(&output_diff.diff, cwd).await, + None => anyhow::bail!("No PR output item found"), + } +} + +async fn apply_diff(diff: &str, cwd: Option) -> anyhow::Result<()> { + let cwd = cwd.unwrap_or(std::env::current_dir().unwrap_or_else(|_| std::env::temp_dir())); + let req = ApplyGitRequest { + cwd, + diff: diff.to_string(), + revert: false, + preflight: false, + }; + let res = apply_git_patch(&req)?; + if res.exit_code != 0 { + anyhow::bail!( + "Git apply failed (applied={}, skipped={}, conflicts={})\nstdout:\n{}\nstderr:\n{}", + res.applied_paths.len(), + res.skipped_paths.len(), + res.conflicted_paths.len(), + res.stdout, + res.stderr + ); + } + println!("Successfully applied diff"); + Ok(()) +} diff --git a/codex-rs/chatgpt/src/chatgpt_client.rs b/codex-rs/chatgpt/src/chatgpt_client.rs new file mode 100644 index 0000000000000000000000000000000000000000..2e6feb8256e29d85f3ed24b643a7d081321e270f --- /dev/null +++ b/codex-rs/chatgpt/src/chatgpt_client.rs @@ -0,0 +1,176 @@ +use codex_core::config::Config; +use codex_http_client::HttpClient; +use codex_http_client::HttpClientFactory; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_login::default_client::RESIDENCY_HEADER_NAME; +use codex_login::default_client::create_client; +use codex_login::default_client::create_client_with_chatgpt_cookies; +use codex_login::default_client::default_headers; + +use anyhow::Context; +use serde::Serialize; +use serde::de::DeserializeOwned; +use std::sync::LazyLock; +use std::sync::Mutex; +use std::time::Duration; + +const OAI_PRODUCT_SKU_HEADER: &str = "OAI-Product-Sku"; +const CODEX_PRODUCT_SKU: &str = "codex"; + +struct CachedChatGptClient { + factory: HttpClientFactory, + residency: Option>, + client: HttpClient, +} + +static PSP_CHATGPT_CLIENT: LazyLock>> = + LazyLock::new(|| Mutex::new(None)); + +/// Reuse the default client while retaining its configured ChatGPT cookies. +fn psp_chatgpt_client(factory: HttpClientFactory) -> HttpClient { + let residency = default_headers() + .get(RESIDENCY_HEADER_NAME) + .map(|value| value.as_bytes().to_vec()); + let mut cached = PSP_CHATGPT_CLIENT + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(cached_client) = cached.as_ref() + && cached_client.factory == factory + && cached_client.residency == residency + { + return cached_client.client.clone(); + } + + let client = create_client_with_chatgpt_cookies(&factory); + *cached = Some(CachedChatGptClient { + factory, + residency, + client: client.clone(), + }); + client +} + +/// Make a GET request to the ChatGPT backend API. +pub(crate) async fn chatgpt_get_request( + config: &Config, + path: String, +) -> anyhow::Result { + chatgpt_get_request_with_timeout(config, path, /*timeout*/ None).await +} + +pub(crate) async fn chatgpt_get_request_with_timeout( + config: &Config, + path: String, + timeout: Option, +) -> anyhow::Result { + let chatgpt_base_url = &config.chatgpt_base_url; + let auth_manager = + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await?; + let auth = auth_manager + .auth() + .await + .ok_or_else(|| anyhow::anyhow!("ChatGPT auth not available"))?; + anyhow::ensure!( + auth.uses_codex_backend(), + "ChatGPT backend requests require Codex backend auth" + ); + anyhow::ensure!( + auth.get_account_id().is_some(), + "ChatGPT account ID not available, please re-run `codex login`" + ); + + let url = format!( + "{}/{}", + chatgpt_base_url.trim_end_matches('/'), + path.trim_start_matches('/') + ); + + let http_client_factory = config.http_client_factory(); + let client = if http_client_factory.has_chatgpt_cookies() { + psp_chatgpt_client(http_client_factory) + } else { + create_client() + }; + let mut request = client + .get(&url) + .headers(default_headers()) + .headers(codex_model_provider::auth_provider_from_auth(&auth).to_auth_headers()) + .header(OAI_PRODUCT_SKU_HEADER, CODEX_PRODUCT_SKU) + .header("Content-Type", "application/json"); + if let Some(timeout) = timeout { + request = request.timeout(timeout); + } + let response = request.send().await.context("Failed to send request")?; + + if response.status().is_success() { + let result: T = response + .json() + .await + .context("Failed to parse JSON response")?; + Ok(result) + } else { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + anyhow::bail!("Request failed with status {status}: {body}") + } +} + +/// Make a POST request to the ChatGPT backend API with an already-captured auth identity. +/// +/// Callers that bind other state to the auth snapshot should pass that same snapshot here rather +/// than reacquiring auth while the request is in flight. +pub(crate) async fn chatgpt_post_request_with_timeout< + TResponse: DeserializeOwned, + TRequest: Serialize + ?Sized, +>( + config: &Config, + auth: &CodexAuth, + path: String, + body: &TRequest, + timeout: Duration, + product_sku: &str, +) -> anyhow::Result { + anyhow::ensure!( + auth.uses_codex_backend(), + "ChatGPT backend requests require Codex backend auth" + ); + anyhow::ensure!( + auth.get_account_id().is_some(), + "ChatGPT account ID not available, please re-run codex login" + ); + + let url = format!( + "{}/{}", + config.chatgpt_base_url.trim_end_matches('/'), + path.trim_start_matches('/') + ); + let http_client_factory = config.http_client_factory(); + let client = if http_client_factory.has_chatgpt_cookies() { + psp_chatgpt_client(http_client_factory) + } else { + create_client() + }; + let response = client + .post(&url) + .headers(default_headers()) + .headers(codex_model_provider::auth_provider_from_auth(auth).to_auth_headers()) + .header(OAI_PRODUCT_SKU_HEADER, product_sku) + .header("Content-Type", "application/json") + .timeout(timeout) + .json(body) + .send() + .await + .context("Failed to send request")?; + + if response.status().is_success() { + response + .json() + .await + .context("Failed to parse JSON response") + } else { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + anyhow::bail!("Request failed with status {status}: {body}") + } +} diff --git a/codex-rs/chatgpt/src/connectors.rs b/codex-rs/chatgpt/src/connectors.rs new file mode 100644 index 0000000000000000000000000000000000000000..ae6cbce5c8110a4bab6b0dd807ff39d02e43a048 --- /dev/null +++ b/codex-rs/chatgpt/src/connectors.rs @@ -0,0 +1,505 @@ +use std::collections::HashMap; +use std::collections::HashSet; +use std::time::Duration; + +use crate::chatgpt_client::chatgpt_get_request_with_timeout; +use crate::chatgpt_client::chatgpt_post_request_with_timeout; + +use codex_connectors::AppInfo; +use codex_connectors::AppToolPolicyEvaluator; +use codex_connectors::ConnectorDirectoryCacheContext; +use codex_connectors::ConnectorDirectoryCacheKey; +use codex_connectors::ConnectorMetadata; +use codex_connectors::ConnectorMetadataStore; +use codex_connectors::ConnectorToolSummary; +use codex_connectors::DirectoryListResponse; +use codex_connectors::merge::merge_connectors; +use codex_connectors::merge::merge_plugin_connectors; +use codex_core::config::Config; +pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools; +pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_environment_manager; +pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_mcp_manager; +pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_options; +pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_options_and_status; +pub use codex_core::connectors::list_cached_accessible_connectors_from_mcp_tools; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_plugin::AppConnectorId; +use serde::Deserialize; +use serde::Serialize; + +const DIRECTORY_CONNECTORS_TIMEOUT: Duration = Duration::from_secs(60); +const CONNECTOR_METADATA_TIMEOUT: Duration = Duration::from_secs(60); +const DEFAULT_APPS_PRODUCT_SKU: &str = "codex"; + +async fn apps_enabled(config: &Config) -> anyhow::Result { + let auth_manager = + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await?; + let auth = auth_manager.auth().await; + Ok(config + .features + .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend))) +} + +async fn connector_auth(config: &Config) -> anyhow::Result { + let auth_manager = + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await?; + let auth = auth_manager + .auth() + .await + .ok_or_else(|| anyhow::anyhow!("ChatGPT auth not available"))?; + anyhow::ensure!( + auth.uses_codex_backend(), + "ChatGPT connectors require Codex backend auth" + ); + Ok(auth) +} + +pub async fn list_connectors(config: &Config) -> anyhow::Result> { + if !apps_enabled(config).await? { + return Ok(Vec::new()); + } + let (connectors_result, accessible_result) = tokio::join!( + list_all_connectors(config), + list_accessible_connectors_from_mcp_tools(config), + ); + let connectors = connectors_result?; + let accessible = accessible_result?; + Ok( + AppToolPolicyEvaluator::new(&config.config_layer_stack).apply_app_enabled_state( + merge_connectors_with_accessible( + connectors, accessible, /*all_connectors_loaded*/ true, + ), + ), + ) +} + +pub async fn list_all_connectors(config: &Config) -> anyhow::Result> { + list_all_connectors_with_options(config, /*force_refetch*/ false, &[]).await +} + +pub async fn list_cached_all_connectors( + config: &Config, + plugin_apps: &[AppConnectorId], +) -> Option> { + if !apps_enabled(config).await.ok()? { + return Some(Vec::new()); + } + + let auth = connector_auth(config).await.ok()?; + let cache_context = connector_directory_cache_context(config, &auth); + let connectors = codex_connectors::cached_directory_connectors(&cache_context)?; + Some(merge_directory_and_plugin_connectors( + connectors, + plugin_apps, + )) +} + +pub async fn list_all_connectors_with_options( + config: &Config, + force_refetch: bool, + plugin_apps: &[AppConnectorId], +) -> anyhow::Result> { + if !apps_enabled(config).await? { + return Ok(Vec::new()); + } + let auth = connector_auth(config).await?; + let cache_context = connector_directory_cache_context(config, &auth); + let connectors = codex_connectors::list_all_connectors_with_options( + cache_context, + auth.is_workspace_account(), + force_refetch, + |path| async move { + chatgpt_get_request_with_timeout::( + config, + path, + Some(DIRECTORY_CONNECTORS_TIMEOUT), + ) + .await + }, + ) + .await?; + Ok(merge_directory_and_plugin_connectors( + connectors, + plugin_apps, + )) +} + +pub struct ConnectorMetadataReadResult { + pub apps: Vec, + pub missing_app_ids: Vec, +} + +/// Reads display metadata without loading MCP connector tools or runtime state. +/// +/// The store is created before awaiting the backend request, so a response that arrives after an +/// account or backend change can only commit to the scope under which it was requested. +pub async fn read_connector_metadata( + config: &Config, + auth: &CodexAuth, + app_ids: &[String], + include_tools: bool, +) -> anyhow::Result { + anyhow::ensure!( + auth.uses_codex_backend(), + "ChatGPT backend requests require Codex backend auth" + ); + anyhow::ensure!( + auth.get_account_id().is_some(), + "ChatGPT account ID not available, please re-run codex login" + ); + + let store = ConnectorMetadataStore::new( + config.chatgpt_base_url.clone(), + auth.get_account_id(), + auth.get_chatgpt_user_id(), + auth.is_workspace_account(), + ); + let mut metadata_by_id = store.fresh_records(app_ids, include_tools); + let missing_ids = app_ids + .iter() + .filter(|app_id| !metadata_by_id.contains_key(app_id.as_str())) + .cloned() + .collect::>(); + + if !missing_ids.is_empty() { + let product_sku = config + .apps_mcp_product_sku + .as_deref() + .unwrap_or(DEFAULT_APPS_PRODUCT_SKU); + let response: GetAppsResponse = chatgpt_post_request_with_timeout( + config, + auth, + "/ps/apps/batch".to_string(), + &GetAppsRequest { + app_ids: &missing_ids, + include_tools, + }, + CONNECTOR_METADATA_TIMEOUT, + product_sku, + ) + .await?; + let mut requested_ids = missing_ids.iter().cloned().collect::>(); + let fetched = response + .apps + .into_iter() + .map(batch_app_to_metadata) + .filter(|metadata| requested_ids.remove(&metadata.id)) + .collect::>(); + store.commit(&fetched); + metadata_by_id.extend( + fetched + .into_iter() + .map(|metadata| (metadata.id.clone(), metadata)), + ); + } + + let mut apps = Vec::new(); + let mut missing_app_ids = Vec::new(); + for app_id in app_ids { + if let Some(mut metadata) = metadata_by_id.remove(app_id) { + if !include_tools { + metadata.tool_summaries = None; + } + apps.push(metadata); + } else { + missing_app_ids.push(app_id.clone()); + } + } + + Ok(ConnectorMetadataReadResult { + apps, + missing_app_ids, + }) +} + +#[derive(Serialize)] +struct GetAppsRequest<'a> { + app_ids: &'a [String], + include_tools: bool, +} + +#[derive(Deserialize)] +struct GetAppsResponse { + apps: Vec, +} + +/// The explicit metadata-only projection of Plugin Service's public app response. +/// +/// Serde ignores all other backend fields, including full actions, model descriptions, and +/// runtime state. +#[derive(Deserialize)] +struct BatchApp { + id: String, + name: String, + description: Option, + icon_url: Option, + #[serde(default, rename = "icon_dark_url", alias = "icon_url_dark")] + icon_url_dark: Option, + #[serde(default)] + distribution_channel: Option, + #[serde(default)] + tools: Option>, +} + +#[derive(Deserialize)] +struct BatchAppToolSummary { + name: String, + title: Option, + description: String, + #[serde(default)] + is_enabled: Option, + #[serde(default)] + disabled_reason: Option, + #[serde(default)] + is_read_only: bool, +} + +fn batch_app_to_metadata(app: BatchApp) -> ConnectorMetadata { + let BatchApp { + id, + name, + description, + icon_url, + icon_url_dark, + distribution_channel, + tools, + } = app; + ConnectorMetadata { + id, + name, + description, + icon_url, + icon_url_dark, + distribution_channel, + tool_summaries: tools.map(|tools| { + tools + .into_iter() + .map(|tool| { + let BatchAppToolSummary { + name, + title, + description, + is_enabled, + disabled_reason, + is_read_only, + } = tool; + ConnectorToolSummary { + name, + title, + description, + is_enabled: is_enabled.unwrap_or(true), + disabled_reason, + is_read_only, + } + }) + .collect() + }), + } +} + +fn connector_directory_cache_context( + config: &Config, + auth: &CodexAuth, +) -> ConnectorDirectoryCacheContext { + ConnectorDirectoryCacheContext::new( + config.codex_home.to_path_buf(), + ConnectorDirectoryCacheKey::new( + config.chatgpt_base_url.clone(), + auth.get_account_id(), + auth.get_chatgpt_user_id(), + auth.is_workspace_account(), + ), + ) +} + +fn merge_directory_and_plugin_connectors( + connectors: Vec, + plugin_apps: &[AppConnectorId], +) -> Vec { + merge_plugin_connectors( + connectors, + plugin_apps + .iter() + .map(|connector_id| connector_id.0.clone()), + ) +} + +pub fn connectors_for_plugin_apps( + connectors: Vec, + plugin_apps: &[AppConnectorId], +) -> Vec { + let connectors = merge_plugin_connectors( + connectors, + plugin_apps + .iter() + .map(|connector_id| connector_id.0.clone()), + ); + let mut connectors_by_id = connectors + .into_iter() + .map(|connector| (connector.id.clone(), connector)) + .collect::>(); + + plugin_apps + .iter() + .filter_map(|connector_id| connectors_by_id.remove(connector_id.0.as_str())) + .collect() +} + +pub fn merge_connectors_with_accessible( + connectors: Vec, + accessible_connectors: Vec, + all_connectors_loaded: bool, +) -> Vec { + let accessible_connectors = if all_connectors_loaded { + let connector_ids: HashSet<&str> = connectors + .iter() + .map(|connector| connector.id.as_str()) + .collect(); + accessible_connectors + .into_iter() + .filter(|connector| connector_ids.contains(connector.id.as_str())) + .collect() + } else { + accessible_connectors + }; + merge_connectors(connectors, accessible_connectors) +} + +#[cfg(test)] +mod tests { + use super::*; + use codex_connectors::metadata::connector_install_url; + use codex_plugin::AppConnectorId; + use pretty_assertions::assert_eq; + use serde_json::json; + + #[test] + fn batch_app_accepts_missing_optional_metadata() { + let app = serde_json::from_value::(json!({ + "id": "alpha", + "name": "Alpha", + "description": "Alpha description", + "icon_url": null, + "tools": [{ + "name": "search", + "title": "Search", + "description": "Search Alpha", + }], + })) + .expect("valid legacy batch app"); + + assert_eq!( + batch_app_to_metadata(app), + ConnectorMetadata { + id: "alpha".to_string(), + name: "Alpha".to_string(), + description: Some("Alpha description".to_string()), + icon_url: None, + icon_url_dark: None, + distribution_channel: None, + tool_summaries: Some(vec![ConnectorToolSummary { + name: "search".to_string(), + title: Some("Search".to_string()), + description: "Search Alpha".to_string(), + is_enabled: true, + disabled_reason: None, + is_read_only: false, + }]), + } + ); + } + + fn app(id: &str) -> AppInfo { + AppInfo { + id: id.to_string(), + name: id.to_string(), + description: None, + logo_url: None, + logo_url_dark: None, + icon_assets: None, + icon_dark_assets: None, + distribution_channel: None, + branding: None, + app_metadata: None, + labels: None, + install_url: None, + is_accessible: false, + is_enabled: true, + plugin_display_names: Vec::new(), + } + } + + fn merged_app(id: &str, is_accessible: bool) -> AppInfo { + AppInfo { + id: id.to_string(), + name: id.to_string(), + description: None, + logo_url: None, + logo_url_dark: None, + icon_assets: None, + icon_dark_assets: None, + distribution_channel: None, + branding: None, + app_metadata: None, + labels: None, + install_url: Some(connector_install_url(id, id)), + is_accessible, + is_enabled: true, + plugin_display_names: Vec::new(), + } + } + + #[test] + fn excludes_accessible_connectors_not_in_all_when_all_loaded() { + let merged = merge_connectors_with_accessible( + vec![app("alpha")], + vec![app("alpha"), app("beta")], + /*all_connectors_loaded*/ true, + ); + assert_eq!(merged, vec![merged_app("alpha", /*is_accessible*/ true)]); + } + + #[test] + fn keeps_accessible_connectors_not_in_all_while_all_loading() { + let merged = merge_connectors_with_accessible( + vec![app("alpha")], + vec![app("alpha"), app("beta")], + /*all_connectors_loaded*/ false, + ); + assert_eq!( + merged, + vec![ + merged_app("alpha", /*is_accessible*/ true), + merged_app("beta", /*is_accessible*/ true) + ] + ); + } + + #[test] + fn connectors_for_plugin_apps_returns_only_requested_plugin_apps() { + let connectors = connectors_for_plugin_apps( + vec![app("alpha"), app("beta")], + &[ + AppConnectorId("gmail".to_string()), + AppConnectorId("alpha".to_string()), + AppConnectorId("gmail".to_string()), + ], + ); + assert_eq!( + connectors, + vec![merged_app("gmail", /*is_accessible*/ false), app("alpha")] + ); + } + + #[test] + fn connectors_for_plugin_apps_preserves_formerly_disallowed_plugin_apps() { + let connector_id = "asdk_app_6938a94a61d881918ef32cb999ff937c"; + let connectors = + connectors_for_plugin_apps(Vec::new(), &[AppConnectorId(connector_id.to_string())]); + assert_eq!( + connectors, + vec![merged_app(connector_id, /*is_accessible*/ false)] + ); + } +} diff --git a/codex-rs/chatgpt/src/get_task.rs b/codex-rs/chatgpt/src/get_task.rs new file mode 100644 index 0000000000000000000000000000000000000000..9301ffc38d32fe6d67089534ef775597167fffeb --- /dev/null +++ b/codex-rs/chatgpt/src/get_task.rs @@ -0,0 +1,40 @@ +use codex_core::config::Config; +use serde::Deserialize; + +use crate::chatgpt_client::chatgpt_get_request; + +#[derive(Debug, Deserialize)] +pub struct GetTaskResponse { + pub current_diff_task_turn: Option, +} + +// Only relevant fields for our extraction +#[derive(Debug, Deserialize)] +pub struct AssistantTurn { + pub output_items: Vec, +} + +#[derive(Debug, Deserialize)] +#[serde(tag = "type")] +pub enum OutputItem { + #[serde(rename = "pr")] + Pr(PrOutputItem), + + #[serde(other)] + Other, +} + +#[derive(Debug, Deserialize)] +pub struct PrOutputItem { + pub output_diff: OutputDiff, +} + +#[derive(Debug, Deserialize)] +pub struct OutputDiff { + pub diff: String, +} + +pub(crate) async fn get_task(config: &Config, task_id: String) -> anyhow::Result { + let path = format!("/wham/tasks/{task_id}"); + chatgpt_get_request(config, path).await +} diff --git a/codex-rs/chatgpt/src/lib.rs b/codex-rs/chatgpt/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..057478db1856c84a88eff8353d652be7280b8ae3 --- /dev/null +++ b/codex-rs/chatgpt/src/lib.rs @@ -0,0 +1,4 @@ +pub mod apply_command; +mod chatgpt_client; +pub mod connectors; +pub mod get_task; diff --git a/codex-rs/chatgpt/tests/all.rs b/codex-rs/chatgpt/tests/all.rs new file mode 100644 index 0000000000000000000000000000000000000000..7e136e4cce2a99915ff1b487d1f2b61c500f133a --- /dev/null +++ b/codex-rs/chatgpt/tests/all.rs @@ -0,0 +1,3 @@ +// Single integration test binary that aggregates all test modules. +// The submodules live in `tests/suite/`. +mod suite; diff --git a/codex-rs/chatgpt/tests/suite/apply_command_e2e.rs b/codex-rs/chatgpt/tests/suite/apply_command_e2e.rs new file mode 100644 index 0000000000000000000000000000000000000000..c2d570528ce87db7cacca9f17e4ad4d16e8f40ff --- /dev/null +++ b/codex-rs/chatgpt/tests/suite/apply_command_e2e.rs @@ -0,0 +1,188 @@ +use codex_chatgpt::apply_command::apply_diff_from_task; +use codex_chatgpt::get_task::GetTaskResponse; +use codex_utils_cargo_bin::find_resource; +use tempfile::TempDir; +use tokio::process::Command; + +/// Creates a temporary git repository with initial commit +async fn create_temp_git_repo() -> anyhow::Result { + let temp_dir = TempDir::new()?; + let repo_path = temp_dir.path(); + let envs = vec![ + ("GIT_CONFIG_GLOBAL", "/dev/null"), + ("GIT_CONFIG_NOSYSTEM", "1"), + ]; + + let output = Command::new("git") + .envs(envs.clone()) + .args(["init"]) + .current_dir(repo_path) + .output() + .await?; + + if !output.status.success() { + anyhow::bail!( + "Failed to initialize git repo: {}", + String::from_utf8_lossy(&output.stderr) + ); + } + + Command::new("git") + .envs(envs.clone()) + .args(["config", "user.email", "test@example.com"]) + .current_dir(repo_path) + .output() + .await?; + + Command::new("git") + .envs(envs.clone()) + .args(["config", "user.name", "Test User"]) + .current_dir(repo_path) + .output() + .await?; + + std::fs::write(repo_path.join("README.md"), "# Test Repo\n")?; + + Command::new("git") + .envs(envs.clone()) + .args(["add", "README.md"]) + .current_dir(repo_path) + .output() + .await?; + + let output = Command::new("git") + .envs(envs.clone()) + .args(["commit", "-m", "Initial commit"]) + .current_dir(repo_path) + .output() + .await?; + + if !output.status.success() { + anyhow::bail!( + "Failed to create initial commit: {}", + String::from_utf8_lossy(&output.stderr) + ); + } + + Ok(temp_dir) +} + +async fn mock_get_task_with_fixture() -> anyhow::Result { + let fixture_path = find_resource!("tests/task_turn_fixture.json")?; + let fixture_content = tokio::fs::read_to_string(fixture_path).await?; + let response: GetTaskResponse = serde_json::from_str(&fixture_content)?; + Ok(response) +} + +#[tokio::test] +async fn test_apply_command_creates_fibonacci_file() { + let temp_repo = create_temp_git_repo() + .await + .expect("Failed to create temp git repo"); + let repo_path = temp_repo.path(); + + let task_response = mock_get_task_with_fixture() + .await + .expect("Failed to load fixture"); + + apply_diff_from_task(task_response, Some(repo_path.to_path_buf())) + .await + .expect("Failed to apply diff from task"); + + // Assert that fibonacci.js was created in scripts/ directory + let fibonacci_path = repo_path.join("scripts/fibonacci.js"); + assert!(fibonacci_path.exists(), "fibonacci.js was not created"); + + // Verify the file contents match expected + let contents = std::fs::read_to_string(&fibonacci_path).expect("Failed to read fibonacci.js"); + assert!( + contents.contains("function fibonacci(n)"), + "fibonacci.js doesn't contain expected function" + ); + assert!( + contents.contains("#!/usr/bin/env node"), + "fibonacci.js doesn't have shebang" + ); + assert!( + contents.contains("module.exports = fibonacci;"), + "fibonacci.js doesn't export function" + ); + + // Verify file has correct number of lines (31 as specified in fixture) + let line_count = contents.lines().count(); + assert_eq!( + line_count, 31, + "fibonacci.js should have 31 lines, got {line_count}", + ); +} + +#[tokio::test] +async fn test_apply_command_with_merge_conflicts() { + let temp_repo = create_temp_git_repo() + .await + .expect("Failed to create temp git repo"); + let repo_path = temp_repo.path(); + + // Create conflicting fibonacci.js file first + let scripts_dir = repo_path.join("scripts"); + std::fs::create_dir_all(&scripts_dir).expect("Failed to create scripts directory"); + + let conflicting_content = r#"#!/usr/bin/env node + +// This is a different fibonacci implementation +function fib(num) { + if (num <= 1) return num; + return fib(num - 1) + fib(num - 2); +} + +console.log("Running fibonacci..."); +console.log(fib(10)); +"#; + + let fibonacci_path = scripts_dir.join("fibonacci.js"); + std::fs::write(&fibonacci_path, conflicting_content).expect("Failed to write conflicting file"); + + Command::new("git") + .args(["add", "scripts/fibonacci.js"]) + .current_dir(repo_path) + .output() + .await + .expect("Failed to add fibonacci.js"); + + Command::new("git") + .args(["commit", "-m", "Add conflicting fibonacci implementation"]) + .current_dir(repo_path) + .output() + .await + .expect("Failed to commit conflicting file"); + + let original_dir = std::env::current_dir().expect("Failed to get current dir"); + std::env::set_current_dir(repo_path).expect("Failed to change directory"); + struct DirGuard(std::path::PathBuf); + impl Drop for DirGuard { + fn drop(&mut self) { + let _ = std::env::set_current_dir(&self.0); + } + } + let _guard = DirGuard(original_dir); + + let task_response = mock_get_task_with_fixture() + .await + .expect("Failed to load fixture"); + + let apply_result = apply_diff_from_task(task_response, Some(repo_path.to_path_buf())).await; + + assert!( + apply_result.is_err(), + "Expected apply to fail due to merge conflicts" + ); + + let contents = std::fs::read_to_string(&fibonacci_path).expect("Failed to read fibonacci.js"); + + assert!( + contents.contains("<<<<<<< HEAD") + || contents.contains("=======") + || contents.contains(">>>>>>> "), + "fibonacci.js should contain merge conflict markers, got: {contents}", + ); +} diff --git a/codex-rs/chatgpt/tests/suite/mod.rs b/codex-rs/chatgpt/tests/suite/mod.rs new file mode 100644 index 0000000000000000000000000000000000000000..40b4a59a0d19c2b38e4cf3cab2bf02818a281a6b --- /dev/null +++ b/codex-rs/chatgpt/tests/suite/mod.rs @@ -0,0 +1,2 @@ +// Aggregates all former standalone integration tests as modules. +mod apply_command_e2e; diff --git a/codex-rs/chatgpt/tests/task_turn_fixture.json b/codex-rs/chatgpt/tests/task_turn_fixture.json new file mode 100644 index 0000000000000000000000000000000000000000..3750f550395aaea491f3d2287e1a18002820876e --- /dev/null +++ b/codex-rs/chatgpt/tests/task_turn_fixture.json @@ -0,0 +1,65 @@ +{ + "current_diff_task_turn": { + "output_items": [ + { + "type": "pr", + "pr_title": "Add fibonacci script", + "pr_message": "## Summary\n- add a basic Fibonacci script under `scripts/`\n\n## Testing\n- `node scripts/fibonacci.js 10`\n- `npm run lint` *(fails: next not found)*", + "output_diff": { + "type": "output_diff", + "repo_id": "/workspace/rddit-vercel", + "base_commit_sha": "1a2e9baf2ce2fdd0c126b47b1bcfd512de2a9f7b", + "diff": "diff --git a/scripts/fibonacci.js b/scripts/fibonacci.js\nnew file mode 100644\nindex 0000000000000000000000000000000000000000..6c9fdfdbf8669b7968936411050525b995d0a9a6\n--- /dev/null\n+++ b/scripts/fibonacci.js\n@@ -0,0 +1,31 @@\n+#!/usr/bin/env node\n+\n+function fibonacci(n) {\n+ if (n < 0) {\n+ throw new Error(\"n must be non-negative\");\n+ }\n+ let a = 0;\n+ let b = 1;\n+ for (let i = 0; i < n; i++) {\n+ const next = a + b;\n+ a = b;\n+ b = next;\n+ }\n+ return a;\n+}\n+\n+function printUsage() {\n+ console.log(\"Usage: node scripts/fibonacci.js \");\n+}\n+\n+if (require.main === module) {\n+ const arg = process.argv[2];\n+ if (arg === undefined || isNaN(Number(arg))) {\n+ printUsage();\n+ process.exit(1);\n+ }\n+ const n = Number(arg);\n+ console.log(fibonacci(n));\n+}\n+\n+module.exports = fibonacci;\n", + "external_storage_diff": { + "file_id": "file_00000000114c61f786900f8c2130ace7", + "ttl": null + }, + "files_modified": 1, + "lines_added": 31, + "lines_removed": 0, + "commit_message": "Add fibonacci script" + } + }, + { + "type": "message", + "role": "assistant", + "content": [ + { + "content_type": "text", + "text": "**Summary**\n\n- Created a command-line Fibonacci script that validates input and prints the result when executed with Node" + }, + { + "content_type": "repo_file_citation", + "path": "scripts/fibonacci.js", + "line_range_start": 1, + "line_range_end": 31 + }, + { + "content_type": "text", + "text": "\n\n**Testing**\n\n- ❌ `npm run lint` (failed to run `next lint`)" + }, + { + "content_type": "terminal_chunk_citation", + "terminal_chunk_id": "7dd543", + "line_range_start": 1, + "line_range_end": 5 + }, + { + "content_type": "text", + "text": "\n- ✅ `node scripts/fibonacci.js 10` produced “55”" + }, + { + "content_type": "terminal_chunk_citation", + "terminal_chunk_id": "6ee559", + "line_range_start": 1, + "line_range_end": 3 + }, + { + "content_type": "text", + "text": "\n\nCodex couldn't run certain commands due to environment limitations. Consider configuring a setup script or internet access in your Codex environment to install dependencies." + } + ] + } + ] + } +} diff --git a/codex-rs/codex-client/src/lib.rs b/codex-rs/codex-client/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..60b8d0006394e456754d06c214ed90a2c55ee5e9 --- /dev/null +++ b/codex-rs/codex-client/src/lib.rs @@ -0,0 +1,14 @@ +mod retry; +mod sse; +mod telemetry; + +pub use crate::retry::RetryOn; +pub use crate::retry::RetryOperation; +pub use crate::retry::RetryPolicy; +pub use crate::retry::backoff; +pub use crate::retry::run_with_retry; +pub use crate::sse::sse_stream; +pub use crate::telemetry::RequestTelemetry; +pub use codex_http_client::HttpClient as CodexHttpClient; +pub use codex_http_client::RequestBuilder as CodexRequestBuilder; +pub use codex_http_client::*; diff --git a/codex-rs/codex-client/src/retry.rs b/codex-rs/codex-client/src/retry.rs new file mode 100644 index 0000000000000000000000000000000000000000..86071851f61094a595af5024bbbb047f29275255 --- /dev/null +++ b/codex-rs/codex-client/src/retry.rs @@ -0,0 +1,109 @@ +use codex_http_client::Request; +use codex_http_client::TransportError; +use rand::Rng; +use std::future::Future; +use std::time::Duration; + +#[derive(Debug, Clone)] +pub struct RetryPolicy { + pub max_attempts: u64, + pub base_delay: Duration, + pub retry_on: RetryOn, +} + +#[derive(Debug, Clone)] +pub struct RetryOn { + pub retry_429: bool, + pub retry_5xx: bool, + pub retry_transport: bool, +} + +impl RetryOn { + pub fn should_retry(&self, err: &TransportError, attempt: u64, max_attempts: u64) -> bool { + if attempt >= max_attempts { + return false; + } + match err { + TransportError::Http { status, .. } => { + (self.retry_429 && status.as_u16() == 429) + || (self.retry_5xx && status.is_server_error()) + } + TransportError::Timeout + | TransportError::Connection(_) + | TransportError::Network(_) => self.retry_transport, + TransportError::Build(_) + | TransportError::RetryLimit + | TransportError::ResponseTooLarge { .. } => false, + } + } +} + +pub fn backoff(base: Duration, attempt: u64) -> Duration { + if attempt == 0 { + return base; + } + let exp = 2u64.saturating_pow(attempt as u32 - 1); + let millis = base.as_millis() as u64; + let raw = millis.saturating_mul(exp); + let jitter: f64 = rand::rng().random_range(0.9..1.1); + Duration::from_millis((raw as f64 * jitter) as u64) +} + +/// Identifies a retry path and its associated trace-event layer. +#[derive(Debug, Clone, Copy)] +pub enum RetryOperation { + HttpRequest, + Sampling, + RemoteCompactionV2, +} + +/// Emits retry telemetry at the caller's source location without adding it to normal OTEL logs. +#[macro_export] +macro_rules! record_retry { + ($attempt:expr, $delay:expr, $operation:expr $(,)?) => {{ + let (layer, operation) = match $operation { + $crate::RetryOperation::HttpRequest => ("http", "request"), + $crate::RetryOperation::Sampling => ("stream", "sampling"), + $crate::RetryOperation::RemoteCompactionV2 => ("stream", "remote_compaction_v2"), + }; + + ::tracing::event!( + target: "codex_otel.trace_safe", + ::tracing::Level::TRACE, + event.name = "codex.retry", + retry.attempt = $attempt, + retry.delay_ms = ($delay).as_millis() as u64, + retry.layer = layer, + retry.operation = operation, + ); + }}; +} + +pub async fn run_with_retry( + policy: RetryPolicy, + mut make_req: impl FnMut() -> Request, + op: F, +) -> Result +where + F: Fn(Request, u64) -> Fut, + Fut: Future>, +{ + for attempt in 0..=policy.max_attempts { + let req = make_req(); + match op(req, attempt).await { + Ok(resp) => return Ok(resp), + Err(err) + if policy + .retry_on + .should_retry(&err, attempt, policy.max_attempts) => + { + let retry_attempt = attempt + 1; + let delay = backoff(policy.base_delay, retry_attempt); + crate::record_retry!(retry_attempt, delay, RetryOperation::HttpRequest); + tokio::time::sleep(delay).await; + } + Err(err) => return Err(err), + } + } + Err(TransportError::RetryLimit) +} diff --git a/codex-rs/codex-client/src/sse.rs b/codex-rs/codex-client/src/sse.rs new file mode 100644 index 0000000000000000000000000000000000000000..ed1591f6de4bf4c7d7cef3ecbba8e52d07d2571c --- /dev/null +++ b/codex-rs/codex-client/src/sse.rs @@ -0,0 +1,48 @@ +use codex_http_client::ByteStream; +use codex_http_client::StreamError; +use eventsource_stream::Eventsource; +use futures::StreamExt; +use tokio::sync::mpsc; +use tokio::time::Duration; +use tokio::time::timeout; + +/// Minimal SSE helper that forwards raw `data:` frames as UTF-8 strings. +/// +/// Errors and idle timeouts are sent as `Err(StreamError)` before the task exits. +pub fn sse_stream( + stream: ByteStream, + idle_timeout: Duration, + tx: mpsc::Sender>, +) { + tokio::spawn(async move { + let mut stream = stream + .map(|res| res.map_err(|e| StreamError::Stream(e.to_string()))) + .eventsource(); + + loop { + match timeout(idle_timeout, stream.next()).await { + Ok(Some(Ok(ev))) => { + if tx.send(Ok(ev.data.clone())).await.is_err() { + return; + } + } + Ok(Some(Err(e))) => { + let _ = tx.send(Err(StreamError::Stream(e.to_string()))).await; + return; + } + Ok(None) => { + let _ = tx + .send(Err(StreamError::Stream( + "stream closed before completion".into(), + ))) + .await; + return; + } + Err(_) => { + let _ = tx.send(Err(StreamError::Timeout)).await; + return; + } + } + } + }); +} diff --git a/codex-rs/codex-client/src/telemetry.rs b/codex-rs/codex-client/src/telemetry.rs new file mode 100644 index 0000000000000000000000000000000000000000..b856414d3332b25f88982dce9c97d3f5be3caa04 --- /dev/null +++ b/codex-rs/codex-client/src/telemetry.rs @@ -0,0 +1,14 @@ +use codex_http_client::TransportError; +use http::StatusCode; +use std::time::Duration; + +/// API specific telemetry. +pub trait RequestTelemetry: Send + Sync { + fn on_request( + &self, + attempt: u64, + status: Option, + error: Option<&TransportError>, + duration: Duration, + ); +} diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..a8facc3ae11fe761edd260ed51f187ded6491482 --- /dev/null +++ b/codex-rs/core-api/src/lib.rs @@ -0,0 +1,137 @@ +//! Public facade for thread management APIs built on `codex-core`. + +#![deny(private_bounds, private_interfaces, unreachable_pub)] + +pub use codex_analytics::AnalyticsEventsClient; +pub use codex_app_server_protocol::ServerNotification; +pub use codex_app_server_protocol::item_event_to_server_notification; +pub use codex_arg0::Arg0DispatchPaths; +pub use codex_arg0::arg0_dispatch_or_else; +pub use codex_config::ConfigLayerStack; +pub use codex_config::config_toml::ProjectConfig; +pub use codex_config::config_toml::RealtimeAudioConfig; +pub use codex_config::config_toml::RealtimeConfig; +pub use codex_config::types::AuthCredentialsStoreMode; +pub use codex_config::types::AuthKeyringBackendKind; +pub use codex_config::types::History; +pub use codex_config::types::MemoriesConfig; +pub use codex_config::types::ModelAvailabilityNuxConfig; +pub use codex_config::types::Notice; +pub use codex_config::types::OAuthCredentialsStoreMode; +pub use codex_config::types::OtelConfig; +pub use codex_config::types::SessionPickerViewMode; +pub use codex_config::types::ToolSuggestConfig; +pub use codex_config::types::TuiKeymap; +pub use codex_config::types::TuiNotificationSettings; +pub use codex_config::types::TuiPetAnchor; +pub use codex_config::types::UriBasedFileOpener; +pub use codex_core::CodexAppsToolsCache; +pub use codex_core::CodexThread; +pub use codex_core::CodexThreadSettingsOverrides; +pub use codex_core::EnvironmentConfig; +pub use codex_core::EnvironmentMcpPolicy; +pub use codex_core::EnvironmentNetworkPolicy; +pub use codex_core::ForkSnapshot; +pub use codex_core::LoadedAgentsMd; +pub use codex_core::McpManager; +pub use codex_core::NetworkDomainPermission; +pub use codex_core::NetworkDomainPermissionEntry; +pub use codex_core::NetworkDomainPermissions; +pub use codex_core::NetworkUnixSocketPermission; +pub use codex_core::NetworkUnixSocketPermissions; +pub use codex_core::NewThread; +pub use codex_core::NotSubmittedReason; +pub use codex_core::RecoverTurnRequest; +pub use codex_core::StartIfIdleSubmission; +pub use codex_core::StartThreadOptions; +pub use codex_core::StateDbHandle; +pub use codex_core::SteerSubmission; +pub use codex_core::SuspendTurnOutcome; +pub use codex_core::ThreadManager; +pub use codex_core::ThreadShutdownReport; +pub use codex_core::TurnInput; +pub use codex_core::TurnInputRequest; +pub use codex_core::TurnInputSubmission; +pub use codex_core::TurnStartOptions; +pub use codex_core::WaitForEnvironmentToolConfig; +pub use codex_core::build_models_manager; +pub use codex_core::config::Config; +pub use codex_core::config::Constrained; +pub use codex_core::config::ExtraConfig; +pub use codex_core::config::GhostSnapshotConfig; +pub use codex_core::config::MultiAgentV2Config; +pub use codex_core::config::NetworkProxySpec; +pub use codex_core::config::Permissions; +pub use codex_core::config::TerminalResizeReflowConfig; +pub use codex_core::config::ThreadStoreConfig; +pub use codex_core::config::find_codex_home; +pub use codex_core::init_state_db; +pub use codex_core::local_agent_graph_store_from_state_db; +pub use codex_core::passthrough_image_store; +pub use codex_core::resolve_installation_id; +pub use codex_core::thread_store_from_config; +pub use codex_exec_server::EnvironmentManager; +pub use codex_exec_server::EnvironmentRegistryConnectRequest; +pub use codex_exec_server::EnvironmentRegistryConnectResponse; +pub use codex_exec_server::EnvironmentRegistryHarnessKeyValidationRequest; +pub use codex_exec_server::EnvironmentRegistryHarnessKeyValidationResponse; +pub use codex_exec_server::EnvironmentRegistryRegistrationRequest; +pub use codex_exec_server::EnvironmentRegistryRegistrationResponse; +pub use codex_exec_server::ExecServerError; +pub use codex_exec_server::ExecServerRuntimePaths; +pub use codex_exec_server::NoiseChannelIdentity; +pub use codex_exec_server::NoiseChannelPublicKey; +pub use codex_exec_server::NoiseRendezvousConnectBundle; +pub use codex_exec_server::NoiseRendezvousConnectProvider; +pub use codex_extension_api::ExtensionRegistryBuilder; +pub use codex_extension_api::Instructions; +pub use codex_extension_api::LoadInstructionsFuture; +pub use codex_extension_api::LoadedUserInstructions; +pub use codex_extension_api::ThreadInstructionsProvider; +pub use codex_extension_api::UserInstructionsProvider; +pub use codex_extension_api::empty_extension_registry; +pub use codex_features::Feature; +pub use codex_features::Features; +pub use codex_history::InitialHistory; +pub use codex_home::CodexHomeUserInstructionsProvider; +pub use codex_image_generation_extension::install as install_image_generation_extension; +pub use codex_login::AuthHeaders; +pub use codex_login::AuthManager; +pub use codex_login::CodexAuth; +pub use codex_login::ExternalAuth; +pub use codex_login::ExternalAuthFuture; +pub use codex_login::ExternalAuthRefreshContext; +pub use codex_login::ExternalAuthRefreshReason; +pub use codex_login::default_client::set_default_originator; +pub use codex_model_provider_info::OPENAI_PROVIDER_ID; +pub use codex_model_provider_info::built_in_model_providers; +pub use codex_models_manager::manager::RefreshStrategy; +pub use codex_models_manager::manager::SharedModelsManager; +pub use codex_protocol::ThreadId; +pub use codex_protocol::capabilities::CapabilityRootLocation; +pub use codex_protocol::capabilities::SelectedCapabilityRoot; +pub use codex_protocol::config_types::AltScreenMode; +pub use codex_protocol::config_types::ApprovalsReviewer; +pub use codex_protocol::config_types::AutoCompactTokenLimitScope; +pub use codex_protocol::config_types::CollaborationModeMask; +pub use codex_protocol::config_types::ShellEnvironmentPolicy; +pub use codex_protocol::config_types::WebSearchMode; +pub use codex_protocol::dynamic_tools::DynamicToolFunctionSpec; +pub use codex_protocol::dynamic_tools::DynamicToolNamespaceSpec; +pub use codex_protocol::dynamic_tools::DynamicToolNamespaceTool; +pub use codex_protocol::dynamic_tools::DynamicToolSpec; +pub use codex_protocol::error::Result as CodexResult; +pub use codex_protocol::models::PermissionProfile; +pub use codex_protocol::models::PermissionProfileSnapshot; +pub use codex_protocol::openai_models::ModelPreset; +pub use codex_protocol::protocol::AskForApproval; +pub use codex_protocol::protocol::EnvironmentConfigState; +pub use codex_protocol::protocol::EventMsg; +pub use codex_protocol::protocol::Op; +pub use codex_protocol::protocol::SessionConfiguredEvent; +pub use codex_protocol::protocol::SessionSource; +pub use codex_protocol::protocol::TurnEnvironmentSelection; +pub use codex_protocol::protocol::W3cTraceContext; +pub use codex_protocol::user_input::UserInput; +pub use codex_state::SqliteConfig; +pub use codex_utils_absolute_path::AbsolutePathBuf; diff --git a/codex-rs/diagnostics/src/lib.rs b/codex-rs/diagnostics/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..a966e4aee546e586d141ee5b7f70bc82d377911f --- /dev/null +++ b/codex-rs/diagnostics/src/lib.rs @@ -0,0 +1,210 @@ +use std::sync::Mutex; +use std::sync::Once; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +static GAUGES: Mutex> = Mutex::new(Vec::new()); + +/// A process-wide gauge that registers itself the first time it is used. +pub struct Gauge { + name: &'static str, + value: AtomicU64, + registered: Once, +} + +impl Gauge { + /// Creates a gauge suitable for use in a `static` declaration. + pub const fn new(name: &'static str) -> Self { + Self { + name, + value: AtomicU64::new(0), + registered: Once::new(), + } + } + + /// Increments the gauge and registers it if needed. + pub fn increment(&'static self) { + self.register(); + self.value.fetch_add(1, Ordering::Relaxed); + } + + /// Decrements the gauge without allowing an underflow. + pub fn decrement(&self) { + let _ = self + .value + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |value| { + Some(value.saturating_sub(1)) + }); + } + + /// Increments the gauge for the lifetime of the returned guard. + pub fn track(&'static self) -> GaugeGuard { + self.increment(); + GaugeGuard { gauge: self } + } + + fn register(&'static self) { + self.registered.call_once(|| { + GAUGES + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(self); + }); + } +} + +/// Decrements its gauge when the measured object is dropped. +pub struct GaugeGuard { + gauge: &'static Gauge, +} + +impl Drop for GaugeGuard { + fn drop(&mut self) { + self.gauge.decrement(); + } +} + +/// The current value of one registered diagnostic gauge. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct GaugeSnapshot { + pub name: &'static str, + pub value: u64, +} + +/// Best-effort operating-system measurements for the current process. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct ProcessSnapshot { + pub id: u32, + pub resident_memory_bytes: Option, + pub physical_footprint_bytes: Option, +} + +/// Content-free diagnostic values contributed by this process. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct DiagnosticsSnapshot { + pub process: ProcessSnapshot, + pub gauges: Vec, +} + +/// Collects built-in process measurements and every registered gauge. +pub fn snapshot() -> DiagnosticsSnapshot { + let mut gauges = GAUGES + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .iter() + .map(|gauge| GaugeSnapshot { + name: gauge.name, + value: gauge.value.load(Ordering::Relaxed), + }) + .collect::>(); + gauges.sort_unstable_by_key(|gauge| gauge.name); + + DiagnosticsSnapshot { + process: process_snapshot(), + gauges, + } +} + +#[cfg(target_os = "macos")] +fn process_snapshot() -> ProcessSnapshot { + let usage = unsafe { + let mut usage = std::mem::MaybeUninit::::zeroed(); + // SAFETY: the kernel initializes this correctly sized buffer on success. + if libc::proc_pid_rusage( + libc::getpid(), + libc::RUSAGE_INFO_V0, + usage.as_mut_ptr().cast(), + ) != 0 + { + return empty_process_snapshot(); + } + usage.assume_init() + }; + + ProcessSnapshot { + id: std::process::id(), + resident_memory_bytes: Some(usage.ri_resident_size), + physical_footprint_bytes: Some(usage.ri_phys_footprint), + } +} + +#[cfg(target_os = "linux")] +fn process_snapshot() -> ProcessSnapshot { + // SAFETY: querying the system page size does not access caller-owned memory. + let page_size = u64::try_from(unsafe { libc::sysconf(libc::_SC_PAGESIZE) }) + .ok() + .filter(|page_size| *page_size > 0); + let resident_pages = std::fs::read_to_string("/proc/self/statm") + .ok() + .and_then(|statm| statm.split_whitespace().nth(1)?.parse::().ok()); + + ProcessSnapshot { + id: std::process::id(), + resident_memory_bytes: resident_pages + .zip(page_size) + .map(|(pages, page_size)| pages.saturating_mul(page_size)), + physical_footprint_bytes: None, + } +} + +#[cfg(target_os = "windows")] +fn process_snapshot() -> ProcessSnapshot { + #[repr(C)] + struct ProcessMemoryCounters { + size: u32, + page_fault_count: u32, + peak_working_set_size: usize, + working_set_size: usize, + quota_peak_paged_pool_usage: usize, + quota_paged_pool_usage: usize, + quota_peak_non_paged_pool_usage: usize, + quota_non_paged_pool_usage: usize, + pagefile_usage: usize, + peak_pagefile_usage: usize, + } + + #[link(name = "kernel32")] + unsafe extern "system" { + fn GetCurrentProcess() -> *mut std::ffi::c_void; + fn K32GetProcessMemoryInfo( + process: *mut std::ffi::c_void, + counters: *mut ProcessMemoryCounters, + size: u32, + ) -> i32; + } + + let counters = unsafe { + let mut counters = std::mem::MaybeUninit::::zeroed(); + let size = u32::try_from(std::mem::size_of::()).unwrap_or(u32::MAX); + // SAFETY: the pseudo-handle is valid and the kernel initializes this + // correctly sized writable buffer on success. + if K32GetProcessMemoryInfo(GetCurrentProcess(), counters.as_mut_ptr(), size) == 0 { + return empty_process_snapshot(); + } + counters.assume_init() + }; + + ProcessSnapshot { + id: std::process::id(), + resident_memory_bytes: u64::try_from(counters.working_set_size).ok(), + physical_footprint_bytes: None, + } +} + +#[cfg(not(any(target_os = "macos", target_os = "linux", target_os = "windows")))] +fn process_snapshot() -> ProcessSnapshot { + empty_process_snapshot() +} + +#[cfg(not(target_os = "linux"))] +fn empty_process_snapshot() -> ProcessSnapshot { + ProcessSnapshot { + id: std::process::id(), + resident_memory_bytes: None, + physical_footprint_bytes: None, + } +} + +#[cfg(test)] +#[path = "tests.rs"] +mod tests; diff --git a/codex-rs/diagnostics/src/tests.rs b/codex-rs/diagnostics/src/tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..d5ac2f2a5acbabe1b5e65e8bf6e29f57e5e10c50 --- /dev/null +++ b/codex-rs/diagnostics/src/tests.rs @@ -0,0 +1,58 @@ +use super::Gauge; +use super::snapshot; +use pretty_assertions::assert_eq; + +static GUARD_GAUGE: Gauge = Gauge::new("diagnostics.tests.guard"); +static REGISTERED_GAUGE: Gauge = Gauge::new("diagnostics.tests.registered"); + +#[test] +fn gauge_guards_follow_the_measured_lifetime() { + let first = GUARD_GAUGE.track(); + let second = GUARD_GAUGE.track(); + let value = || { + snapshot() + .gauges + .into_iter() + .find(|gauge| gauge.name == "diagnostics.tests.guard") + .expect("used gauge should be registered") + .value + }; + + assert_eq!(value(), 2); + drop(first); + assert_eq!(value(), 1); + drop(second); + assert_eq!(value(), 0); +} + +#[test] +fn snapshot_includes_process_memory_and_registers_gauges_once() { + REGISTERED_GAUGE.increment(); + REGISTERED_GAUGE.increment(); + let diagnostics = snapshot(); + let registered = diagnostics + .gauges + .iter() + .filter(|gauge| gauge.name == "diagnostics.tests.registered") + .collect::>(); + + assert_eq!(diagnostics.process.id, std::process::id()); + assert_eq!(registered.len(), 1); + assert_eq!(registered[0].value, 2); + #[cfg(any(target_os = "macos", target_os = "linux", target_os = "windows"))] + assert!( + diagnostics + .process + .resident_memory_bytes + .is_some_and(|bytes| bytes > 0) + ); + #[cfg(target_os = "macos")] + assert!( + diagnostics + .process + .physical_footprint_bytes + .is_some_and(|bytes| bytes > 0) + ); + #[cfg(not(target_os = "macos"))] + assert_eq!(diagnostics.process.physical_footprint_bytes, None); +} diff --git a/codex-rs/exec/src/cli.rs b/codex-rs/exec/src/cli.rs new file mode 100644 index 0000000000000000000000000000000000000000..0ef0be6af9cc177e86da68e9fa20df05f9ac813c --- /dev/null +++ b/codex-rs/exec/src/cli.rs @@ -0,0 +1,319 @@ +use clap::Args; +use clap::FromArgMatches; +use clap::Parser; +use clap::ValueEnum; +use codex_protocol::protocol::ThreadSource; +use codex_utils_cli::CliConfigOverrides; +use codex_utils_cli::SharedCliOptions; +use std::path::PathBuf; + +#[derive(Parser, Debug)] +#[command( + version, + override_usage = "codex exec [OPTIONS] [PROMPT]\n codex exec [OPTIONS] [ARGS]" +)] +pub struct Cli { + /// Action to perform. If omitted, runs a new non-interactive session. + #[command(subcommand)] + pub command: Option, + + /// Error out when config.toml contains fields that are not recognized by this version of Codex. + #[arg(long = "strict-config", global = true, default_value_t = false)] + pub strict_config: bool, + + #[clap(flatten)] + pub shared: ExecSharedCliOptions, + + /// Source classification for newly created or forked threads. + #[arg(long = "thread-source", value_name = "SOURCE", global = true)] + pub thread_source: Option, + + /// Allow running Codex outside a Git repository. + #[arg(long = "skip-git-repo-check", global = true, default_value_t = false)] + pub skip_git_repo_check: bool, + + /// Run without persisting session files to disk. + #[arg(long = "ephemeral", global = true, default_value_t = false)] + pub ephemeral: bool, + + /// Do not load `$CODEX_HOME/config.toml`; auth still uses `CODEX_HOME`. + #[arg(long = "ignore-user-config", global = true, default_value_t = false)] + pub ignore_user_config: bool, + + /// Do not load user or project execpolicy `.rules` files. + #[arg(long = "ignore-rules", global = true, default_value_t = false)] + pub ignore_rules: bool, + + /// Path to a JSON Schema file describing the model's final response shape. + #[arg(long = "output-schema", value_name = "FILE", global = true)] + pub output_schema: Option, + + #[clap(skip)] + pub config_overrides: CliConfigOverrides, + + /// Specifies color settings for use in the output. + #[arg(long = "color", value_enum, default_value_t = Color::Auto)] + pub color: Color, + + /// Print events to stdout as JSONL. + #[arg( + long = "json", + alias = "experimental-json", + default_value_t = false, + global = true + )] + pub json: bool, + + /// Specifies file where the last message from the agent should be written. + #[arg( + long = "output-last-message", + short = 'o', + value_name = "FILE", + global = true + )] + pub last_message_file: Option, + + /// Initial instructions for the agent. If not provided as an argument (or + /// if `-` is used), instructions are read from stdin. If stdin is piped and + /// a prompt is also provided, stdin is appended as a `` block. + #[arg(value_name = "PROMPT", value_hint = clap::ValueHint::Other)] + pub prompt: Option, +} + +impl std::ops::Deref for Cli { + type Target = SharedCliOptions; + + fn deref(&self) -> &Self::Target { + &self.shared.0 + } +} + +impl std::ops::DerefMut for Cli { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.shared.0 + } +} + +#[derive(Debug, Default)] +pub struct ExecSharedCliOptions(SharedCliOptions); + +impl ExecSharedCliOptions { + pub fn into_inner(self) -> SharedCliOptions { + self.0 + } +} + +impl std::ops::Deref for ExecSharedCliOptions { + type Target = SharedCliOptions; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl std::ops::DerefMut for ExecSharedCliOptions { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} + +impl Args for ExecSharedCliOptions { + fn augment_args(cmd: clap::Command) -> clap::Command { + mark_exec_global_args(SharedCliOptions::augment_args(cmd)) + } + + fn augment_args_for_update(cmd: clap::Command) -> clap::Command { + mark_exec_global_args(SharedCliOptions::augment_args_for_update(cmd)) + } +} + +impl FromArgMatches for ExecSharedCliOptions { + fn from_arg_matches(matches: &clap::ArgMatches) -> Result { + SharedCliOptions::from_arg_matches(matches).map(Self) + } + + fn update_from_arg_matches(&mut self, matches: &clap::ArgMatches) -> Result<(), clap::Error> { + self.0.update_from_arg_matches(matches) + } +} + +fn mark_exec_global_args(cmd: clap::Command) -> clap::Command { + cmd.mut_arg("model", |arg| arg.global(true)) + .mut_arg("dangerously_bypass_approvals_and_sandbox", |arg| { + arg.global(true) + }) + .mut_arg("bypass_hook_trust", |arg| arg.global(true)) + .mut_arg("worktree", |arg| arg.global(true)) +} + +#[derive(Debug, clap::Subcommand)] +pub enum Command { + /// Resume a previous session by id or pick the most recent with --last. + Resume(ResumeArgs), + + /// Fork a previous session by id into a new session. + Fork(ForkArgs), + + /// Run a code review against the current repository. + Review(ReviewArgs), +} + +#[derive(Args, Debug)] +pub struct ForkArgs { + /// Conversation/session id (UUID) or thread name to fork. + #[arg(value_name = "SESSION_ID")] + pub session_id: String, + + /// Optional image(s) to attach to the prompt sent after forking. + #[arg( + long = "image", + short = 'i', + value_name = "FILE", + value_delimiter = ',', + num_args = 1 + )] + pub images: Vec, + + /// Optional prompt to send after forking. If `-` is used, read from stdin. + #[arg(value_name = "PROMPT", value_hint = clap::ValueHint::Other)] + pub prompt: Option, +} + +#[derive(Args, Debug)] +struct ResumeArgsRaw { + // Note: This is the direct clap shape. We reinterpret the positional when --last is set + // so "codex resume --last " treats the positional as a prompt, not a session id. + /// Conversation/session id (UUID) or thread name. UUIDs take precedence if it parses. + /// If omitted, use --last to pick the most recent recorded session. + #[arg(value_name = "SESSION_ID")] + session_id: Option, + + /// Resume the most recent recorded session (newest) without specifying an id. + #[arg(long = "last", default_value_t = false)] + last: bool, + + /// Show all sessions (disables cwd filtering). + #[arg(long = "all", default_value_t = false)] + all: bool, + + /// Optional image(s) to attach to the prompt sent after resuming. + #[arg( + long = "image", + short = 'i', + value_name = "FILE", + value_delimiter = ',', + num_args = 1 + )] + images: Vec, + + /// Prompt to send after resuming the session. If `-` is used, read from stdin. + #[arg(value_name = "PROMPT", value_hint = clap::ValueHint::Other)] + prompt: Option, +} + +#[derive(Debug)] +pub struct ResumeArgs { + /// Conversation/session id (UUID) or thread name. UUIDs take precedence if it parses. + /// If omitted, use --last to pick the most recent recorded session. + pub session_id: Option, + + /// Resume the most recent recorded session (newest) without specifying an id. + pub last: bool, + + /// Show all sessions (disables cwd filtering). + pub all: bool, + + /// Optional image(s) to attach to the prompt sent after resuming. + pub images: Vec, + + /// Prompt to send after resuming the session. If `-` is used, read from stdin. + pub prompt: Option, +} + +impl From for ResumeArgs { + fn from(raw: ResumeArgsRaw) -> Self { + // When --last is used without an explicit prompt, treat the positional as the prompt + // (clap can’t express this conditional positional meaning cleanly). + let (session_id, prompt) = if raw.last && raw.prompt.is_none() { + (None, raw.session_id) + } else { + (raw.session_id, raw.prompt) + }; + Self { + session_id, + last: raw.last, + all: raw.all, + images: raw.images, + prompt, + } + } +} + +impl Args for ResumeArgs { + fn augment_args(cmd: clap::Command) -> clap::Command { + ResumeArgsRaw::augment_args(cmd) + } + + fn augment_args_for_update(cmd: clap::Command) -> clap::Command { + ResumeArgsRaw::augment_args_for_update(cmd) + } +} + +impl FromArgMatches for ResumeArgs { + fn from_arg_matches(matches: &clap::ArgMatches) -> Result { + ResumeArgsRaw::from_arg_matches(matches).map(Self::from) + } + + fn update_from_arg_matches(&mut self, matches: &clap::ArgMatches) -> Result<(), clap::Error> { + *self = ResumeArgsRaw::from_arg_matches(matches).map(Self::from)?; + Ok(()) + } +} + +#[derive(Args, Debug)] +pub struct ReviewArgs { + /// Review staged, unstaged, and untracked changes. + #[arg( + long = "uncommitted", + default_value_t = false, + conflicts_with_all = ["base", "commit", "prompt"] + )] + pub uncommitted: bool, + + /// Review changes against the given base branch. + #[arg( + long = "base", + value_name = "BRANCH", + conflicts_with_all = ["uncommitted", "commit", "prompt"] + )] + pub base: Option, + + /// Review the changes introduced by a commit. + #[arg( + long = "commit", + value_name = "SHA", + conflicts_with_all = ["uncommitted", "base", "prompt"] + )] + pub commit: Option, + + /// Optional commit title to display in the review summary. + #[arg(long = "title", value_name = "TITLE", requires = "commit")] + pub commit_title: Option, + + /// Custom review instructions. If `-` is used, read from stdin. + #[arg(value_name = "PROMPT", value_hint = clap::ValueHint::Other)] + pub prompt: Option, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, ValueEnum)] +#[value(rename_all = "kebab-case")] +pub enum Color { + Always, + Never, + #[default] + Auto, +} + +#[cfg(test)] +#[path = "cli_tests.rs"] +mod tests; diff --git a/codex-rs/exec/src/cli_tests.rs b/codex-rs/exec/src/cli_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..423935904acf25d95a348ef26334d161089b492f --- /dev/null +++ b/codex-rs/exec/src/cli_tests.rs @@ -0,0 +1,147 @@ +use super::*; +use pretty_assertions::assert_eq; + +#[test] +fn resume_parses_prompt_after_global_flags() { + const PROMPT: &str = "echo resume-with-global-flags-after-subcommand"; + let cli = Cli::parse_from([ + "codex-exec", + "resume", + "--last", + "--json", + "--model", + "gpt-5.2-codex", + "--dangerously-bypass-approvals-and-sandbox", + "--skip-git-repo-check", + "--ephemeral", + "--ignore-user-config", + "--ignore-rules", + PROMPT, + ]); + + assert!(cli.ephemeral); + assert!(cli.ignore_user_config); + assert!(cli.ignore_rules); + let Some(Command::Resume(args)) = cli.command else { + panic!("expected resume command"); + }; + let effective_prompt = args.prompt.clone().or_else(|| { + if args.last { + args.session_id.clone() + } else { + None + } + }); + assert_eq!(effective_prompt.as_deref(), Some(PROMPT)); +} + +#[test] +fn resume_accepts_output_flags_after_subcommand() { + const PROMPT: &str = "echo resume-with-output-file"; + let cli = Cli::parse_from([ + "codex-exec", + "resume", + "session-123", + "-o", + "/tmp/resume-output.md", + "--output-schema", + "/tmp/schema.json", + PROMPT, + ]); + + assert_eq!( + cli.last_message_file, + Some(PathBuf::from("/tmp/resume-output.md")) + ); + assert_eq!(cli.output_schema, Some(PathBuf::from("/tmp/schema.json"))); + let Some(Command::Resume(args)) = cli.command else { + panic!("expected resume command"); + }; + assert_eq!(args.session_id.as_deref(), Some("session-123")); + assert_eq!(args.prompt.as_deref(), Some(PROMPT)); +} + +#[test] +fn fork_parses_prompt_after_global_flags() { + const PROMPT: &str = "continue on the fork"; + let cli = Cli::parse_from([ + "codex-exec", + "fork", + "session-123", + "--json", + "--model", + "gpt-5.2-codex", + "--thread-source", + "automated_review", + "--skip-git-repo-check", + "--ephemeral", + PROMPT, + ]); + + assert!(cli.json); + assert!(cli.ephemeral); + assert_eq!( + cli.thread_source, + Some(ThreadSource::Feature("automated_review".to_string())) + ); + let Some(Command::Fork(args)) = cli.command else { + panic!("expected fork command"); + }; + assert_eq!(args.session_id, "session-123"); + assert_eq!(args.prompt.as_deref(), Some(PROMPT)); +} + +#[test] +fn parses_config_isolation_flags() { + let cli = Cli::parse_from([ + "codex-exec", + "--ignore-user-config", + "--ignore-rules", + "summarize", + ]); + + assert!(cli.ignore_user_config); + assert!(cli.ignore_rules); +} + +#[test] +fn approve_for_me_flag_applies_to_resume_when_passed_at_exec_root() { + for flag in ["--approve-for-me", "--not-so-yolo"] { + let cli = Cli::parse_from(["codex-exec", flag, "resume", "--last"]); + + assert!(cli.auto_review); + } +} + +#[test] +fn approve_for_me_flag_conflicts_with_other_sandbox_modes() { + for conflicting_args in [ + vec!["--sandbox", "read-only"], + vec!["--dangerously-bypass-approvals-and-sandbox"], + ] { + let mut args = vec!["codex-exec", "--approve-for-me"]; + args.extend(conflicting_args); + args.push("summarize"); + + let error = Cli::try_parse_from(args).expect_err("flags should conflict"); + assert_eq!(error.kind(), clap::error::ErrorKind::ArgumentConflict); + } +} + +#[test] +fn worktree_flag_is_accepted_after_fork_subcommand() { + let cli = Cli::try_parse_from(["codex-exec", "fork", "session-id", "--worktree"]) + .expect("worktree should be a global exec argument"); + + assert!(cli.worktree); + assert!(matches!(cli.command, Some(Command::Fork(_)))); +} + +#[test] +fn worktree_flag_is_accepted_before_fork_subcommand() { + let cli = Cli::try_parse_from(["codex-exec", "--worktree", "fork", "session-id"]) + .expect("worktree should be accepted before the fork subcommand"); + + assert!(cli.worktree); + assert!(matches!(cli.command, Some(Command::Fork(_)))); +} diff --git a/codex-rs/exec/src/event_processor.rs b/codex-rs/exec/src/event_processor.rs new file mode 100644 index 0000000000000000000000000000000000000000..4d1da09001261651a550e36cf2ceaf1d6235c83c --- /dev/null +++ b/codex-rs/exec/src/event_processor.rs @@ -0,0 +1,48 @@ +use std::path::Path; + +use codex_app_server_protocol::ServerNotification; +use codex_core::config::Config; +use codex_protocol::protocol::SessionConfiguredEvent; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CodexStatus { + Running, + InitiateShutdown, +} + +pub(crate) trait EventProcessor { + /// Print summary of effective configuration and user prompt. + fn print_config_summary( + &mut self, + config: &Config, + prompt: &str, + session_configured: &SessionConfiguredEvent, + ); + + /// Handle a single typed app-server notification emitted by the agent. + fn process_server_notification(&mut self, notification: ServerNotification) -> CodexStatus; + + /// Handle a local exec warning that is not represented as an app-server notification. + fn process_warning(&mut self, message: String) -> CodexStatus; + + fn print_final_output(&mut self) {} +} + +pub(crate) fn handle_last_message(last_agent_message: Option<&str>, output_file: &Path) { + let message = last_agent_message.unwrap_or_default(); + write_last_message_file(message, Some(output_file)); + if last_agent_message.is_none() { + eprintln!( + "Warning: no last agent message; wrote empty content to {}", + output_file.display() + ); + } +} + +fn write_last_message_file(contents: &str, last_message_path: Option<&Path>) { + if let Some(path) = last_message_path + && let Err(e) = std::fs::write(path, contents) + { + eprintln!("Failed to write last message file {path:?}: {e}"); + } +} diff --git a/codex-rs/exec/src/event_processor_with_human_output.rs b/codex-rs/exec/src/event_processor_with_human_output.rs new file mode 100644 index 0000000000000000000000000000000000000000..222bc485ab425ea4ad0bad8107cc5e2303bd20d9 --- /dev/null +++ b/codex-rs/exec/src/event_processor_with_human_output.rs @@ -0,0 +1,535 @@ +use std::io::IsTerminal; +use std::path::PathBuf; + +use codex_app_server_protocol::CommandExecutionStatus; +use codex_app_server_protocol::McpToolCallStatus; +use codex_app_server_protocol::PatchApplyStatus; +use codex_app_server_protocol::ServerNotification; +use codex_app_server_protocol::ThreadItem; +use codex_app_server_protocol::ThreadTokenUsage; +use codex_app_server_protocol::TurnStatus; +use codex_core::config::Config; +use codex_model_provider_info::WireApi; +use codex_protocol::num_format::format_with_separators; +use codex_protocol::protocol::SessionConfiguredEvent; +use codex_utils_path_uri::PathUri; +use codex_utils_sandbox_summary::summarize_permission_profile; +use owo_colors::OwoColorize; +use owo_colors::Style; + +use crate::event_processor::CodexStatus; +use crate::event_processor::EventProcessor; +use crate::event_processor::handle_last_message; + +pub(crate) struct EventProcessorWithHumanOutput { + bold: Style, + cyan: Style, + dimmed: Style, + green: Style, + italic: Style, + magenta: Style, + red: Style, + yellow: Style, + show_agent_reasoning: bool, + show_raw_agent_reasoning: bool, + last_message_path: Option, + final_message: Option, + final_message_rendered: bool, + emit_final_message_on_shutdown: bool, + last_total_token_usage: Option, +} + +impl EventProcessorWithHumanOutput { + pub(crate) fn create_with_ansi( + with_ansi: bool, + config: &Config, + last_message_path: Option, + ) -> Self { + let style = |styled: Style, plain: Style| if with_ansi { styled } else { plain }; + Self { + bold: style(Style::new().bold(), Style::new()), + cyan: style(Style::new().cyan(), Style::new()), + dimmed: style(Style::new().dimmed(), Style::new()), + green: style(Style::new().green(), Style::new()), + italic: style(Style::new().italic(), Style::new()), + magenta: style(Style::new().magenta(), Style::new()), + red: style(Style::new().red(), Style::new()), + yellow: style(Style::new().yellow(), Style::new()), + show_agent_reasoning: !config.hide_agent_reasoning, + show_raw_agent_reasoning: config.show_raw_agent_reasoning, + last_message_path, + final_message: None, + final_message_rendered: false, + emit_final_message_on_shutdown: false, + last_total_token_usage: None, + } + } + + fn render_item_started(&self, item: &ThreadItem) { + match item { + ThreadItem::CommandExecution { command, cwd, .. } => { + eprintln!( + "{}\n{} in {cwd}", + "exec".style(self.italic).style(self.magenta), + command.style(self.bold), + ); + } + ThreadItem::McpToolCall { server, tool, .. } => { + eprintln!( + "{} {} {}", + "mcp:".style(self.bold), + format!("{server}/{tool}").style(self.cyan), + "started".style(self.dimmed) + ); + } + ThreadItem::WebSearch(item) => { + eprintln!("{} {}", "web search:".style(self.bold), item.query); + } + ThreadItem::FileChange { .. } => { + eprintln!("{}", "apply patch".style(self.bold)); + } + ThreadItem::CollabAgentToolCall { tool, .. } => { + eprintln!("{} {:?}", "collab:".style(self.bold), tool); + } + _ => {} + } + } + + fn render_item_completed(&mut self, item: ThreadItem) { + match item { + ThreadItem::AgentMessage { text, .. } => { + eprintln!( + "{}\n{}", + "codex".style(self.italic).style(self.magenta), + text + ); + self.final_message = Some(text); + self.final_message_rendered = true; + } + ThreadItem::Reasoning { + summary, content, .. + } => { + if self.show_agent_reasoning + && let Some(text) = + reasoning_text(&summary, &content, self.show_raw_agent_reasoning) + && !text.trim().is_empty() + { + eprintln!("{}", text.style(self.dimmed)); + } + } + ThreadItem::CommandExecution { + command: _, + aggregated_output, + exit_code, + status, + duration_ms, + .. + } => { + let duration_suffix = duration_ms + .map(|duration_ms| format!(" in {duration_ms}ms")) + .unwrap_or_default(); + match status { + CommandExecutionStatus::Completed => { + eprintln!( + "{}", + format!(" succeeded{duration_suffix}:").style(self.green) + ); + } + CommandExecutionStatus::Failed => { + let exit_code = exit_code.unwrap_or(1); + eprintln!( + "{}", + format!(" exited {exit_code}{duration_suffix}:").style(self.red) + ); + } + CommandExecutionStatus::Declined => { + eprintln!( + "{}", + format!(" declined{duration_suffix}:").style(self.yellow) + ); + } + CommandExecutionStatus::InProgress => { + eprintln!( + "{}", + format!(" in progress{duration_suffix}:").style(self.dimmed) + ); + } + } + if let Some(output) = aggregated_output + && !output.trim().is_empty() + { + eprintln!("{output}"); + } + } + ThreadItem::FileChange { + changes, status, .. + } => { + let status_text = match status { + PatchApplyStatus::Completed => "completed", + PatchApplyStatus::Failed => "failed", + PatchApplyStatus::Declined => "declined", + PatchApplyStatus::InProgress => "in_progress", + }; + eprintln!("{} {}", "patch:".style(self.bold), status_text); + for change in changes { + eprintln!("{}", change.path.style(self.dimmed)); + } + } + ThreadItem::McpToolCall { + server, + tool, + status, + error, + .. + } => { + let status_text = match status { + McpToolCallStatus::Completed => "completed".style(self.green), + McpToolCallStatus::Failed => "failed".style(self.red), + McpToolCallStatus::InProgress => "in_progress".style(self.dimmed), + }; + eprintln!( + "{} {} {}", + "mcp:".style(self.bold), + format!("{server}/{tool}").style(self.cyan), + format!("({status_text})").style(self.dimmed) + ); + if let Some(error) = error { + eprintln!("{}", error.message.style(self.red)); + } + } + ThreadItem::WebSearch(item) => { + eprintln!("{} {}", "web search:".style(self.bold), item.query); + } + ThreadItem::ContextCompaction { .. } => { + eprintln!("{}", "context compacted".style(self.dimmed)); + } + _ => {} + } + } +} + +impl EventProcessor for EventProcessorWithHumanOutput { + fn print_config_summary( + &mut self, + config: &Config, + prompt: &str, + session_configured_event: &SessionConfiguredEvent, + ) { + const VERSION: &str = env!("CARGO_PKG_VERSION"); + eprintln!("OpenAI Codex v{VERSION}\n--------"); + for (key, value) in config_summary_entries(config, session_configured_event) { + eprintln!("{} {}", format!("{key}:").style(self.bold), value); + } + eprintln!("--------"); + eprintln!("{}\n{}", "user".style(self.cyan), prompt); + } + + fn process_server_notification(&mut self, notification: ServerNotification) -> CodexStatus { + match notification { + ServerNotification::ConfigWarning(notification) => { + let details = notification + .details + .map(|details| format!(" ({details})")) + .unwrap_or_default(); + eprintln!( + "{} {}{}", + "warning:".style(self.yellow).style(self.bold), + notification.summary, + details + ); + CodexStatus::Running + } + ServerNotification::Warning(notification) => self.process_warning(notification.message), + ServerNotification::AuthRecoveryStarted(notification) => { + eprintln!("{}", notification.message); + CodexStatus::Running + } + ServerNotification::AuthRecoveryCompleted(notification) => { + eprintln!("{}", notification.message.style(self.green)); + CodexStatus::Running + } + ServerNotification::Error(notification) => { + eprintln!( + "{} {}", + "ERROR:".style(self.red).style(self.bold), + notification.error + ); + CodexStatus::Running + } + ServerNotification::DeprecationNotice(notification) => { + eprintln!( + "{} {}", + "deprecated:".style(self.yellow).style(self.bold), + notification.summary + ); + if let Some(details) = notification.details { + eprintln!("{}", details.style(self.dimmed)); + } + CodexStatus::Running + } + ServerNotification::HookStarted(notification) => { + eprintln!( + "{} {}", + "hook:".style(self.bold), + format!("{:?}", notification.run.event_name).style(self.dimmed) + ); + CodexStatus::Running + } + ServerNotification::HookCompleted(notification) => { + eprintln!( + "{} {} {:?}", + "hook:".style(self.bold), + format!("{:?}", notification.run.event_name).style(self.dimmed), + notification.run.status + ); + CodexStatus::Running + } + ServerNotification::ItemStarted(notification) => { + self.render_item_started(¬ification.item); + CodexStatus::Running + } + ServerNotification::ItemCompleted(notification) => { + self.render_item_completed(notification.item); + CodexStatus::Running + } + ServerNotification::ModelRerouted(notification) => { + eprintln!( + "{} {} -> {}", + "model rerouted:".style(self.yellow).style(self.bold), + notification.from_model, + notification.to_model + ); + CodexStatus::Running + } + ServerNotification::ModelVerification(_) => CodexStatus::Running, + ServerNotification::ThreadTokenUsageUpdated(notification) => { + self.last_total_token_usage = Some(notification.token_usage); + CodexStatus::Running + } + ServerNotification::TurnCompleted(notification) => match notification.turn.status { + TurnStatus::Completed => { + let rendered_message = self + .final_message_rendered + .then(|| self.final_message.clone()) + .flatten(); + if let Some(final_message) = + final_message_from_turn_items(notification.turn.items.as_slice()) + { + self.final_message_rendered = + rendered_message.as_deref() == Some(final_message.as_str()); + self.final_message = Some(final_message); + } + self.emit_final_message_on_shutdown = true; + CodexStatus::InitiateShutdown + } + TurnStatus::Failed => { + self.final_message = None; + self.final_message_rendered = false; + self.emit_final_message_on_shutdown = false; + if let Some(error) = notification.turn.error { + eprintln!("{} {}", "ERROR:".style(self.red).style(self.bold), error); + } + CodexStatus::InitiateShutdown + } + TurnStatus::Interrupted => { + self.final_message = None; + self.final_message_rendered = false; + self.emit_final_message_on_shutdown = false; + eprintln!("{}", "turn interrupted".style(self.dimmed)); + CodexStatus::InitiateShutdown + } + TurnStatus::InProgress => CodexStatus::Running, + }, + ServerNotification::TurnDiffUpdated(notification) => { + if !notification.diff.trim().is_empty() { + eprintln!("{}", notification.diff); + } + CodexStatus::Running + } + ServerNotification::TurnPlanUpdated(notification) => { + if let Some(explanation) = notification.explanation { + eprintln!("{}", explanation.style(self.italic)); + } + for step in notification.plan { + match step.status { + codex_app_server_protocol::TurnPlanStepStatus::Completed => { + eprintln!(" {} {}", "✓".style(self.green), step.step); + } + codex_app_server_protocol::TurnPlanStepStatus::InProgress => { + eprintln!(" {} {}", "→".style(self.cyan), step.step); + } + codex_app_server_protocol::TurnPlanStepStatus::Pending => { + eprintln!( + " {} {}", + "•".style(self.dimmed), + step.step.style(self.dimmed) + ); + } + } + } + CodexStatus::Running + } + ServerNotification::TurnStarted(_) => CodexStatus::Running, + _ => CodexStatus::Running, + } + } + + fn process_warning(&mut self, message: String) -> CodexStatus { + eprintln!( + "{} {message}", + "warning:".style(self.yellow).style(self.bold) + ); + CodexStatus::Running + } + + fn print_final_output(&mut self) { + if self.emit_final_message_on_shutdown + && let Some(path) = self.last_message_path.as_deref() + { + handle_last_message(self.final_message.as_deref(), path); + } + + if let Some(usage) = &self.last_total_token_usage { + eprintln!( + "{}\n{}", + "tokens used".style(self.dimmed), + format_with_separators(blended_total(usage)) + ); + } + + #[allow(clippy::print_stdout)] + if should_print_final_message_to_stdout( + self.emit_final_message_on_shutdown + .then_some(self.final_message.as_deref()) + .flatten(), + std::io::stdout().is_terminal(), + std::io::stderr().is_terminal(), + ) && let Some(message) = self.final_message.as_deref() + { + println!("{message}"); + } else if should_print_final_message_to_tty( + self.emit_final_message_on_shutdown + .then_some(self.final_message.as_deref()) + .flatten(), + self.final_message_rendered, + std::io::stdout().is_terminal(), + std::io::stderr().is_terminal(), + ) && let Some(message) = self.final_message.as_deref() + { + eprintln!( + "{}\n{}", + "codex".style(self.italic).style(self.magenta), + message + ); + } + } +} + +fn config_summary_entries( + config: &Config, + session_configured_event: &SessionConfiguredEvent, +) -> Vec<(&'static str, String)> { + let permission_profile = config.permissions.effective_permission_profile(); + let mut entries = vec![ + ("workdir", config.cwd.display().to_string()), + ("model", session_configured_event.model.clone()), + ( + "provider", + session_configured_event.model_provider_id.clone(), + ), + ( + "approval", + config.permissions.approval_policy.value().to_string(), + ), + ( + "sandbox", + summarize_permission_profile( + &permission_profile, + &PathUri::from_abs_path(&config.cwd), + &config.effective_workspace_roots(), + ), + ), + ]; + if config.model_provider.wire_api == WireApi::Responses { + entries.push(( + "reasoning effort", + config + .model_reasoning_effort + .as_ref() + .map(std::string::ToString::to_string) + .unwrap_or_else(|| "none".to_string()), + )); + entries.push(( + "reasoning summaries", + config + .model_reasoning_summary + .map(|summary| summary.to_string()) + .unwrap_or_else(|| "none".to_string()), + )); + } + entries.push(( + "session id", + session_configured_event.session_id.to_string(), + )); + entries +} + +fn reasoning_text( + summary: &[String], + content: &[String], + show_raw_agent_reasoning: bool, +) -> Option { + let entries = if show_raw_agent_reasoning && !content.is_empty() { + content + } else { + summary + }; + if entries.is_empty() { + None + } else { + Some(entries.join("\n")) + } +} + +fn final_message_from_turn_items(items: &[ThreadItem]) -> Option { + items + .iter() + .rev() + .find_map(|item| match item { + ThreadItem::AgentMessage { text, .. } => Some(text.clone()), + _ => None, + }) + .or_else(|| { + items.iter().rev().find_map(|item| match item { + ThreadItem::Plan { text, .. } => Some(text.clone()), + _ => None, + }) + }) +} + +fn blended_total(usage: &ThreadTokenUsage) -> i64 { + let cached_input = usage.total.cached_input_tokens.max(0); + let non_cached_input = (usage.total.input_tokens - cached_input).max(0); + (non_cached_input + usage.total.output_tokens.max(0)).max(0) +} + +fn should_print_final_message_to_stdout( + final_message: Option<&str>, + stdout_is_terminal: bool, + stderr_is_terminal: bool, +) -> bool { + final_message.is_some() && !(stdout_is_terminal && stderr_is_terminal) +} + +fn should_print_final_message_to_tty( + final_message: Option<&str>, + final_message_rendered: bool, + stdout_is_terminal: bool, + stderr_is_terminal: bool, +) -> bool { + final_message.is_some() && !final_message_rendered && stdout_is_terminal && stderr_is_terminal +} + +#[cfg(test)] +#[path = "event_processor_with_human_output_tests.rs"] +mod tests; diff --git a/codex-rs/exec/src/event_processor_with_human_output_tests.rs b/codex-rs/exec/src/event_processor_with_human_output_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..594e2d3bb6aac3dd7088ab644bbc776e1fd2934f --- /dev/null +++ b/codex-rs/exec/src/event_processor_with_human_output_tests.rs @@ -0,0 +1,532 @@ +use codex_app_server_protocol::ServerNotification; +use codex_app_server_protocol::ThreadItem; +use codex_app_server_protocol::Turn; +use codex_app_server_protocol::TurnStatus; +use codex_core::config::ConfigBuilder; +use codex_protocol::SessionId; +use codex_protocol::ThreadId; +use codex_protocol::models::PermissionProfile; +use codex_protocol::permissions::FileSystemAccessMode; +use codex_protocol::permissions::FileSystemPath; +use codex_protocol::permissions::FileSystemSandboxEntry; +use codex_protocol::permissions::FileSystemSandboxPolicy; +use codex_protocol::permissions::NetworkSandboxPolicy; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::SessionConfiguredEvent; +use codex_utils_absolute_path::test_support::PathBufExt; +use codex_utils_absolute_path::test_support::test_path_buf; +use codex_utils_path_uri::PathUri; +use codex_utils_sandbox_summary::summarize_permission_profile; +use owo_colors::Style; +use pretty_assertions::assert_eq; + +use super::EventProcessorWithHumanOutput; +use super::config_summary_entries; +use super::final_message_from_turn_items; +use super::reasoning_text; +use super::should_print_final_message_to_stdout; +use super::should_print_final_message_to_tty; +use crate::event_processor::EventProcessor; + +#[test] +fn suppresses_final_stdout_message_when_both_streams_are_terminals() { + assert!(!should_print_final_message_to_stdout( + Some("hello"), + /*stdout_is_terminal*/ true, + /*stderr_is_terminal*/ true + )); +} + +#[test] +fn prints_final_stdout_message_when_stdout_is_not_terminal() { + assert!(should_print_final_message_to_stdout( + Some("hello"), + /*stdout_is_terminal*/ false, + /*stderr_is_terminal*/ true + )); +} + +#[test] +fn prints_final_stdout_message_when_stderr_is_not_terminal() { + assert!(should_print_final_message_to_stdout( + Some("hello"), + /*stdout_is_terminal*/ true, + /*stderr_is_terminal*/ false + )); +} + +#[test] +fn suppresses_final_stdout_message_when_missing() { + assert!(!should_print_final_message_to_stdout( + /*final_message*/ None, /*stdout_is_terminal*/ false, + /*stderr_is_terminal*/ false + )); +} + +#[test] +fn prints_final_tty_message_when_not_yet_rendered() { + assert!(should_print_final_message_to_tty( + Some("hello"), + /*final_message_rendered*/ false, + /*stdout_is_terminal*/ true, + /*stderr_is_terminal*/ true + )); +} + +#[test] +fn suppresses_final_tty_message_when_already_rendered() { + assert!(!should_print_final_message_to_tty( + Some("hello"), + /*final_message_rendered*/ true, + /*stdout_is_terminal*/ true, + /*stderr_is_terminal*/ true + )); +} + +#[test] +fn reasoning_text_prefers_summary_when_raw_reasoning_is_hidden() { + let text = reasoning_text( + &["summary".to_string()], + &["raw".to_string()], + /*show_raw_agent_reasoning*/ false, + ); + + assert_eq!(text.as_deref(), Some("summary")); +} + +#[test] +fn reasoning_text_uses_raw_content_when_enabled() { + let text = reasoning_text( + &["summary".to_string()], + &["raw".to_string()], + /*show_raw_agent_reasoning*/ true, + ); + + assert_eq!(text.as_deref(), Some("raw")); +} + +#[test] +fn summarizes_disabled_permission_profile_as_danger_full_access() { + let cwd = PathUri::from_abs_path(&test_path_buf("/tmp").abs()); + + assert_eq!( + summarize_permission_profile( + &PermissionProfile::Disabled, + &cwd, + std::slice::from_ref(&cwd), + ), + "danger-full-access" + ); +} + +#[test] +fn summarizes_external_permission_profile() { + let cwd = PathUri::from_abs_path(&test_path_buf("/tmp").abs()); + + assert_eq!( + summarize_permission_profile( + &PermissionProfile::External { + network: NetworkSandboxPolicy::Enabled, + }, + &cwd, + std::slice::from_ref(&cwd), + ), + "external-sandbox (network access enabled)" + ); +} + +#[test] +fn summarizes_managed_workspace_write_permission_profile() { + let cwd = test_path_buf("/tmp/project").abs(); + let cache_root = test_path_buf("/tmp/cache").abs(); + let profile = PermissionProfile::from_runtime_permissions( + &FileSystemSandboxPolicy::restricted(vec![ + FileSystemSandboxEntry { + path: cwd.clone().into(), + access: FileSystemAccessMode::Write, + missing_path_behavior: None, + }, + FileSystemSandboxEntry { + path: FileSystemPath::Path { + path: cache_root.clone().into(), + }, + access: FileSystemAccessMode::Write, + missing_path_behavior: None, + }, + ]), + NetworkSandboxPolicy::Restricted, + ); + + assert_eq!( + summarize_permission_profile( + &profile, + &PathUri::from_abs_path(&cwd), + &[cwd.clone().into(), cache_root.clone().into()], + ), + format!("workspace-write [workdir, {}]", cache_root.display()) + ); +} + +#[test] +fn summarizes_managed_read_only_permission_profile() { + let cwd = PathUri::from_abs_path(&test_path_buf("/tmp/project").abs()); + let profile = PermissionProfile::from_runtime_permissions( + &FileSystemSandboxPolicy::restricted(Vec::new()), + NetworkSandboxPolicy::Restricted, + ); + + assert_eq!( + summarize_permission_profile(&profile, &cwd, std::slice::from_ref(&cwd)), + "read-only" + ); +} + +#[tokio::test] +async fn config_summary_entries_include_runtime_workspace_roots() { + let codex_home = tempfile::tempdir().expect("create codex home"); + let cwd = tempfile::tempdir().expect("create cwd"); + let extra_root = tempfile::tempdir().expect("create extra root"); + let mut config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .fallback_cwd(Some(cwd.path().to_path_buf())) + .build() + .await + .expect("build default config"); + let cwd = cwd.path().to_path_buf().abs(); + let extra_root = extra_root.path().to_path_buf().abs(); + let expected_extra_root_name = extra_root + .file_name() + .expect("extra root should have file name") + .to_string_lossy() + .to_string(); + config.cwd = cwd.clone(); + config.workspace_roots = vec![cwd.clone(), extra_root]; + config + .permissions + .set_workspace_roots(config.workspace_roots.clone()); + config + .permissions + .set_permission_profile(PermissionProfile::workspace_write_with( + &[], + NetworkSandboxPolicy::Restricted, + /*exclude_tmpdir_env_var*/ true, + /*exclude_slash_tmp*/ true, + )) + .expect("set permission profile"); + + let session_configured_event = SessionConfiguredEvent { + session_id: SessionId::new(), + thread_id: ThreadId::new(), + forked_from_id: None, + parent_thread_id: None, + thread_source: None, + thread_name: None, + model: "gpt-5.4".to_string(), + model_provider_id: config.model_provider_id.clone(), + service_tier: None, + approval_policy: AskForApproval::Never, + approvals_reviewer: config.approvals_reviewer, + permission_profile: config.permissions.effective_permission_profile(), + active_permission_profile: None, + cwd, + reasoning_effort: None, + initial_messages: None, + network_proxy: None, + rollout_path: None, + }; + + let summary_entries = config_summary_entries(&config, &session_configured_event); + let sandbox_summary = summary_entries + .iter() + .find_map(|(key, value)| (*key == "sandbox").then_some(value)) + .expect("sandbox summary entry"); + assert!( + sandbox_summary.starts_with("workspace-write [workdir, ") + && sandbox_summary.contains(&expected_extra_root_name), + "expected runtime workspace root in sandbox summary: {summary_entries:?}" + ); +} + +#[test] +fn final_message_from_turn_items_uses_latest_agent_message() { + let message = final_message_from_turn_items(&[ + ThreadItem::AgentMessage { + id: "msg-1".to_string(), + text: "first".to_string(), + phase: None, + memory_citation: None, + delivery: None, + questions: None, + }, + ThreadItem::Plan { + id: "plan-1".to_string(), + text: "plan".to_string(), + }, + ThreadItem::AgentMessage { + id: "msg-2".to_string(), + text: "second".to_string(), + phase: None, + memory_citation: None, + delivery: None, + questions: None, + }, + ]); + + assert_eq!(message.as_deref(), Some("second")); +} + +#[test] +fn final_message_from_turn_items_falls_back_to_latest_plan() { + let message = final_message_from_turn_items(&[ + ThreadItem::Reasoning { + id: "reasoning-1".to_string(), + summary: vec!["inspect".to_string()], + content: Vec::new(), + }, + ThreadItem::Plan { + id: "plan-1".to_string(), + text: "first plan".to_string(), + }, + ThreadItem::Plan { + id: "plan-2".to_string(), + text: "final plan".to_string(), + }, + ]); + + assert_eq!(message.as_deref(), Some("final plan")); +} + +#[test] +fn turn_completed_recovers_final_message_from_turn_items() { + let mut processor = EventProcessorWithHumanOutput { + bold: Style::new(), + cyan: Style::new(), + dimmed: Style::new(), + green: Style::new(), + italic: Style::new(), + magenta: Style::new(), + red: Style::new(), + yellow: Style::new(), + show_agent_reasoning: true, + show_raw_agent_reasoning: false, + last_message_path: None, + final_message: None, + final_message_rendered: false, + emit_final_message_on_shutdown: false, + last_total_token_usage: None, + }; + + let status = processor.process_server_notification(ServerNotification::TurnCompleted( + codex_app_server_protocol::TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: vec![ThreadItem::AgentMessage { + id: "msg-1".to_string(), + text: "final answer".to_string(), + phase: None, + memory_citation: None, + delivery: None, + questions: None, + }], + status: TurnStatus::Completed, + error: None, + started_at: None, + completed_at: Some(0), + duration_ms: None, + }, + }, + )); + + assert_eq!( + status, + crate::event_processor::CodexStatus::InitiateShutdown + ); + assert_eq!(processor.final_message.as_deref(), Some("final answer")); +} + +#[test] +fn turn_completed_overwrites_stale_final_message_from_turn_items() { + let mut processor = EventProcessorWithHumanOutput { + bold: Style::new(), + cyan: Style::new(), + dimmed: Style::new(), + green: Style::new(), + italic: Style::new(), + magenta: Style::new(), + red: Style::new(), + yellow: Style::new(), + show_agent_reasoning: true, + show_raw_agent_reasoning: false, + last_message_path: None, + final_message: Some("stale answer".to_string()), + final_message_rendered: true, + emit_final_message_on_shutdown: false, + last_total_token_usage: None, + }; + + let status = processor.process_server_notification(ServerNotification::TurnCompleted( + codex_app_server_protocol::TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: vec![ThreadItem::AgentMessage { + id: "msg-1".to_string(), + text: "final answer".to_string(), + phase: None, + memory_citation: None, + delivery: None, + questions: None, + }], + status: TurnStatus::Completed, + error: None, + started_at: None, + completed_at: Some(0), + duration_ms: None, + }, + }, + )); + + assert_eq!( + status, + crate::event_processor::CodexStatus::InitiateShutdown + ); + assert_eq!(processor.final_message.as_deref(), Some("final answer")); + assert!(!processor.final_message_rendered); +} + +#[test] +fn turn_completed_preserves_streamed_final_message_when_turn_items_are_empty() { + let mut processor = EventProcessorWithHumanOutput { + bold: Style::new(), + cyan: Style::new(), + dimmed: Style::new(), + green: Style::new(), + italic: Style::new(), + magenta: Style::new(), + red: Style::new(), + yellow: Style::new(), + show_agent_reasoning: true, + show_raw_agent_reasoning: false, + last_message_path: None, + final_message: Some("streamed answer".to_string()), + final_message_rendered: false, + emit_final_message_on_shutdown: false, + last_total_token_usage: None, + }; + + let status = processor.process_server_notification(ServerNotification::TurnCompleted( + codex_app_server_protocol::TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: Vec::new(), + status: TurnStatus::Completed, + error: None, + started_at: None, + completed_at: Some(0), + duration_ms: None, + }, + }, + )); + + assert_eq!( + status, + crate::event_processor::CodexStatus::InitiateShutdown + ); + assert_eq!(processor.final_message.as_deref(), Some("streamed answer")); + assert!(processor.emit_final_message_on_shutdown); +} + +#[test] +fn turn_failed_clears_stale_final_message() { + let mut processor = EventProcessorWithHumanOutput { + bold: Style::new(), + cyan: Style::new(), + dimmed: Style::new(), + green: Style::new(), + italic: Style::new(), + magenta: Style::new(), + red: Style::new(), + yellow: Style::new(), + show_agent_reasoning: true, + show_raw_agent_reasoning: false, + last_message_path: None, + final_message: Some("partial answer".to_string()), + final_message_rendered: true, + emit_final_message_on_shutdown: true, + last_total_token_usage: None, + }; + + let status = processor.process_server_notification(ServerNotification::TurnCompleted( + codex_app_server_protocol::TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: Vec::new(), + status: TurnStatus::Failed, + error: None, + started_at: None, + completed_at: Some(0), + duration_ms: None, + }, + }, + )); + + assert_eq!( + status, + crate::event_processor::CodexStatus::InitiateShutdown + ); + assert_eq!(processor.final_message, None); + assert!(!processor.final_message_rendered); + assert!(!processor.emit_final_message_on_shutdown); +} + +#[test] +fn turn_interrupted_clears_stale_final_message() { + let mut processor = EventProcessorWithHumanOutput { + bold: Style::new(), + cyan: Style::new(), + dimmed: Style::new(), + green: Style::new(), + italic: Style::new(), + magenta: Style::new(), + red: Style::new(), + yellow: Style::new(), + show_agent_reasoning: true, + show_raw_agent_reasoning: false, + last_message_path: None, + final_message: Some("partial answer".to_string()), + final_message_rendered: true, + emit_final_message_on_shutdown: true, + last_total_token_usage: None, + }; + + let status = processor.process_server_notification(ServerNotification::TurnCompleted( + codex_app_server_protocol::TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: Vec::new(), + status: TurnStatus::Interrupted, + error: None, + started_at: None, + completed_at: Some(0), + duration_ms: None, + }, + }, + )); + + assert_eq!( + status, + crate::event_processor::CodexStatus::InitiateShutdown + ); + assert_eq!(processor.final_message, None); + assert!(!processor.final_message_rendered); + assert!(!processor.emit_final_message_on_shutdown); +} diff --git a/codex-rs/exec/src/event_processor_with_jsonl_output.rs b/codex-rs/exec/src/event_processor_with_jsonl_output.rs new file mode 100644 index 0000000000000000000000000000000000000000..488cbc52e9ca31cea1203fb8ba923af50e2b39df --- /dev/null +++ b/codex-rs/exec/src/event_processor_with_jsonl_output.rs @@ -0,0 +1,635 @@ +use std::collections::HashMap; +use std::path::PathBuf; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use codex_app_server_protocol::CollabAgentTool; +use codex_app_server_protocol::CollabAgentToolCallStatus; +use codex_app_server_protocol::CommandExecutionStatus; +use codex_app_server_protocol::McpToolCallStatus; +use codex_app_server_protocol::PatchApplyStatus; +use codex_app_server_protocol::PatchChangeKind; +use codex_app_server_protocol::ServerNotification; +use codex_app_server_protocol::ThreadItem; +use codex_app_server_protocol::ThreadTokenUsage; +use codex_app_server_protocol::TurnStatus; +use codex_core::config::Config; +use codex_protocol::models::WebSearchAction; +use codex_protocol::protocol::SessionConfiguredEvent; +use serde_json::json; + +pub use crate::event_processor::CodexStatus; +use crate::event_processor::EventProcessor; +use crate::event_processor::handle_last_message; +use crate::exec_events::AgentMessageItem; +use crate::exec_events::CollabAgentState; +use crate::exec_events::CollabAgentStatus; +use crate::exec_events::CollabTool; +use crate::exec_events::CollabToolCallItem; +use crate::exec_events::CollabToolCallStatus; +use crate::exec_events::CommandExecutionItem; +use crate::exec_events::CommandExecutionStatus as ExecCommandExecutionStatus; +use crate::exec_events::ErrorItem; +use crate::exec_events::FileChangeItem; +use crate::exec_events::FileUpdateChange; +use crate::exec_events::ItemCompletedEvent; +use crate::exec_events::ItemStartedEvent; +use crate::exec_events::ItemUpdatedEvent; +use crate::exec_events::McpToolCallItem; +use crate::exec_events::McpToolCallItemError; +use crate::exec_events::McpToolCallItemResult; +use crate::exec_events::McpToolCallStatus as ExecMcpToolCallStatus; +use crate::exec_events::PatchApplyStatus as ExecPatchApplyStatus; +use crate::exec_events::PatchChangeKind as ExecPatchChangeKind; +use crate::exec_events::ReasoningItem; +use crate::exec_events::ThreadErrorEvent; +use crate::exec_events::ThreadEvent; +use crate::exec_events::ThreadItem as ExecThreadItem; +use crate::exec_events::ThreadItemDetails; +use crate::exec_events::ThreadStartedEvent; +use crate::exec_events::TodoItem; +use crate::exec_events::TodoListItem; +use crate::exec_events::TurnCompletedEvent; +use crate::exec_events::TurnFailedEvent; +use crate::exec_events::TurnStartedEvent; +use crate::exec_events::Usage; +use crate::exec_events::WebSearchItem; + +pub struct EventProcessorWithJsonOutput { + last_message_path: Option, + next_item_id: AtomicU64, + raw_to_exec_item_id: HashMap, + running_todo_list: Option, + last_total_token_usage: Option, + last_critical_error: Option, + final_message: Option, + emit_final_message_on_shutdown: bool, +} + +#[derive(Debug, Clone)] +struct RunningTodoList { + item_id: String, + items: Vec, +} + +#[derive(Debug, PartialEq)] +pub struct CollectedThreadEvents { + pub events: Vec, + pub status: CodexStatus, +} + +impl EventProcessorWithJsonOutput { + pub fn new(last_message_path: Option) -> Self { + Self { + last_message_path, + next_item_id: AtomicU64::new(0), + raw_to_exec_item_id: HashMap::new(), + running_todo_list: None, + last_total_token_usage: None, + last_critical_error: None, + final_message: None, + emit_final_message_on_shutdown: false, + } + } + + pub fn final_message(&self) -> Option<&str> { + self.final_message.as_deref() + } + + fn next_item_id(&self) -> String { + format!("item_{}", self.next_item_id.fetch_add(1, Ordering::SeqCst)) + } + + #[allow(clippy::print_stdout)] + fn emit(&self, event: ThreadEvent) { + println!( + "{}", + serde_json::to_string(&event).unwrap_or_else(|err| { + json!({ + "type": "error", + "message": format!("failed to serialize exec json event: {err}"), + }) + .to_string() + }) + ); + } + + fn usage_from_last_total(&self) -> Usage { + let Some(usage) = self.last_total_token_usage.as_ref() else { + return Usage::default(); + }; + Usage { + input_tokens: usage.total.input_tokens, + cached_input_tokens: usage.total.cached_input_tokens, + cache_write_input_tokens: usage.total.cache_write_input_tokens, + output_tokens: usage.total.output_tokens, + reasoning_output_tokens: usage.total.reasoning_output_tokens, + } + } + + pub fn map_todo_items(plan: &[codex_app_server_protocol::TurnPlanStep]) -> Vec { + plan.iter() + .map(|step| TodoItem { + text: step.step.clone(), + completed: matches!( + step.status, + codex_app_server_protocol::TurnPlanStepStatus::Completed + ), + }) + .collect() + } + + fn map_item_with_id( + item: ThreadItem, + make_id: impl FnOnce() -> String, + ) -> Option { + match item { + ThreadItem::AgentMessage { text, .. } => Some(ExecThreadItem { + id: make_id(), + details: ThreadItemDetails::AgentMessage(AgentMessageItem { text }), + }), + ThreadItem::Reasoning { summary, .. } => { + let text = summary.join("\n"); + if text.trim().is_empty() { + return None; + } + Some(ExecThreadItem { + id: make_id(), + details: ThreadItemDetails::Reasoning(ReasoningItem { text }), + }) + } + ThreadItem::CommandExecution { + command, + aggregated_output, + exit_code, + status, + .. + } => Some(ExecThreadItem { + id: make_id(), + details: ThreadItemDetails::CommandExecution(CommandExecutionItem { + command, + aggregated_output: aggregated_output.unwrap_or_default(), + exit_code, + status: match status { + CommandExecutionStatus::InProgress => ExecCommandExecutionStatus::InProgress, + CommandExecutionStatus::Completed => ExecCommandExecutionStatus::Completed, + CommandExecutionStatus::Failed => ExecCommandExecutionStatus::Failed, + CommandExecutionStatus::Declined => ExecCommandExecutionStatus::Declined, + }, + }), + }), + ThreadItem::FileChange { + changes, status, .. + } => Some(ExecThreadItem { + id: make_id(), + details: ThreadItemDetails::FileChange(FileChangeItem { + changes: changes + .into_iter() + .map(|change| FileUpdateChange { + path: change.path, + kind: match change.kind { + PatchChangeKind::Add => ExecPatchChangeKind::Add, + PatchChangeKind::Delete => ExecPatchChangeKind::Delete, + PatchChangeKind::Update { .. } => ExecPatchChangeKind::Update, + }, + }) + .collect(), + status: match status { + PatchApplyStatus::InProgress => ExecPatchApplyStatus::InProgress, + PatchApplyStatus::Completed => ExecPatchApplyStatus::Completed, + PatchApplyStatus::Failed | PatchApplyStatus::Declined => { + ExecPatchApplyStatus::Failed + } + }, + }), + }), + ThreadItem::McpToolCall { + server, + tool, + status, + arguments, + result, + error, + .. + } => Some(ExecThreadItem { + id: make_id(), + details: ThreadItemDetails::McpToolCall(McpToolCallItem { + server, + tool, + status: match status { + McpToolCallStatus::InProgress => ExecMcpToolCallStatus::InProgress, + McpToolCallStatus::Completed => ExecMcpToolCallStatus::Completed, + McpToolCallStatus::Failed => ExecMcpToolCallStatus::Failed, + }, + arguments, + result: result.map(|result| McpToolCallItemResult { + content: result.content, + meta: result.meta, + structured_content: result.structured_content, + }), + error: error.map(|error| McpToolCallItemError { + message: error.message, + }), + }), + }), + ThreadItem::CollabAgentToolCall { + tool, + sender_thread_id, + receiver_thread_ids, + prompt, + agents_states, + status, + .. + } => Some(ExecThreadItem { + id: make_id(), + details: ThreadItemDetails::CollabToolCall(CollabToolCallItem { + tool: match tool { + CollabAgentTool::SendMessage + | CollabAgentTool::FollowupTask + | CollabAgentTool::InterruptAgent + | CollabAgentTool::ListAgents => return None, + CollabAgentTool::SpawnAgent => CollabTool::SpawnAgent, + CollabAgentTool::SendInput => CollabTool::SendInput, + CollabAgentTool::ResumeAgent => CollabTool::Wait, + CollabAgentTool::Wait => CollabTool::Wait, + CollabAgentTool::CloseAgent => CollabTool::CloseAgent, + }, + sender_thread_id, + receiver_thread_ids, + prompt, + agents_states: agents_states + .into_iter() + .map(|(thread_id, state)| { + ( + thread_id, + CollabAgentState { + status: match state.status { + codex_app_server_protocol::CollabAgentStatus::PendingInit => { + CollabAgentStatus::PendingInit + } + codex_app_server_protocol::CollabAgentStatus::Running => { + CollabAgentStatus::Running + } + codex_app_server_protocol::CollabAgentStatus::Interrupted => { + CollabAgentStatus::Interrupted + } + codex_app_server_protocol::CollabAgentStatus::Completed => { + CollabAgentStatus::Completed + } + codex_app_server_protocol::CollabAgentStatus::Errored => { + CollabAgentStatus::Errored + } + codex_app_server_protocol::CollabAgentStatus::Shutdown => { + CollabAgentStatus::Shutdown + } + codex_app_server_protocol::CollabAgentStatus::NotFound => { + CollabAgentStatus::NotFound + } + }, + message: state.message, + }, + ) + }) + .collect(), + status: match status { + CollabAgentToolCallStatus::InProgress => CollabToolCallStatus::InProgress, + CollabAgentToolCallStatus::Completed => CollabToolCallStatus::Completed, + CollabAgentToolCallStatus::Failed => CollabToolCallStatus::Failed, + CollabAgentToolCallStatus::Interrupted => return None, + }, + }), + }), + ThreadItem::WebSearch(item) => Some(ExecThreadItem { + id: make_id(), + details: ThreadItemDetails::WebSearch(WebSearchItem { + id: item.id, + query: item.query, + action: match item.action { + Some(action) => serde_json::from_value( + serde_json::to_value(action).unwrap_or_else(|_| json!("other")), + ) + .unwrap_or(WebSearchAction::Other), + None => WebSearchAction::Other, + }, + }), + }), + _ => None, + } + } + + fn started_item_id(&mut self, raw_id: &str) -> String { + if let Some(existing) = self.raw_to_exec_item_id.get(raw_id) { + return existing.clone(); + } + let exec_id = self.next_item_id(); + self.raw_to_exec_item_id + .insert(raw_id.to_string(), exec_id.clone()); + exec_id + } + + fn completed_item_id(&mut self, raw_id: &str) -> String { + self.raw_to_exec_item_id + .remove(raw_id) + .unwrap_or_else(|| self.next_item_id()) + } + + fn map_started_item(&mut self, item: ThreadItem) -> Option { + match item { + ThreadItem::AgentMessage { .. } | ThreadItem::Reasoning { .. } => None, + other => { + let raw_id = other.id().to_string(); + Self::map_item_with_id(other, || self.started_item_id(&raw_id)) + } + } + } + + fn map_completed_item_mut(&mut self, item: ThreadItem) -> Option { + if let ThreadItem::Reasoning { summary, .. } = &item + && summary.join("\n").trim().is_empty() + { + return None; + } + match &item { + ThreadItem::AgentMessage { .. } | ThreadItem::Reasoning { .. } => { + Self::map_item_with_id(item, || self.next_item_id()) + } + other => { + let raw_id = other.id().to_string(); + Self::map_item_with_id(item, || self.completed_item_id(&raw_id)) + } + } + } + + fn reconcile_unfinished_started_items( + &mut self, + turn_items: &[ThreadItem], + ) -> Vec { + turn_items + .iter() + .filter_map(|item| { + let raw_id = item.id().to_string(); + if !self.raw_to_exec_item_id.contains_key(&raw_id) { + return None; + } + self.map_completed_item_mut(item.clone()) + .map(|item| ThreadEvent::ItemCompleted(ItemCompletedEvent { item })) + }) + .collect() + } + + fn final_message_from_turn_items(items: &[ThreadItem]) -> Option { + items + .iter() + .rev() + .find_map(|item| match item { + ThreadItem::AgentMessage { text, .. } => Some(text.clone()), + _ => None, + }) + .or_else(|| { + items.iter().rev().find_map(|item| match item { + ThreadItem::Plan { text, .. } => Some(text.clone()), + _ => None, + }) + }) + } + + pub fn thread_started_event(session_configured: &SessionConfiguredEvent) -> ThreadEvent { + ThreadEvent::ThreadStarted(ThreadStartedEvent { + thread_id: session_configured.thread_id.to_string(), + }) + } + + pub fn collect_warning(&mut self, message: String) -> CollectedThreadEvents { + CollectedThreadEvents { + events: vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: self.next_item_id(), + details: ThreadItemDetails::Error(ErrorItem { message }), + }, + })], + status: CodexStatus::Running, + } + } + + pub fn collect_thread_events( + &mut self, + notification: ServerNotification, + ) -> CollectedThreadEvents { + let mut events = Vec::new(); + let status = match notification { + ServerNotification::ConfigWarning(notification) => { + let message = match notification.details { + Some(details) if !details.is_empty() => { + format!("{} ({details})", notification.summary) + } + _ => notification.summary, + }; + events.push(ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: self.next_item_id(), + details: ThreadItemDetails::Error(ErrorItem { message }), + }, + })); + CodexStatus::Running + } + ServerNotification::Warning(notification) => { + let warning = self.collect_warning(notification.message); + events.extend(warning.events); + warning.status + } + ServerNotification::Error(notification) => { + let message = match notification.error.additional_details { + Some(details) if !details.is_empty() => { + format!("{} ({details})", notification.error.message) + } + _ => notification.error.message, + }; + let error = ThreadErrorEvent { message }; + self.last_critical_error = Some(error.clone()); + events.push(ThreadEvent::Error(error)); + CodexStatus::Running + } + ServerNotification::DeprecationNotice(notification) => { + let message = match notification.details { + Some(details) if !details.is_empty() => { + format!("{} ({details})", notification.summary) + } + _ => notification.summary, + }; + events.push(ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: self.next_item_id(), + details: ThreadItemDetails::Error(ErrorItem { message }), + }, + })); + CodexStatus::Running + } + ServerNotification::HookStarted(_) | ServerNotification::HookCompleted(_) => { + CodexStatus::Running + } + ServerNotification::ItemStarted(notification) => { + if let Some(item) = self.map_started_item(notification.item) { + events.push(ThreadEvent::ItemStarted(ItemStartedEvent { item })); + } + CodexStatus::Running + } + ServerNotification::ItemCompleted(notification) => { + if let Some(item) = self.map_completed_item_mut(notification.item) { + if let ThreadItemDetails::AgentMessage(AgentMessageItem { text }) = + &item.details + { + self.final_message = Some(text.clone()); + } + events.push(ThreadEvent::ItemCompleted(ItemCompletedEvent { item })); + } + CodexStatus::Running + } + ServerNotification::ModelRerouted(notification) => { + events.push(ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: self.next_item_id(), + details: ThreadItemDetails::Error(ErrorItem { + message: format!( + "model rerouted: {} -> {} ({:?})", + notification.from_model, notification.to_model, notification.reason + ), + }), + }, + })); + CodexStatus::Running + } + ServerNotification::ModelVerification(_) => CodexStatus::Running, + ServerNotification::ThreadTokenUsageUpdated(notification) => { + self.last_total_token_usage = Some(notification.token_usage); + CodexStatus::Running + } + ServerNotification::TurnCompleted(notification) => { + if let Some(running) = self.running_todo_list.take() { + events.push(ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: running.item_id, + details: ThreadItemDetails::TodoList(TodoListItem { + items: running.items, + }), + }, + })); + } + events.extend(self.reconcile_unfinished_started_items(¬ification.turn.items)); + match notification.turn.status { + TurnStatus::Completed => { + if let Some(final_message) = + Self::final_message_from_turn_items(notification.turn.items.as_slice()) + { + self.final_message = Some(final_message); + } + self.emit_final_message_on_shutdown = true; + events.push(ThreadEvent::TurnCompleted(TurnCompletedEvent { + usage: self.usage_from_last_total(), + })); + CodexStatus::InitiateShutdown + } + TurnStatus::Failed => { + self.final_message = None; + self.emit_final_message_on_shutdown = false; + let error = notification + .turn + .error + .map(|error| ThreadErrorEvent { + message: match error.additional_details { + Some(details) if !details.is_empty() => { + format!("{} ({details})", error.message) + } + _ => error.message, + }, + }) + .or_else(|| self.last_critical_error.clone()) + .unwrap_or_else(|| ThreadErrorEvent { + message: "turn failed".to_string(), + }); + events.push(ThreadEvent::TurnFailed(TurnFailedEvent { error })); + CodexStatus::InitiateShutdown + } + TurnStatus::Interrupted => { + self.final_message = None; + self.emit_final_message_on_shutdown = false; + CodexStatus::InitiateShutdown + } + TurnStatus::InProgress => CodexStatus::Running, + } + } + ServerNotification::TurnDiffUpdated(_) => CodexStatus::Running, + ServerNotification::TurnPlanUpdated(notification) => { + let items = Self::map_todo_items(¬ification.plan); + if let Some(running) = self.running_todo_list.as_mut() { + running.items = items.clone(); + let item_id = running.item_id.clone(); + events.push(ThreadEvent::ItemUpdated(ItemUpdatedEvent { + item: ExecThreadItem { + id: item_id, + details: ThreadItemDetails::TodoList(TodoListItem { items }), + }, + })); + } else { + let item_id = self.next_item_id(); + self.running_todo_list = Some(RunningTodoList { + item_id: item_id.clone(), + items: items.clone(), + }); + events.push(ThreadEvent::ItemStarted(ItemStartedEvent { + item: ExecThreadItem { + id: item_id, + details: ThreadItemDetails::TodoList(TodoListItem { items }), + }, + })); + } + CodexStatus::Running + } + ServerNotification::TurnStarted(_) => { + events.push(ThreadEvent::TurnStarted(TurnStartedEvent {})); + CodexStatus::Running + } + _ => CodexStatus::Running, + }; + + CollectedThreadEvents { events, status } + } +} + +impl EventProcessor for EventProcessorWithJsonOutput { + fn print_config_summary( + &mut self, + _: &Config, + _: &str, + session_configured: &SessionConfiguredEvent, + ) { + self.emit(Self::thread_started_event(session_configured)); + } + + fn process_server_notification(&mut self, notification: ServerNotification) -> CodexStatus { + let collected = self.collect_thread_events(notification); + for event in collected.events { + self.emit(event); + } + collected.status + } + + fn process_warning(&mut self, message: String) -> CodexStatus { + let collected = self.collect_warning(message); + for event in collected.events { + self.emit(event); + } + collected.status + } + + fn print_final_output(&mut self) { + if self.emit_final_message_on_shutdown + && let Some(path) = self.last_message_path.as_deref() + { + handle_last_message(self.final_message.as_deref(), path); + } + } +} + +#[cfg(test)] +#[path = "event_processor_with_jsonl_output_tests.rs"] +mod tests; diff --git a/codex-rs/exec/src/event_processor_with_jsonl_output_tests.rs b/codex-rs/exec/src/event_processor_with_jsonl_output_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..5109848a6ed1c4de6524b9327f2efca5ef83328d --- /dev/null +++ b/codex-rs/exec/src/event_processor_with_jsonl_output_tests.rs @@ -0,0 +1,144 @@ +use super::*; +use pretty_assertions::assert_eq; +use serde_json::json; +use tempfile::tempdir; + +#[test] +fn failed_turn_does_not_overwrite_output_last_message_file() { + let tempdir = tempdir().expect("create tempdir"); + let output_path = tempdir.path().join("last-message.txt"); + std::fs::write(&output_path, "keep existing contents").expect("seed output file"); + + let mut processor = EventProcessorWithJsonOutput::new(Some(output_path.clone())); + + let collected = processor.collect_thread_events(ServerNotification::ItemCompleted( + codex_app_server_protocol::ItemCompletedNotification { + item: ThreadItem::AgentMessage { + id: "msg-1".to_string(), + text: "partial answer".to_string(), + phase: None, + memory_citation: None, + delivery: None, + questions: None, + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!(collected.status, CodexStatus::Running); + assert_eq!(processor.final_message(), Some("partial answer")); + + let status = processor.process_server_notification(ServerNotification::TurnCompleted( + codex_app_server_protocol::TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: codex_app_server_protocol::Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: Vec::new(), + status: TurnStatus::Failed, + error: Some(codex_app_server_protocol::TurnError { + misalignment: None, + message: "turn failed".to_string(), + additional_details: None, + codex_error_info: None, + }), + started_at: None, + completed_at: Some(0), + duration_ms: None, + }, + }, + )); + + assert_eq!(status, CodexStatus::InitiateShutdown); + assert_eq!(processor.final_message(), None); + + EventProcessor::print_final_output(&mut processor); + + assert_eq!( + std::fs::read_to_string(&output_path).expect("read output file"), + "keep existing contents" + ); +} + +#[test] +fn runtime_warning_emits_a_non_fatal_error_item() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let collected = processor.collect_thread_events(ServerNotification::Warning( + codex_app_server_protocol::WarningNotification { + thread_id: Some("thread-1".to_string()), + message: "invalid global instructions".to_string(), + }, + )); + + assert_eq!( + collected, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::Error(ErrorItem { + message: "invalid global instructions".to_string(), + }), + }, + })], + status: CodexStatus::Running, + } + ); +} + +#[test] +fn mcp_tool_call_result_preserves_meta_in_jsonl_event() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let collected = processor.collect_thread_events(ServerNotification::ItemCompleted( + codex_app_server_protocol::ItemCompletedNotification { + item: ThreadItem::McpToolCall { + id: "mcp-1".to_string(), + server: "search service".to_string(), + tool: "web_run".to_string(), + status: McpToolCallStatus::Completed, + arguments: json!({"search_query": [{"q": "OpenAI Codex CLI documentation"}]}), + app_context: None, + mcp_app_resource_uri: None, + mcp_app_ui: None, + plugin_id: None, + read_only_hint: None, + result: Some(Box::new(codex_app_server_protocol::McpToolCallResult { + content: vec![json!({"type": "text", "text": "search result"})], + structured_content: None, + meta: Some(json!({"raw_messages": [{"ref_id": "turn0search0"}]})), + })), + error: None, + duration_ms: Some(42), + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!(collected.status, CodexStatus::Running); + assert_eq!(collected.events.len(), 1); + + let ThreadEvent::ItemCompleted(ItemCompletedEvent { item }) = &collected.events[0] else { + panic!("expected item.completed event"); + }; + let ThreadItemDetails::McpToolCall(item) = &item.details else { + panic!("expected MCP tool call item"); + }; + let result = item.result.as_ref().expect("expected MCP tool result"); + assert_eq!( + result.meta, + Some(json!({"raw_messages": [{"ref_id": "turn0search0"}]})) + ); + + let serialized = serde_json::to_value(&collected.events[0]).expect("serialize event"); + assert_eq!( + serialized["item"]["result"]["_meta"], + json!({"raw_messages": [{"ref_id": "turn0search0"}]}) + ); + assert!(serialized["item"]["result"].get("meta").is_none()); +} diff --git a/codex-rs/exec/src/exec_events.rs b/codex-rs/exec/src/exec_events.rs new file mode 100644 index 0000000000000000000000000000000000000000..30df7f176a02c5283405a70fac2d5ef9acdcb66e --- /dev/null +++ b/codex-rs/exec/src/exec_events.rs @@ -0,0 +1,320 @@ +use codex_protocol::models::WebSearchAction; +use serde::Deserialize; +use serde::Serialize; +use serde_json::Value as JsonValue; +use std::collections::HashMap; +use ts_rs::TS; + +/// Top-level JSONL events emitted by codex exec +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +#[serde(tag = "type")] +pub enum ThreadEvent { + /// Emitted when a new thread is started as the first event. + #[serde(rename = "thread.started")] + ThreadStarted(ThreadStartedEvent), + /// Emitted when a turn is started by sending a new prompt to the model. + /// A turn encompasses all events that happen while agent is processing the prompt. + #[serde(rename = "turn.started")] + TurnStarted(TurnStartedEvent), + /// Emitted when a turn is completed. Typically right after the assistant's response. + #[serde(rename = "turn.completed")] + TurnCompleted(TurnCompletedEvent), + /// Indicates that a turn failed with an error. + #[serde(rename = "turn.failed")] + TurnFailed(TurnFailedEvent), + /// Emitted when a new item is added to the thread. Typically the item will be in an "in progress" state. + #[serde(rename = "item.started")] + ItemStarted(ItemStartedEvent), + /// Emitted when an item is updated. + #[serde(rename = "item.updated")] + ItemUpdated(ItemUpdatedEvent), + /// Signals that an item has reached a terminal state—either success or failure. + #[serde(rename = "item.completed")] + ItemCompleted(ItemCompletedEvent), + /// Represents an unrecoverable error emitted directly by the event stream. + #[serde(rename = "error")] + Error(ThreadErrorEvent), +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct ThreadStartedEvent { + /// The identified of the new thread. Can be used to resume the thread later. + pub thread_id: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS, Default)] + +pub struct TurnStartedEvent {} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct TurnCompletedEvent { + pub usage: Usage, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct TurnFailedEvent { + pub error: ThreadErrorEvent, +} + +/// Describes the usage of tokens during a turn. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS, Default)] +pub struct Usage { + /// The number of input tokens used during the turn. + pub input_tokens: i64, + /// The number of cached input tokens used during the turn. + pub cached_input_tokens: i64, + /// The number of input tokens written to the prompt cache during the turn. + #[serde(default)] + pub cache_write_input_tokens: i64, + /// The number of output tokens used during the turn. + pub output_tokens: i64, + /// The number of reasoning output tokens used during the turn. + pub reasoning_output_tokens: i64, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct ItemStartedEvent { + pub item: ThreadItem, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct ItemCompletedEvent { + pub item: ThreadItem, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct ItemUpdatedEvent { + pub item: ThreadItem, +} + +/// Fatal error emitted by the stream. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct ThreadErrorEvent { + pub message: String, +} + +/// Canonical representation of a thread item and its domain-specific payload. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct ThreadItem { + pub id: String, + #[serde(flatten)] + pub details: ThreadItemDetails, +} + +/// Typed payloads for each supported thread item type. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ThreadItemDetails { + /// Response from the agent. + /// Either a natural-language response or a JSON string when structured output is requested. + AgentMessage(AgentMessageItem), + /// Agent's reasoning summary. + Reasoning(ReasoningItem), + /// Tracks a command executed by the agent. The item starts when the command is + /// spawned, and completes when the process exits with an exit code. + CommandExecution(CommandExecutionItem), + /// Represents a set of file changes by the agent. The item is emitted only as a + /// completed event once the patch succeeds or fails. + FileChange(FileChangeItem), + /// Represents a call to an MCP tool. The item starts when the invocation is + /// dispatched and completes when the MCP server reports success or failure. + McpToolCall(McpToolCallItem), + /// Represents a call to a collab tool. The item starts when the collab tool is + /// invoked and completes when the collab tool reports success or failure. + CollabToolCall(CollabToolCallItem), + /// Captures a web search request. It starts when the search is kicked off + /// and completes when results are returned to the agent. + WebSearch(WebSearchItem), + /// Tracks the agent's running to-do list. It starts when the plan is first + /// issued, updates as steps change state, and completes when the turn ends. + TodoList(TodoListItem), + /// Describes a non-fatal error surfaced as an item. + Error(ErrorItem), +} + +/// Response from the agent. +/// Either a natural-language response or a JSON string when structured output is requested. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct AgentMessageItem { + pub text: String, +} + +/// Agent's reasoning summary. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct ReasoningItem { + pub text: String, +} + +/// The status of a command execution. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default, TS)] +#[serde(rename_all = "snake_case")] +pub enum CommandExecutionStatus { + #[default] + InProgress, + Completed, + Failed, + Declined, +} + +/// A command executed by the agent. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct CommandExecutionItem { + pub command: String, + pub aggregated_output: String, + pub exit_code: Option, + pub status: CommandExecutionStatus, +} + +/// A set of file changes by the agent. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct FileUpdateChange { + pub path: String, + pub kind: PatchChangeKind, +} + +/// The status of a file change. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +#[serde(rename_all = "snake_case")] +pub enum PatchApplyStatus { + InProgress, + Completed, + Failed, +} + +/// A set of file changes by the agent. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct FileChangeItem { + pub changes: Vec, + pub status: PatchApplyStatus, +} + +/// Indicates the type of the file change. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +#[serde(rename_all = "snake_case")] +pub enum PatchChangeKind { + Add, + Delete, + Update, +} + +/// The status of an MCP tool call. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default, TS)] +#[serde(rename_all = "snake_case")] +pub enum McpToolCallStatus { + #[default] + InProgress, + Completed, + Failed, +} + +/// The status of a collab tool call. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default, TS)] +#[serde(rename_all = "snake_case")] +pub enum CollabToolCallStatus { + #[default] + InProgress, + Completed, + Failed, +} + +/// Supported collab tools. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, TS)] +#[serde(rename_all = "snake_case")] +pub enum CollabTool { + SpawnAgent, + SendInput, + Wait, + CloseAgent, +} + +/// The status of a collab agent. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, TS)] +#[serde(rename_all = "snake_case")] +pub enum CollabAgentStatus { + PendingInit, + Running, + Interrupted, + Completed, + Errored, + Shutdown, + NotFound, +} + +/// Last known state of a collab agent. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, TS)] +pub struct CollabAgentState { + pub status: CollabAgentStatus, + pub message: Option, +} + +/// A call to a collab tool. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct CollabToolCallItem { + pub tool: CollabTool, + pub sender_thread_id: String, + pub receiver_thread_ids: Vec, + pub prompt: Option, + pub agents_states: HashMap, + pub status: CollabToolCallStatus, +} + +/// Result payload produced by an MCP tool invocation. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct McpToolCallItemResult { + // NOTE: `rmcp::model::Content` (and its `RawContent` variants) would be a + // more precise Rust representation of MCP content blocks. We intentionally + // use `serde_json::Value` here because this crate exports JSON schema + TS + // types (`schemars`/`ts-rs`), and the rmcp model types aren't set up to be + // schema/TS friendly (and would introduce heavier coupling to rmcp's Rust + // representations). Using `JsonValue` keeps the payload wire-shaped and + // easy to export. + pub content: Vec, + #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub meta: Option, + pub structured_content: Option, +} + +/// Error details reported by a failed MCP tool invocation. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct McpToolCallItemError { + pub message: String, +} + +/// A call to an MCP tool. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct McpToolCallItem { + pub server: String, + pub tool: String, + #[serde(default)] + pub arguments: JsonValue, + pub result: Option, + pub error: Option, + pub status: McpToolCallStatus, +} + +/// A web search request. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct WebSearchItem { + pub id: String, + pub query: String, + pub action: WebSearchAction, +} + +/// An error notification. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct ErrorItem { + pub message: String, +} + +/// An item in agent's to-do list. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct TodoItem { + pub text: String, + pub completed: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] +pub struct TodoListItem { + pub items: Vec, +} diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..b56a56f1ebea4e91b300913af17d6812db35e9f4 --- /dev/null +++ b/codex-rs/exec/src/lib.rs @@ -0,0 +1,2337 @@ +// - In the default output mode, it is paramount that the only thing written to +// stdout is the final message (if any). +// - In --json mode, stdout must be valid JSONL, one event per line. +// For both modes, any other output must be written to stderr. +#![recursion_limit = "256"] +#![deny(clippy::print_stdout)] + +mod cli; +mod event_processor; +mod event_processor_with_human_output; +pub(crate) mod event_processor_with_jsonl_output; +pub(crate) mod exec_events; +mod worktree; + +pub use cli::Cli; +pub use cli::Command; +pub use cli::ReviewArgs; +use codex_app_server_client::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY; +use codex_app_server_client::EnvironmentManager; +use codex_app_server_client::ExecServerRuntimePaths; +use codex_app_server_client::InProcessAppServerClient; +use codex_app_server_client::InProcessClientStartArgs; +use codex_app_server_client::InProcessServerEvent; +use codex_app_server_client::TypedRequestError; +use codex_app_server_protocol::ClientRequest; +use codex_app_server_protocol::ConfigWarningNotification; +use codex_app_server_protocol::JSONRPCErrorError; +use codex_app_server_protocol::McpServerElicitationAction; +use codex_app_server_protocol::McpServerElicitationRequestResponse; +use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ReviewStartParams; +use codex_app_server_protocol::ReviewStartResponse; +use codex_app_server_protocol::ReviewTarget as ApiReviewTarget; +use codex_app_server_protocol::ServerNotification; +use codex_app_server_protocol::ServerRequest; +use codex_app_server_protocol::Thread as AppServerThread; +use codex_app_server_protocol::ThreadForkParams; +use codex_app_server_protocol::ThreadForkResponse; +use codex_app_server_protocol::ThreadHistoryMode; +use codex_app_server_protocol::ThreadItem as AppServerThreadItem; +use codex_app_server_protocol::ThreadListParams; +use codex_app_server_protocol::ThreadListResponse; +use codex_app_server_protocol::ThreadReadParams; +use codex_app_server_protocol::ThreadReadResponse; +use codex_app_server_protocol::ThreadResumeParams; +use codex_app_server_protocol::ThreadResumeResponse; +use codex_app_server_protocol::ThreadSortKey; +use codex_app_server_protocol::ThreadSource; +use codex_app_server_protocol::ThreadSourceKind; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::ThreadStartResponse; +use codex_app_server_protocol::ThreadUnsubscribeParams; +use codex_app_server_protocol::ThreadUnsubscribeResponse; +use codex_app_server_protocol::TurnInterruptParams; +use codex_app_server_protocol::TurnInterruptResponse; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStartResponse; +use codex_app_server_protocol::TurnStartedNotification; +use codex_arg0::Arg0DispatchPaths; +use codex_cloud_config::cloud_config_bundle_loader_for_storage; +use codex_config::CloudConfigBundleLoader; +use codex_config::ConfigLoadError; +use codex_config::ConfigLoadOptions; +use codex_config::LoaderOverrides; +use codex_config::format_config_error_with_source; +use codex_core::StateDbHandle; +use codex_core::check_execpolicy_for_warnings; +use codex_core::config::Config; +use codex_core::config::ConfigBuilder; +use codex_core::config::ConfigOverrides; +use codex_core::config::ConfigTomlLoadResult; +use codex_core::config::bootstrap_auth_config; +use codex_core::config::find_codex_home; +use codex_core::config::load_config_toml_with_layer_stack; +use codex_core::config::resolve_oss_provider; +use codex_core::config::resolve_profile_v2_config_path; +use codex_core::find_thread_meta_by_name_str; +use codex_core::format_exec_policy_error_with_source; +use codex_core::path_utils; +use codex_core::read_session_meta_line; +use codex_features::Feature; +use codex_feedback::CodexFeedback; +use codex_git_utils::get_git_repo_root; +use codex_history::RolloutItem; +use codex_history::RolloutLine; +use codex_login::default_client::set_default_client_residency_requirement; +use codex_login::default_client::set_default_originator; +use codex_login::enforce_login_restrictions; +use codex_login::is_workload_identity_selected; +use codex_model_provider_info::LMSTUDIO_OSS_PROVIDER_ID; +use codex_model_provider_info::OLLAMA_OSS_PROVIDER_ID; +use codex_otel::set_parent_from_context; +use codex_otel::traceparent_context_from_env; +use codex_protocol::SessionId; +use codex_protocol::ThreadId; +use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::config_types::SandboxMode; +use codex_protocol::models::ActivePermissionProfile; +use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::ReviewRequest; +use codex_protocol::protocol::ReviewTarget; +use codex_protocol::protocol::SessionConfiguredEvent; +use codex_protocol::protocol::SessionSource; +use codex_protocol::user_input::UserInput; +use codex_utils_absolute_path::AbsolutePathBuf; +use codex_utils_absolute_path::canonicalize_existing_preserving_symlinks; +use codex_utils_cli::SharedCliOptions; +use codex_utils_oss::ensure_oss_provider_ready; +use codex_utils_oss::get_default_model_for_oss_provider; +use codex_worktree::CreateWorktree; +use codex_worktree::WorktreeManager; +use codex_worktree::WorktreeSettings; +use event_processor_with_human_output::EventProcessorWithHumanOutput; +pub use event_processor_with_jsonl_output::CodexStatus; +pub use event_processor_with_jsonl_output::CollectedThreadEvents; +pub use event_processor_with_jsonl_output::EventProcessorWithJsonOutput; +pub use exec_events::AgentMessageItem; +pub use exec_events::CollabAgentState; +pub use exec_events::CollabAgentStatus; +pub use exec_events::CollabTool; +pub use exec_events::CollabToolCallItem; +pub use exec_events::CollabToolCallStatus; +pub use exec_events::CommandExecutionItem; +pub use exec_events::CommandExecutionStatus; +pub use exec_events::ErrorItem; +pub use exec_events::FileChangeItem; +pub use exec_events::FileUpdateChange; +pub use exec_events::ItemCompletedEvent; +pub use exec_events::ItemStartedEvent; +pub use exec_events::ItemUpdatedEvent; +pub use exec_events::McpToolCallItem; +pub use exec_events::McpToolCallItemError; +pub use exec_events::McpToolCallItemResult; +pub use exec_events::McpToolCallStatus; +pub use exec_events::PatchApplyStatus; +pub use exec_events::PatchChangeKind; +pub use exec_events::ReasoningItem; +pub use exec_events::ThreadErrorEvent; +pub use exec_events::ThreadEvent; +pub use exec_events::ThreadItem as ExecThreadItem; +pub use exec_events::ThreadItemDetails; +pub use exec_events::ThreadStartedEvent; +pub use exec_events::TodoItem; +pub use exec_events::TodoListItem; +pub use exec_events::TurnCompletedEvent; +pub use exec_events::TurnFailedEvent; +pub use exec_events::TurnStartedEvent; +pub use exec_events::Usage; +pub use exec_events::WebSearchItem; +use serde_json::Value; +use std::collections::HashMap; +use std::future::Future; +use std::io::IsTerminal; +use std::io::Read; +use std::path::Path; +use std::path::PathBuf; +use supports_color::Stream; +use tokio::sync::mpsc; +use tracing::Instrument; +use tracing::error; +use tracing::field; +use tracing::info; +use tracing::info_span; +use tracing::warn; +use tracing_subscriber::EnvFilter; +use tracing_subscriber::prelude::*; +use uuid::Uuid; + +use crate::cli::Command as ExecCommand; +use crate::event_processor::EventProcessor; + +const DEFAULT_ANALYTICS_ENABLED: bool = true; +const EXEC_DEFAULT_LOG_FILTER: &str = "error,opentelemetry_sdk=off,opentelemetry_otlp=off"; + +enum InitialOperation { + ForkOnly, + UserTurn { + items: Vec, + output_schema: Option, + }, + Review { + review_request: ReviewRequest, + }, +} + +enum StdinPromptBehavior { + /// Read stdin only when there is no positional prompt, which is the legacy + /// `codex exec` behavior for `codex exec` with piped input. + RequiredIfPiped, + /// Always treat stdin as the prompt, used for the explicit `codex exec -` + /// sentinel and similar forced-stdin call sites. + Forced, + /// If stdin is piped alongside a positional prompt, treat stdin as + /// additional context to append rather than as the primary prompt. + OptionalAppend, +} + +struct RequestIdSequencer { + next: i64, +} + +impl RequestIdSequencer { + fn new() -> Self { + Self { next: 1 } + } + + fn next(&mut self) -> RequestId { + let id = self.next; + self.next += 1; + RequestId::Integer(id) + } +} + +struct ExecRunArgs { + in_process_start_args: InProcessClientStartArgs, + state_db: Option, + command: Option, + config: Config, + resume_approvals_reviewer_override: Option, + dangerously_bypass_approvals_and_sandbox: bool, + exec_span: tracing::Span, + images: Vec, + json_mode: bool, + last_message_file: Option, + model_provider: Option, + managed_worktree: Option, + oss: bool, + output_schema_path: Option, + prompt: Option, + skip_git_repo_check: bool, + stderr_with_ansi: bool, + thread_source: ThreadSource, +} + +struct ManagedExecWorktree { + manager: WorktreeManager, + checkout: PathBuf, + source_cwd: PathBuf, +} + +fn exec_root_span() -> tracing::Span { + info_span!( + "codex.exec", + otel.kind = "internal", + thread.id = field::Empty, + turn.id = field::Empty, + ) +} + +fn exec_stderr_env_filter() -> EnvFilter { + // OTEL export is best-effort; keep exporter self-diagnostics out of + // headless command output unless the caller opts in with RUST_LOG. + EnvFilter::try_from_default_env() + .or_else(|_| EnvFilter::try_new(EXEC_DEFAULT_LOG_FILTER)) + .unwrap_or_else(|_| EnvFilter::new("error")) +} + +pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { + if let Err(err) = set_default_originator("codex_exec".to_string()) { + tracing::warn!(?err, "Failed to set codex exec originator override {err:?}"); + } + + let Cli { + mut command, + strict_config, + shared, + thread_source, + skip_git_repo_check, + ephemeral, + ignore_user_config, + ignore_rules, + color, + last_message_file, + json: json_mode, + prompt, + output_schema: output_schema_path, + mut config_overrides, + } = cli; + let mut shared = shared.into_inner(); + shared.take_auto_review_config_overrides(&mut config_overrides); + let SharedCliOptions { + images, + model: model_cli_arg, + oss, + oss_provider, + config_profile_v2, + sandbox_mode: sandbox_mode_cli_arg, + auto_review: _, + dangerously_bypass_approvals_and_sandbox, + bypass_hook_trust, + cwd, + mut add_dir, + worktree, + } = shared; + + if worktree { + if ignore_user_config { + anyhow::bail!("--worktree cannot be combined with --ignore-user-config"); + } + if ephemeral { + anyhow::bail!("--worktree cannot be combined with --ephemeral"); + } + match command.as_ref() { + Some(ExecCommand::Resume(_)) => { + anyhow::bail!("--worktree is not supported with `codex exec resume`"); + } + Some(ExecCommand::Review(_)) => { + anyhow::bail!("--worktree is not supported with `codex exec review`"); + } + Some(ExecCommand::Fork(_)) | None => {} + } + } + + let (_stdout_with_ansi, stderr_with_ansi) = match color { + cli::Color::Always => (true, true), + cli::Color::Never => (false, false), + cli::Color::Auto => ( + supports_color::on_cached(Stream::Stdout).is_some(), + supports_color::on_cached(Stream::Stderr).is_some(), + ), + }; + let fmt_layer = tracing_subscriber::fmt::layer() + .with_ansi(stderr_with_ansi) + .with_writer(std::io::stderr) + .with_filter(exec_stderr_env_filter()); + + let sandbox_mode = if dangerously_bypass_approvals_and_sandbox { + Some(SandboxMode::DangerFullAccess) + } else { + sandbox_mode_cli_arg.map(Into::::into) + }; + + // Parse `-c` overrides from the CLI. + let cli_kv_overrides = match config_overrides.parse_overrides() { + Ok(v) => v, + #[allow(clippy::print_stderr)] + Err(e) => { + eprintln!("Error parsing -c overrides: {e}"); + std::process::exit(1); + } + }; + + let mut resolved_cwd = cwd.clone(); + let mut config_cwd = match resolved_cwd.as_deref() { + Some(path) => { + AbsolutePathBuf::from_absolute_path(canonicalize_existing_preserving_symlinks(path)?)? + } + None => AbsolutePathBuf::current_dir()?, + }; + + // we load config.toml here to determine project state. + #[allow(clippy::print_stderr)] + let codex_home = match find_codex_home() { + Ok(codex_home) => codex_home, + Err(err) => { + eprintln!("Error finding codex home: {err}"); + std::process::exit(1); + } + }; + let user_config_path = config_profile_v2 + .as_ref() + .map(|profile_v2| resolve_profile_v2_config_path(&codex_home, profile_v2)); + let loader_overrides = LoaderOverrides { + user_config_path, + user_config_profile: config_profile_v2, + ignore_user_config, + ignore_user_and_project_exec_policy_rules: ignore_rules, + ..Default::default() + }; + + if worktree + && EnvironmentManager::prepare_from_codex_home(&codex_home) + .await? + .default_environment_is_remote() + { + anyhow::bail!("--worktree requires local execution"); + } + + let managed_worktree = if worktree { + let gate_bootstrap = load_bootstrap_config_or_exit( + &codex_home, + /*cwd*/ None, + cli_kv_overrides.clone(), + loader_overrides.clone(), + strict_config, + CloudConfigBundleLoader::default(), + ) + .await; + let gate_cloud_config = cloud_config_bundle_loader_for_storage( + bootstrap_auth_config(&codex_home, &gate_bootstrap)?, + /*enable_codex_api_key_env*/ false, + ) + .await?; + let gate_config = ConfigBuilder::default() + .codex_home(codex_home.to_path_buf()) + .cli_overrides(cli_kv_overrides.clone()) + .loader_overrides(LoaderOverrides { + ignore_project_config: true, + ..loader_overrides.clone() + }) + .fallback_cwd(Some(config_cwd.to_path_buf())) + .strict_config(strict_config) + .cloud_config_bundle(gate_cloud_config.clone()) + .build() + .await?; + if !gate_config.features.enabled(Feature::Worktrees) { + anyhow::bail!( + "--worktree requires the worktrees feature; enable it with --enable worktrees" + ); + } + if let Some(ExecCommand::Fork(args)) = command.as_mut() { + let saved_cwd = worktree::fork_source( + args, + &gate_config, + &arg0_paths, + &cli_kv_overrides, + &loader_overrides, + gate_cloud_config.clone(), + strict_config, + ) + .await?; + if resolved_cwd.is_none() { + config_cwd = AbsolutePathBuf::from_absolute_path(saved_cwd)?; + } + } + let source_config = ConfigBuilder::default() + .codex_home(codex_home.to_path_buf()) + .cli_overrides(cli_kv_overrides.clone()) + .loader_overrides(LoaderOverrides { + ignore_project_config: true, + ..loader_overrides.clone() + }) + .fallback_cwd(Some(config_cwd.to_path_buf())) + .strict_config(strict_config) + .cloud_config_bundle(gate_cloud_config) + .build() + .await?; + if source_config.active_project.is_untrusted() { + anyhow::bail!("--worktree requires a source that is not explicitly untrusted"); + } + for path in &mut add_dir { + if path.is_relative() { + *path = std::env::current_dir()?.join(&*path); + } + } + // Allocation belongs to the host, not this session's project, profile, or overrides. + let host_config = load_bootstrap_config_or_exit( + &codex_home, + /*cwd*/ None, + Vec::new(), + LoaderOverrides::default(), + strict_config, + CloudConfigBundleLoader::default(), + ) + .await; + let settings = + WorktreeSettings::for_cli(&codex_home, host_config.config_toml.desktop.as_ref())?; + let manager = WorktreeManager::new(settings); + let checkout = manager.create(&CreateWorktree { + source_cwd: config_cwd.as_path().to_path_buf(), + base: None, + })?; + resolved_cwd = Some(checkout.cwd.clone()); + config_cwd = AbsolutePathBuf::from_absolute_path(checkout.cwd.clone())?; + Some(ManagedExecWorktree { + manager, + checkout: checkout.root, + source_cwd: checkout.source_cwd, + }) + } else { + None + }; + let bootstrap_config = load_bootstrap_config_or_exit( + &codex_home, + Some(&config_cwd), + cli_kv_overrides.clone(), + loader_overrides.clone(), + strict_config, + CloudConfigBundleLoader::default(), + ) + .await; + let bootstrap_config_toml = &bootstrap_config.config_toml; + let bootstrap_auth_config = bootstrap_auth_config(&codex_home, &bootstrap_config)?; + // API keys cannot fetch workspace-managed configuration. Preserve the + // existing ChatGPT bootstrap identity even when model requests allow + // CODEX_API_KEY. + let cloud_config_bundle = cloud_config_bundle_loader_for_storage( + bootstrap_auth_config, + /*enable_codex_api_key_env*/ false, + ) + .await?; + if let Some(worktree) = managed_worktree.as_ref() { + // Destination auth can fetch source policy that the host bootstrap could not. + let source_config = ConfigBuilder::default() + .codex_home(codex_home.to_path_buf()) + .cli_overrides(cli_kv_overrides.clone()) + .loader_overrides(LoaderOverrides { + ignore_project_config: true, + ..loader_overrides.clone() + }) + .fallback_cwd(Some(worktree.source_cwd.clone())) + .strict_config(strict_config) + .cloud_config_bundle(cloud_config_bundle.clone()) + .build() + .await?; + if source_config.active_project.is_untrusted() { + anyhow::bail!( + "--worktree requires a source that is not explicitly untrusted; unused checkout at {} remains. Remove it manually with `git worktree remove` when safe", + worktree.checkout.display() + ); + } + } + let run_cli_overrides = cli_kv_overrides.clone(); + let run_loader_overrides = loader_overrides.clone(); + let run_cloud_config_bundle = cloud_config_bundle.clone(); + + let model_provider = if oss { + let bootstrap_config_with_cloud_config; + let config_toml_for_oss = if oss_provider.is_none() { + // The first load intentionally skips cloud config so we can read + // auth/base-url settings needed to fetch the bundle. If OSS mode + // needs a default provider from config, reload with the bundle. + bootstrap_config_with_cloud_config = load_bootstrap_config_or_exit( + &codex_home, + Some(&config_cwd), + cli_kv_overrides.clone(), + loader_overrides.clone(), + strict_config, + cloud_config_bundle.clone(), + ) + .await; + &bootstrap_config_with_cloud_config.config_toml + } else { + bootstrap_config_toml + }; + + let resolved = resolve_oss_provider(oss_provider.as_deref(), config_toml_for_oss); + + if let Some(provider) = resolved { + Some(provider) + } else { + return Err(anyhow::anyhow!( + "No default OSS provider configured. Use --local-provider=provider or set oss_provider to one of: {LMSTUDIO_OSS_PROVIDER_ID}, {OLLAMA_OSS_PROVIDER_ID} in config.toml" + )); + } + } else { + None // No OSS mode enabled + }; + + // When using `--oss`, let the bootstrapper pick the model based on selected provider + let model = if let Some(model) = model_cli_arg { + Some(model) + } else if oss { + model_provider + .as_ref() + .and_then(|provider_id| get_default_model_for_oss_provider(provider_id)) + .map(std::borrow::ToOwned::to_owned) + } else { + None // No model specified, will use the default. + }; + + let overrides = ConfigOverrides { + model, + review_model: None, + // Default to never ask for approvals in headless mode. Rebuild below if + // the fully resolved reviewer is AutoReview. + approval_policy: Some(AskForApproval::Never), + approvals_reviewer: None, + sandbox_mode, + permission_profile: None, + default_permissions: None, + persisted_permission_profile_id: None, + cwd: resolved_cwd, + workspace_roots: None, + model_provider: model_provider.clone(), + service_tier: None, + codex_self_exe: arg0_paths.codex_self_exe.clone(), + codex_linux_sandbox_exe: arg0_paths.codex_linux_sandbox_exe.clone(), + main_execve_wrapper_exe: arg0_paths.main_execve_wrapper_exe.clone(), + default_zsh_path: None, + base_instructions: None, + developer_instructions: None, + personality: None, + compact_prompt: None, + show_raw_agent_reasoning: oss.then_some(true), + tools_web_search_request: None, + ephemeral: ephemeral.then_some(true), + bypass_hook_trust: bypass_hook_trust.then_some(true), + additional_writable_roots: add_dir, + }; + + let build_config = |overrides| { + ConfigBuilder::default() + .codex_home(codex_home.to_path_buf()) + .cli_overrides(cli_kv_overrides.clone()) + .harness_overrides(overrides) + .loader_overrides(loader_overrides.clone()) + .strict_config(strict_config) + .cloud_config_bundle(cloud_config_bundle.clone()) + .build() + }; + let config = build_exec_config( + overrides, + dangerously_bypass_approvals_and_sandbox, + build_config, + ) + .await?; + let resume_approvals_reviewer_override = cli_kv_overrides + .iter() + .any(|(key, _)| key == "approvals_reviewer") + .then(|| config.approvals_reviewer.into()); + + #[allow(clippy::print_stderr)] + match check_execpolicy_for_warnings(&config.config_layer_stack).await { + Ok(None) => {} + Ok(Some(err)) | Err(err) => { + eprintln!( + "Error loading rules:\n{}", + format_exec_policy_error_with_source(&err) + ); + std::process::exit(1); + } + } + + set_default_client_residency_requirement(config.enforce_residency.value()); + + if !is_workload_identity_selected() + && let Err(err) = enforce_login_restrictions(&config.auth_config()).await + { + eprintln!("{err}"); + std::process::exit(1); + } + + let otel = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + codex_core::otel_init::build_provider( + &config, + env!("CARGO_PKG_VERSION"), + /*service_name_override*/ None, + DEFAULT_ANALYTICS_ENABLED, + ) + })) { + Ok(Ok(otel)) => otel, + Ok(Err(e)) => { + eprintln!("Could not create otel exporter: {e}"); + None + } + Err(_) => { + eprintln!("Could not create otel exporter: panicked during initialization"); + None + } + }; + codex_core::otel_init::record_process_start(otel.as_ref(), "codex_exec"); + codex_core::otel_init::install_sqlite_telemetry(otel.as_ref(), "codex_exec"); + + let otel_logger_layer = otel.as_ref().and_then(|o| o.logger_layer()); + + let otel_tracing_layer = otel.as_ref().and_then(|o| o.tracing_layer()); + + let _ = tracing_subscriber::registry() + .with(fmt_layer) + .with(otel_tracing_layer) + .with(otel_logger_layer) + .try_init(); + + let exec_span = exec_root_span(); + if let Some(context) = traceparent_context_from_env() { + set_parent_from_context(&exec_span, context); + } + let config_warnings: Vec = config + .startup_warnings + .iter() + .map(|warning| ConfigWarningNotification { + summary: warning.clone(), + details: None, + path: None, + range: None, + }) + .collect(); + let local_runtime_paths = ExecServerRuntimePaths::from_optional_paths( + arg0_paths.codex_self_exe.clone(), + arg0_paths.codex_linux_sandbox_exe.clone(), + )?; + #[cfg(target_os = "macos")] + let local_runtime_paths = local_runtime_paths.with_allowed_symlinked_codex_home( + codex_config::allowed_symlinked_codex_home(&config.config_layer_stack, &config.codex_home), + ); + let state_db = codex_core::init_state_db(&config).await; + let environment_manager = if run_loader_overrides.ignore_user_config { + EnvironmentManager::from_env(Some(local_runtime_paths), config.http_client_factory()) + .await? + } else { + EnvironmentManager::from_codex_home( + config.codex_home.clone(), + Some(local_runtime_paths), + config.http_client_factory(), + ) + .await? + }; + let in_process_start_args = InProcessClientStartArgs { + arg0_paths, + config: std::sync::Arc::new(config.clone()), + cli_overrides: run_cli_overrides, + loader_overrides: run_loader_overrides, + strict_config, + cloud_config_bundle: run_cloud_config_bundle, + feedback: CodexFeedback::new(), + log_db: None, + state_db: state_db.clone(), + environment_manager: std::sync::Arc::new(environment_manager), + config_warnings, + session_source: SessionSource::Exec, + enable_codex_api_key_env: true, + client_name: "codex_exec".to_string(), + client_version: env!("CARGO_PKG_VERSION").to_string(), + experimental_api: true, + mcp_server_openai_form_elicitation: false, + opt_out_notification_methods: Vec::new(), + channel_capacity: DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, + }; + run_exec_session(ExecRunArgs { + in_process_start_args, + state_db, + command, + config, + resume_approvals_reviewer_override, + dangerously_bypass_approvals_and_sandbox, + exec_span: exec_span.clone(), + images, + json_mode, + last_message_file, + model_provider, + managed_worktree, + oss, + output_schema_path, + prompt, + skip_git_repo_check, + stderr_with_ansi, + thread_source: thread_source.map(Into::into).unwrap_or(ThreadSource::User), + }) + .instrument(exec_span) + .await +} + +async fn build_exec_config( + overrides: ConfigOverrides, + preserve_headless_approval_policy: bool, + build_config: BuildConfig, +) -> std::io::Result +where + BuildConfig: Fn(ConfigOverrides) -> BuildFuture, + BuildFuture: Future>, +{ + let build_without_headless_approval_policy = || { + build_config(ConfigOverrides { + approval_policy: None, + ..overrides.clone() + }) + }; + match build_config(overrides.clone()).await { + Ok(config) + if config.approvals_reviewer == ApprovalsReviewer::AutoReview + && !preserve_headless_approval_policy => + { + build_without_headless_approval_policy().await + } + Ok(config) => Ok(config), + Err(headless_error) if !preserve_headless_approval_policy => { + match build_without_headless_approval_policy().await { + Ok(config) if config.approvals_reviewer == ApprovalsReviewer::AutoReview => { + Ok(config) + } + Ok(_) | Err(_) => Err(headless_error), + } + } + Err(headless_error) => Err(headless_error), + } +} + +#[allow(clippy::print_stderr)] +async fn load_bootstrap_config_or_exit( + codex_home: &Path, + cwd: Option<&AbsolutePathBuf>, + cli_kv_overrides: Vec<(String, codex_config::TomlValue)>, + loader_overrides: LoaderOverrides, + strict_config: bool, + cloud_config_bundle: CloudConfigBundleLoader, +) -> ConfigTomlLoadResult { + match load_config_toml_with_layer_stack( + codex_home, + cwd, + cli_kv_overrides, + ConfigLoadOptions { + loader_overrides, + strict_config, + cloud_config_bundle, + }, + ) + .await + { + Ok(config_toml) => config_toml, + Err(err) => { + let config_error = err + .get_ref() + .and_then(|err| err.downcast_ref::()) + .map(ConfigLoadError::config_error); + if let Some(config_error) = config_error { + eprintln!( + "Error loading config.toml:\n{}", + format_config_error_with_source(config_error) + ); + } else { + eprintln!("Error loading config.toml: {err}"); + } + std::process::exit(1); + } + } +} + +async fn run_exec_session(args: ExecRunArgs) -> anyhow::Result<()> { + let ExecRunArgs { + in_process_start_args, + state_db, + command, + config, + resume_approvals_reviewer_override, + dangerously_bypass_approvals_and_sandbox, + exec_span, + images, + json_mode, + last_message_file, + model_provider, + managed_worktree, + oss, + output_schema_path, + prompt, + skip_git_repo_check, + stderr_with_ansi, + thread_source, + } = args; + + let mut event_processor: Box = match json_mode { + true => Box::new(EventProcessorWithJsonOutput::new(last_message_file.clone())), + _ => Box::new(EventProcessorWithHumanOutput::create_with_ansi( + stderr_with_ansi, + &config, + last_message_file.clone(), + )), + }; + if oss { + // We're in the oss section, so provider_id should be Some + // Let's handle None case gracefully though just in case + let provider_id = match model_provider.as_ref() { + Some(id) => id, + None => { + error!("OSS provider unexpectedly not set when oss flag is used"); + return Err(anyhow::anyhow!( + "OSS provider not set but oss flag was used" + )); + } + }; + ensure_oss_provider_ready(provider_id, &config) + .await + .map_err(|e| anyhow::anyhow!("OSS setup failed: {e}"))?; + } + + let default_cwd = config.cwd.to_path_buf(); + let default_approval_policy = config.permissions.approval_policy.value(); + let default_effort = config.model_reasoning_effort.clone(); + + let (initial_operation, prompt_summary) = match (command.as_ref(), prompt, images) { + (Some(ExecCommand::Review(review_cli)), _, _) => { + let review_request = build_review_request(review_cli)?; + let summary = codex_core::review_prompts::user_facing_hint(&review_request.target); + (InitialOperation::Review { review_request }, summary) + } + (Some(ExecCommand::Resume(args)), root_prompt, imgs) => { + let prompt_arg = args + .prompt + .clone() + .or_else(|| { + if args.last { + args.session_id.clone() + } else { + None + } + }) + .or(root_prompt); + let prompt_text = resolve_prompt(prompt_arg); + let mut items: Vec = imgs + .into_iter() + .chain(args.images.iter().cloned()) + .map(|path| UserInput::LocalImage { path, detail: None }) + .collect(); + items.push(UserInput::Text { + text: prompt_text.clone(), + // CLI input doesn't track UI element ranges, so none are available here. + text_elements: Vec::new(), + }); + let output_schema = load_output_schema(output_schema_path.clone()); + ( + InitialOperation::UserTurn { + items, + output_schema, + }, + prompt_text, + ) + } + (Some(ExecCommand::Fork(args)), root_prompt, imgs) => { + let prompt_arg = args.prompt.clone().or(root_prompt); + if let Some(prompt_arg) = prompt_arg { + let prompt_text = resolve_prompt(Some(prompt_arg)); + let mut items: Vec = imgs + .into_iter() + .chain(args.images.iter().cloned()) + .map(|path| UserInput::LocalImage { path, detail: None }) + .collect(); + items.push(UserInput::Text { + text: prompt_text.clone(), + text_elements: Vec::new(), + }); + let output_schema = load_output_schema(output_schema_path); + ( + InitialOperation::UserTurn { + items, + output_schema, + }, + prompt_text, + ) + } else if !imgs.is_empty() || !args.images.is_empty() { + anyhow::bail!("Forking with images requires a prompt"); + } else if output_schema_path.is_some() || last_message_file.is_some() { + anyhow::bail!("Forking with output options requires a prompt"); + } else if config.ephemeral { + anyhow::bail!("Ephemeral forks require a prompt"); + } else { + (InitialOperation::ForkOnly, String::new()) + } + } + (None, root_prompt, imgs) => { + let prompt_text = resolve_root_prompt(root_prompt); + let mut items: Vec = imgs + .into_iter() + .map(|path| UserInput::LocalImage { path, detail: None }) + .collect(); + items.push(UserInput::Text { + text: prompt_text.clone(), + // CLI input doesn't track UI element ranges, so none are available here. + text_elements: Vec::new(), + }); + let output_schema = load_output_schema(output_schema_path); + ( + InitialOperation::UserTurn { + items, + output_schema, + }, + prompt_text, + ) + } + }; + + // When --yolo (dangerously_bypass_approvals_and_sandbox) is set, also skip the git repo check + // since the user is explicitly running in an externally sandboxed environment. + if !skip_git_repo_check + && !dangerously_bypass_approvals_and_sandbox + && get_git_repo_root(&default_cwd).is_none() + { + eprintln!("Not inside a trusted directory and --skip-git-repo-check was not specified."); + std::process::exit(1); + } + + let mut request_ids = RequestIdSequencer::new(); + let mut client = InProcessAppServerClient::start(in_process_start_args) + .await + .map_err(|err| { + anyhow::anyhow!("failed to initialize in-process app-server client: {err}") + })?; + + // Resolve resume and fork through existing app-server thread lifecycle APIs. + let (primary_thread_id, fallback_session_configured) = if let Some(ExecCommand::Resume(args)) = + command.as_ref() + { + if let Some(thread_id) = + resolve_resume_thread_id(&client, &config, state_db.as_ref(), args).await? + { + let response: ThreadResumeResponse = send_request_with_response( + &client, + ClientRequest::ThreadResume { + request_id: request_ids.next(), + params: thread_resume_params_from_config( + &config, + thread_id, + resume_approvals_reviewer_override, + ), + }, + "thread/resume", + ) + .await + .map_err(anyhow::Error::msg)?; + let session_configured = + session_configured_from_thread_resume_response(&response, &config) + .map_err(anyhow::Error::msg)?; + (session_configured.thread_id, session_configured) + } else { + let response = start_thread(&client, &mut request_ids, &config, &thread_source) + .await + .map_err(anyhow::Error::msg)?; + let session_configured = + session_configured_from_thread_start_response(&response, &config) + .map_err(anyhow::Error::msg)?; + (session_configured.thread_id, session_configured) + } + } else if let Some(ExecCommand::Fork(args)) = command.as_ref() { + let source_args = crate::cli::ResumeArgs { + session_id: Some(args.session_id.clone()), + last: false, + all: true, + images: Vec::new(), + prompt: None, + }; + let source_thread_id = + resolve_resume_thread_id(&client, &config, state_db.as_ref(), &source_args) + .await? + .ok_or_else(|| anyhow::anyhow!("Session not found: {}", args.session_id))?; + let permissions = permissions_selection_from_config(&config); + let sandbox = permissions.is_none().then(|| { + sandbox_mode_from_permission_profile( + &config.permissions.effective_permission_profile(), + config.cwd.as_path(), + ) + }); + let response: ThreadForkResponse = send_request_with_response( + &client, + ClientRequest::ThreadFork { + request_id: request_ids.next(), + params: ThreadForkParams { + thread_id: source_thread_id, + model: config.model.clone(), + model_provider: Some(config.model_provider_id.clone()), + cwd: Some(config.cwd.to_string_lossy().to_string()), + runtime_workspace_roots: Some(config.workspace_roots.clone()), + approval_policy: Some(config.permissions.approval_policy.value().into()), + approvals_reviewer: resume_approvals_reviewer_override, + sandbox: sandbox.flatten(), + permissions, + config: thread_config_overrides_from_config(&config), + ephemeral: config.ephemeral, + thread_source: Some(thread_source.clone()), + exclude_turns: true, + defer_goal_continuation: !config.ephemeral, + ..ThreadForkParams::default() + }, + }, + "thread/fork", + ) + .await + .map_err(anyhow::Error::msg)?; + let session_configured = session_configured_from_thread_response( + &response.thread.session_id, + &response.thread.id, + response.thread.forked_from_id.as_deref(), + response.thread.parent_thread_id.as_deref(), + response.thread.thread_source.clone().map(Into::into), + response.thread.name.clone(), + response.thread.path.clone(), + response.model, + response.model_provider, + response.service_tier, + response.approval_policy.to_core(), + response.approvals_reviewer.to_core(), + config.permissions.effective_permission_profile(), + response.active_permission_profile.map(Into::into), + response.cwd, + response.reasoning_effort, + ) + .map_err(anyhow::Error::msg)?; + (session_configured.thread_id, session_configured) + } else { + let response = start_thread(&client, &mut request_ids, &config, &thread_source) + .await + .map_err(anyhow::Error::msg)?; + let session_configured = session_configured_from_thread_start_response(&response, &config) + .map_err(anyhow::Error::msg)?; + (session_configured.thread_id, session_configured) + }; + + if let Some(worktree) = managed_worktree.as_ref() { + worktree + .manager + .bind_thread(&worktree.checkout, &primary_thread_id.to_string())?; + } + + let primary_thread_id_for_span = primary_thread_id.to_string(); + // Use the start/resume response as the authoritative bootstrap payload. + // Waiting for a later streamed `SessionConfigured` event adds up to 10s of + // avoidable startup latency on the in-process path. + let session_configured = fallback_session_configured; + + exec_span.record("thread.id", primary_thread_id_for_span.as_str()); + + // Print the effective configuration and initial request so users can see what Codex + // is using. + event_processor.print_config_summary(&config, &prompt_summary, &session_configured); + if !json_mode + && let Some(message) = + codex_core::config::system_bwrap_warning(config.permissions.permission_profile()) + { + event_processor.process_warning(message); + } + + info!("Codex initialized with event: {session_configured:?}"); + + let (interrupt_tx, mut interrupt_rx) = mpsc::unbounded_channel::<()>(); + tokio::spawn(async move { + if tokio::signal::ctrl_c().await.is_ok() { + tracing::debug!("Keyboard interrupt"); + let _ = interrupt_tx.send(()); + } + }); + + let task_id = match initial_operation { + InitialOperation::ForkOnly => { + request_shutdown(&client, &mut request_ids, &primary_thread_id_for_span) + .await + .map_err(anyhow::Error::msg)?; + client + .shutdown() + .await + .map_err(|err| anyhow::anyhow!("in-process app-server shutdown failed: {err}"))?; + event_processor.print_final_output(); + return Ok(()); + } + InitialOperation::UserTurn { + items, + output_schema, + } => { + let response: TurnStartResponse = send_request_with_response( + &client, + ClientRequest::TurnStart { + request_id: request_ids.next(), + params: TurnStartParams { + disabled_plugin_ids: None, + thread_id: primary_thread_id_for_span.clone(), + turn_trigger: None, + client_user_message_id: None, + input: items.into_iter().map(Into::into).collect(), + tool_output: None, + responsesapi_client_metadata: None, + additional_context: None, + environments: None, + cwd: Some(default_cwd), + runtime_workspace_roots: None, + approval_policy: Some(default_approval_policy.into()), + approvals_reviewer: None, + sandbox_policy: None, + permissions: None, + model: None, + service_tier: None, + service_tier_for_turn: None, + effort: default_effort, + summary: None, + personality: None, + output_schema, + collaboration_mode: None, + multi_agent_mode: None, + cyber_access_program: None, + }, + }, + "turn/start", + ) + .await + .map_err(anyhow::Error::msg)?; + let task_id = response.turn.id; + info!("Sent prompt with event ID: {task_id}"); + task_id + } + InitialOperation::Review { review_request } => { + let response: ReviewStartResponse = send_request_with_response( + &client, + ClientRequest::ReviewStart { + request_id: request_ids.next(), + params: ReviewStartParams { + thread_id: primary_thread_id_for_span.clone(), + target: review_target_to_api(review_request.target), + delivery: None, + }, + }, + "review/start", + ) + .await + .map_err(anyhow::Error::msg)?; + let _ = event_processor.process_server_notification(ServerNotification::TurnStarted( + TurnStartedNotification { + thread_id: response.review_thread_id.clone(), + turn: response.turn.clone(), + }, + )); + let task_id = response.turn.id; + info!("Sent review request with event ID: {task_id}"); + task_id + } + }; + exec_span.record("turn.id", task_id.as_str()); + + // Run the loop until the task is complete. + // Track whether a fatal error was reported by the server so we can + // exit with a non-zero status for automation-friendly signaling. + let mut error_seen = false; + let mut interrupt_channel_open = true; + let primary_thread_id_for_requests = primary_thread_id.to_string(); + loop { + let server_event = tokio::select! { + maybe_interrupt = interrupt_rx.recv(), if interrupt_channel_open => { + if maybe_interrupt.is_none() { + interrupt_channel_open = false; + continue; + } + if let Err(err) = send_request_with_response::( + &client, + ClientRequest::TurnInterrupt { + request_id: request_ids.next(), + params: TurnInterruptParams { + thread_id: primary_thread_id_for_requests.clone(), + turn_id: task_id.clone(), + }, + }, + "turn/interrupt", + ) + .await + { + warn!("turn/interrupt failed: {err}"); + } + continue; + } + maybe_event = client.next_event() => maybe_event, + }; + + let Some(server_event) = server_event else { + break; + }; + + match server_event { + InProcessServerEvent::ServerRequest(request) => { + handle_server_request(&client, *request, &mut error_seen).await; + } + InProcessServerEvent::ServerNotification(notification) => { + let mut notification = *notification; + if let ServerNotification::Error(payload) = ¬ification { + if payload.thread_id == primary_thread_id_for_requests + && payload.turn_id == task_id + && !payload.will_retry + { + error_seen = true; + } + } else if let ServerNotification::TurnCompleted(payload) = ¬ification + && payload.thread_id == primary_thread_id_for_requests + && payload.turn.id == task_id + && matches!( + payload.turn.status, + codex_app_server_protocol::TurnStatus::Failed + | codex_app_server_protocol::TurnStatus::Interrupted + ) + { + error_seen = true; + } + + if should_process_notification( + ¬ification, + &primary_thread_id_for_requests, + &task_id, + ) { + maybe_backfill_turn_completed_items( + config.ephemeral, + &client, + &mut request_ids, + &mut notification, + ) + .await; + + match event_processor.process_server_notification(notification) { + CodexStatus::Running => {} + CodexStatus::InitiateShutdown => { + if let Err(err) = request_shutdown( + &client, + &mut request_ids, + &primary_thread_id_for_requests, + ) + .await + { + warn!("thread/unsubscribe failed during shutdown: {err}"); + } + break; + } + } + } + } + InProcessServerEvent::Lagged { skipped } => { + let message = lagged_event_warning_message(skipped); + warn!("{message}"); + event_processor.process_warning(message); + } + } + } + + if let Err(err) = client.shutdown().await { + warn!("in-process app-server shutdown failed: {err}"); + } + event_processor.print_final_output(); + if error_seen { + std::process::exit(1); + } + + Ok(()) +} + +async fn start_thread( + client: &InProcessAppServerClient, + request_ids: &mut RequestIdSequencer, + config: &Config, + thread_source: &ThreadSource, +) -> Result { + let mut params = thread_start_params_from_config(config, thread_source); + loop { + match client + .request_typed(ClientRequest::ThreadStart { + request_id: request_ids.next(), + params: params.clone(), + }) + .await + { + Ok(response) => return Ok(response), + Err(TypedRequestError::Server { source, .. }) + if params.history_mode.is_some() + && source.code == -32600 + && source.message + == "paginated threads require thread/turns/list and thread/items/list support" => + { + params.history_mode = None; + } + Err(err) => return Err(format!("thread/start: {err}")), + } + } +} + +fn thread_start_params_from_config( + config: &Config, + thread_source: &ThreadSource, +) -> ThreadStartParams { + let permissions = permissions_selection_from_config(config); + let sandbox = permissions.is_none().then(|| { + sandbox_mode_from_permission_profile( + &config.permissions.effective_permission_profile(), + config.cwd.as_path(), + ) + }); + ThreadStartParams { + model: config.model.clone(), + model_provider: Some(config.model_provider_id.clone()), + cwd: Some(config.cwd.to_string_lossy().to_string()), + runtime_workspace_roots: Some(config.workspace_roots.clone()), + approval_policy: Some(config.permissions.approval_policy.value().into()), + approvals_reviewer: Some(config.approvals_reviewer.into()), + sandbox: sandbox.flatten(), + permissions, + config: thread_config_overrides_from_config(config), + ephemeral: Some(config.ephemeral), + history_mode: (!config.ephemeral).then_some(ThreadHistoryMode::Paginated), + thread_source: Some(thread_source.clone()), + ..ThreadStartParams::default() + } +} + +fn thread_resume_params_from_config( + config: &Config, + thread_id: String, + approvals_reviewer_override: Option, +) -> ThreadResumeParams { + let permissions = permissions_selection_from_config(config); + let sandbox = permissions.is_none().then(|| { + sandbox_mode_from_permission_profile( + &config.permissions.effective_permission_profile(), + config.cwd.as_path(), + ) + }); + ThreadResumeParams { + thread_id, + model: config.model.clone(), + model_provider: Some(config.model_provider_id.clone()), + cwd: Some(config.cwd.to_string_lossy().to_string()), + runtime_workspace_roots: Some(config.workspace_roots.clone()), + approval_policy: Some(config.permissions.approval_policy.value().into()), + approvals_reviewer: approvals_reviewer_override, + sandbox: sandbox.flatten(), + permissions, + config: thread_config_overrides_from_config(config), + exclude_turns: true, + ..ThreadResumeParams::default() + } +} + +fn thread_config_overrides_from_config(config: &Config) -> Option> { + config + .bypass_hook_trust + .then(|| HashMap::from([("bypass_hook_trust".to_string(), Value::Bool(true))])) +} + +fn permissions_selection_from_config(config: &Config) -> Option { + config + .permissions + .active_permission_profile() + .map(permission_profile_id_from_active_profile) +} + +fn permission_profile_id_from_active_profile(active: ActivePermissionProfile) -> String { + active.id +} + +fn sandbox_mode_from_permission_profile( + permission_profile: &PermissionProfile, + cwd: &Path, +) -> Option { + match permission_profile { + PermissionProfile::Disabled => { + Some(codex_app_server_protocol::SandboxMode::DangerFullAccess) + } + PermissionProfile::External { .. } => None, + PermissionProfile::Managed { .. } => { + let file_system_policy = permission_profile.file_system_sandbox_policy(); + if file_system_policy.has_full_disk_write_access() { + permission_profile + .network_sandbox_policy() + .is_enabled() + .then_some(codex_app_server_protocol::SandboxMode::DangerFullAccess) + } else if file_system_policy.can_write_local_path_with_cwd(cwd, cwd) { + Some(codex_app_server_protocol::SandboxMode::WorkspaceWrite) + } else { + Some(codex_app_server_protocol::SandboxMode::ReadOnly) + } + } + } +} + +async fn send_request_with_response( + client: &InProcessAppServerClient, + request: ClientRequest, + method: &str, +) -> Result +where + T: serde::de::DeserializeOwned, +{ + client.request_typed(request).await.map_err(|err| { + if method.is_empty() { + err.to_string() + } else { + format!("{method}: {err}") + } + }) +} + +fn session_configured_from_thread_start_response( + response: &ThreadStartResponse, + config: &Config, +) -> Result { + session_configured_from_thread_response( + &response.thread.session_id, + &response.thread.id, + response.thread.forked_from_id.as_deref(), + response.thread.parent_thread_id.as_deref(), + response.thread.thread_source.clone().map(Into::into), + response.thread.name.clone(), + response.thread.path.clone(), + response.model.clone(), + response.model_provider.clone(), + response.service_tier.clone(), + response.approval_policy.to_core(), + response.approvals_reviewer.to_core(), + config.permissions.effective_permission_profile(), + response.active_permission_profile.clone().map(Into::into), + response.cwd.clone(), + response.reasoning_effort.clone(), + ) +} + +fn session_configured_from_thread_resume_response( + response: &ThreadResumeResponse, + config: &Config, +) -> Result { + session_configured_from_thread_response( + &response.thread.session_id, + &response.thread.id, + response.thread.forked_from_id.as_deref(), + response.thread.parent_thread_id.as_deref(), + response.thread.thread_source.clone().map(Into::into), + response.thread.name.clone(), + response.thread.path.clone(), + response.model.clone(), + response.model_provider.clone(), + response.service_tier.clone(), + response.approval_policy.to_core(), + response.approvals_reviewer.to_core(), + config.permissions.effective_permission_profile(), + response.active_permission_profile.clone().map(Into::into), + response.cwd.clone(), + response.reasoning_effort.clone(), + ) +} + +fn review_target_to_api(target: ReviewTarget) -> ApiReviewTarget { + match target { + ReviewTarget::UncommittedChanges => ApiReviewTarget::UncommittedChanges, + ReviewTarget::BaseBranch { branch } => ApiReviewTarget::BaseBranch { branch }, + ReviewTarget::Commit { sha, title } => ApiReviewTarget::Commit { sha, title }, + ReviewTarget::Custom { instructions } => ApiReviewTarget::Custom { instructions }, + } +} + +#[expect( + clippy::too_many_arguments, + reason = "session mapping keeps explicit fields" +)] +fn session_configured_from_thread_response( + session_id: &str, + thread_id: &str, + forked_from_id: Option<&str>, + parent_thread_id: Option<&str>, + thread_source: Option, + thread_name: Option, + rollout_path: Option, + model: String, + model_provider_id: String, + service_tier: Option, + approval_policy: AskForApproval, + approvals_reviewer: codex_protocol::config_types::ApprovalsReviewer, + permission_profile: PermissionProfile, + active_permission_profile: Option, + cwd: AbsolutePathBuf, + reasoning_effort: Option, +) -> Result { + let session_id = SessionId::from_string(session_id) + .map_err(|err| format!("session id `{session_id}` is invalid: {err}"))?; + let thread_id = ThreadId::from_string(thread_id) + .map_err(|err| format!("thread id `{thread_id}` is invalid: {err}"))?; + let forked_from_id = forked_from_id + .map(ThreadId::from_string) + .transpose() + .map_err(|err| format!("forked-from thread id is invalid: {err}"))?; + let parent_thread_id = parent_thread_id + .map(ThreadId::from_string) + .transpose() + .map_err(|err| format!("parent thread id is invalid: {err}"))?; + + Ok(SessionConfiguredEvent { + session_id, + thread_id, + forked_from_id, + parent_thread_id, + thread_source, + thread_name, + model, + model_provider_id, + service_tier, + approval_policy, + approvals_reviewer, + permission_profile, + active_permission_profile, + cwd, + reasoning_effort, + initial_messages: None, + network_proxy: None, + rollout_path, + }) +} + +fn lagged_event_warning_message(skipped: usize) -> String { + format!("in-process app-server event stream lagged; dropped {skipped} events") +} + +fn should_process_notification( + notification: &ServerNotification, + thread_id: &str, + turn_id: &str, +) -> bool { + match notification { + ServerNotification::ConfigWarning(_) | ServerNotification::DeprecationNotice(_) => true, + // TODO(anp) resolve duplicate startup warnings + ServerNotification::Warning(notification) => notification + .thread_id + .as_deref() + .is_none_or(|candidate| candidate == thread_id), + ServerNotification::AuthRecoveryStarted(notification) + | ServerNotification::AuthRecoveryCompleted(notification) => { + notification.thread_id == thread_id && notification.turn_id == turn_id + } + ServerNotification::Error(notification) => { + notification.thread_id == thread_id && notification.turn_id == turn_id + } + ServerNotification::HookCompleted(notification) => { + notification.thread_id == thread_id + && notification + .turn_id + .as_deref() + .is_none_or(|candidate| candidate == turn_id) + } + ServerNotification::HookStarted(notification) => { + notification.thread_id == thread_id + && notification + .turn_id + .as_deref() + .is_none_or(|candidate| candidate == turn_id) + } + ServerNotification::ItemCompleted(notification) => { + notification.thread_id == thread_id && notification.turn_id == turn_id + } + ServerNotification::ItemStarted(notification) => { + notification.thread_id == thread_id && notification.turn_id == turn_id + } + ServerNotification::ModelRerouted(notification) => { + notification.thread_id == thread_id && notification.turn_id == turn_id + } + ServerNotification::ModelVerification(notification) => { + notification.thread_id == thread_id && notification.turn_id == turn_id + } + ServerNotification::ThreadTokenUsageUpdated(notification) => { + notification.thread_id == thread_id && notification.turn_id == turn_id + } + ServerNotification::TurnCompleted(notification) => { + notification.thread_id == thread_id && notification.turn.id == turn_id + } + ServerNotification::TurnDiffUpdated(notification) => { + notification.thread_id == thread_id && notification.turn_id == turn_id + } + ServerNotification::TurnPlanUpdated(notification) => { + notification.thread_id == thread_id && notification.turn_id == turn_id + } + ServerNotification::TurnStarted(notification) => { + notification.thread_id == thread_id && notification.turn.id == turn_id + } + _ => false, + } +} + +async fn maybe_backfill_turn_completed_items( + thread_ephemeral: bool, + client: &InProcessAppServerClient, + request_ids: &mut RequestIdSequencer, + notification: &mut ServerNotification, +) { + // In-process delivery may drop non-terminal item notifications under backpressure while still + // guaranteeing `turn/completed`. Because app-server currently emits that completion with an + // empty `turn.items`, exec does one last `thread/read` here so human/json output can recover + // the final message and reconcile any still-running items before shutdown. + if !should_backfill_turn_completed_items(thread_ephemeral, notification) { + return; + } + + let ServerNotification::TurnCompleted(payload) = notification else { + return; + }; + + let response = send_request_with_response::( + client, + ClientRequest::ThreadRead { + request_id: request_ids.next(), + params: ThreadReadParams { + thread_id: payload.thread_id.clone(), + include_turns: true, + }, + }, + "thread/read", + ) + .await; + + match response { + Ok(response) => { + if let Some(items) = turn_items_for_thread(&response.thread, &payload.turn.id) { + payload.turn.items = items; + } + } + Err(err) => { + warn!("thread/read failed while backfilling turn items for turn completion: {err}"); + } + } +} + +/// Returns true only when `exec` can safely recover missing turn items from +/// rollout-backed thread history. +fn should_backfill_turn_completed_items( + thread_ephemeral: bool, + notification: &ServerNotification, +) -> bool { + let ServerNotification::TurnCompleted(payload) = notification else { + return false; + }; + + !thread_ephemeral && payload.turn.items_view != codex_app_server_protocol::TurnItemsView::Full +} + +fn turn_items_for_thread( + thread: &AppServerThread, + turn_id: &str, +) -> Option> { + thread + .turns + .iter() + .find(|turn| turn.id == turn_id) + .map(|turn| turn.items.clone()) +} + +fn all_thread_source_kinds() -> Vec { + vec![ + ThreadSourceKind::Cli, + ThreadSourceKind::VsCode, + ThreadSourceKind::Exec, + ThreadSourceKind::AppServer, + ThreadSourceKind::SubAgent, + ThreadSourceKind::SubAgentReview, + ThreadSourceKind::SubAgentCompact, + ThreadSourceKind::SubAgentThreadSpawn, + ThreadSourceKind::SubAgentOther, + ThreadSourceKind::Unknown, + ] +} + +async fn latest_thread_cwd(thread: &AppServerThread) -> PathBuf { + if let Some(path) = thread.path.as_deref() + && let Some(cwd) = parse_latest_turn_context_cwd(path).await + { + return cwd; + } + thread.cwd.to_path_buf() +} + +async fn parse_latest_turn_context_cwd(path: &Path) -> Option { + let path = path.to_path_buf(); + tokio::task::spawn_blocking(move || { + let reader = codex_rollout::open_rollout_seekable_reader(&path).ok()?; + let mut scanner = codex_rollout::ReverseJsonlScanner::new(reader).ok()?; + while let Some(outcome) = scanner.scan_next_rollout_line().ok()? { + if let codex_rollout::ScanOutcome::Parsed(RolloutLine { + item: RolloutItem::TurnContext(item), + .. + }) = outcome + { + return Some(item.cwd.into_path_buf()); + } + } + None + }) + .await + .ok() + .flatten() +} + +fn cwds_match(current_cwd: &Path, session_cwd: &Path) -> bool { + path_utils::paths_match_after_normalization(current_cwd, session_cwd) +} + +async fn resolve_resume_thread_id( + client: &InProcessAppServerClient, + config: &Config, + state_db: Option<&StateDbHandle>, + args: &crate::cli::ResumeArgs, +) -> anyhow::Result> { + let model_providers = resume_lookup_model_providers(config, args); + + if args.last { + let mut use_state_db_only = state_db.is_some(); + let mut cursor = None; + loop { + let response: ThreadListResponse = send_request_with_response( + client, + ClientRequest::ThreadList { + request_id: RequestId::Integer(0), + params: ThreadListParams { + originators: None, + cursor, + limit: Some(100), + sort_key: Some(ThreadSortKey::UpdatedAt), + sort_direction: None, + model_providers: model_providers.clone(), + source_kinds: Some(all_thread_source_kinds()), + archived: Some(false), + section_id: None, + project_id: None, + parent_thread_id: None, + ancestor_thread_id: None, + cwd: None, + use_state_db_only, + search_term: None, + }, + }, + "thread/list", + ) + .await + .map_err(anyhow::Error::msg)?; + for thread in response.data { + if use_state_db_only && let Some(path) = thread.path.as_deref() { + let Ok(session_meta) = read_session_meta_line(path).await else { + continue; + }; + if session_meta.meta.id.to_string() != thread.id { + continue; + } + } + let latest_cwd = latest_thread_cwd(&thread).await; + if args.all || cwds_match(config.cwd.as_path(), latest_cwd.as_path()) { + // A usable SQLite candidate is authoritative. Scanning is reserved for a + // complete miss so successful `--last` lookups avoid auditing every rollout. + return Ok(Some(thread.id)); + } + } + let Some(next_cursor) = response.next_cursor else { + if use_state_db_only { + // Repair from rollouts before giving up on a missing SQLite match. + use_state_db_only = false; + cursor = None; + continue; + } + return Ok(None); + }; + cursor = Some(next_cursor); + } + } + + let Some(session_id) = args.session_id.as_deref() else { + return Ok(None); + }; + if Uuid::parse_str(session_id).is_ok() { + return Ok(Some(session_id.to_string())); + } + if let Some(state_db) = state_db { + let cwd = (!args.all).then_some(config.cwd.as_path()); + let resolved = state_db + .find_thread_by_exact_title( + session_id, + &[], + /*model_providers*/ None, + /*archived_only*/ false, + cwd, + ) + .await?; + if let Some(thread) = resolved { + return Ok(Some(thread.id.to_string())); + } + if let Some((_, session_meta)) = + find_thread_meta_by_name_str(&config.codex_home, session_id, Some(state_db.as_ref())) + .await? + && (args.all || cwds_match(config.cwd.as_path(), &session_meta.meta.cwd)) + { + return Ok(Some(session_meta.meta.id.to_string())); + } + } + + let mut cursor = None; + loop { + let response: ThreadListResponse = send_request_with_response( + client, + ClientRequest::ThreadList { + request_id: RequestId::Integer(0), + params: ThreadListParams { + originators: None, + cursor, + limit: Some(100), + sort_key: Some(ThreadSortKey::UpdatedAt), + sort_direction: None, + model_providers: model_providers.clone(), + source_kinds: Some(all_thread_source_kinds()), + archived: Some(false), + section_id: None, + project_id: None, + parent_thread_id: None, + ancestor_thread_id: None, + cwd: None, + use_state_db_only: false, + search_term: Some(session_id.to_string()), + }, + }, + "thread/list", + ) + .await + .map_err(anyhow::Error::msg)?; + for thread in response.data { + if thread.name.as_deref() != Some(session_id) { + continue; + } + let latest_cwd = latest_thread_cwd(&thread).await; + if args.all || cwds_match(config.cwd.as_path(), latest_cwd.as_path()) { + return Ok(Some(thread.id)); + } + } + let Some(next_cursor) = response.next_cursor else { + return Ok(None); + }; + cursor = Some(next_cursor); + } +} + +fn resume_lookup_model_providers( + config: &Config, + args: &crate::cli::ResumeArgs, +) -> Option> { + if args.last { + Some(vec![config.model_provider_id.clone()]) + } else { + None + } +} + +fn canceled_mcp_server_elicitation_response() -> Result { + serde_json::to_value(McpServerElicitationRequestResponse { + action: McpServerElicitationAction::Cancel, + content: None, + meta: None, + }) + .map_err(|err| format!("failed to encode mcp elicitation response: {err}")) +} + +async fn request_shutdown( + client: &InProcessAppServerClient, + request_ids: &mut RequestIdSequencer, + thread_id: &str, +) -> Result<(), String> { + let request = ClientRequest::ThreadUnsubscribe { + request_id: request_ids.next(), + params: ThreadUnsubscribeParams { + thread_id: thread_id.to_string(), + }, + }; + send_request_with_response::(client, request, "thread/unsubscribe") + .await + .map(|_| ()) +} + +async fn resolve_server_request( + client: &InProcessAppServerClient, + request_id: RequestId, + value: serde_json::Value, + method: &str, +) -> Result<(), String> { + client + .resolve_server_request(request_id, value) + .await + .map_err(|err| format!("failed to resolve `{method}` server request: {err}")) +} + +async fn reject_server_request( + client: &InProcessAppServerClient, + request_id: RequestId, + method: &str, + reason: String, +) -> Result<(), String> { + client + .reject_server_request( + request_id, + JSONRPCErrorError { + code: -32000, + message: reason, + data: None, + }, + ) + .await + .map_err(|err| format!("failed to reject `{method}` server request: {err}")) +} + +fn server_request_method_name(request: &ServerRequest) -> String { + serde_json::to_value(request) + .ok() + .and_then(|value| { + value + .get("method") + .and_then(serde_json::Value::as_str) + .map(str::to_owned) + }) + .unwrap_or_else(|| "unknown".to_string()) +} + +async fn handle_server_request( + client: &InProcessAppServerClient, + request: ServerRequest, + error_seen: &mut bool, +) { + let method = server_request_method_name(&request); + let handle_result = match request { + ServerRequest::McpServerElicitationRequest { request_id, .. } => { + // Exec auto-cancels elicitation instead of surfacing it + // interactively. Preserve that behavior for attached subagent + // threads too so we do not turn a cancel into a decline/error. + match canceled_mcp_server_elicitation_response() { + Ok(value) => { + resolve_server_request( + client, + request_id, + value, + "mcpServer/elicitation/request", + ) + .await + } + Err(err) => Err(err), + } + } + ServerRequest::CommandExecutionRequestApproval { request_id, params } => { + reject_server_request( + client, + request_id, + &method, + format!( + "command execution approval is not supported in exec mode for thread `{}`", + params.thread_id + ), + ) + .await + } + ServerRequest::FileChangeRequestApproval { request_id, params } => { + reject_server_request( + client, + request_id, + &method, + format!( + "file change approval is not supported in exec mode for thread `{}`", + params.thread_id + ), + ) + .await + } + ServerRequest::ToolRequestUserInput { request_id, params } => { + reject_server_request( + client, + request_id, + &method, + format!( + "request_user_input is not supported in exec mode for thread `{}`", + params.thread_id + ), + ) + .await + } + ServerRequest::DynamicToolCall { request_id, params } => { + reject_server_request( + client, + request_id, + &method, + format!( + "dynamic tool calls are not supported in exec mode for thread `{}`", + params.thread_id + ), + ) + .await + } + ServerRequest::ChatgptAuthTokensRefresh { request_id, .. } => { + reject_server_request( + client, + request_id, + &method, + "chatgpt auth token refresh is not supported in exec mode".to_string(), + ) + .await + } + ServerRequest::AttestationGenerate { request_id, .. } => { + reject_server_request( + client, + request_id, + &method, + "attestation generation is not supported in exec mode".to_string(), + ) + .await + } + ServerRequest::CurrentTimeRead { request_id, .. } => { + reject_server_request( + client, + request_id, + &method, + "external current time is not supported in exec mode".to_string(), + ) + .await + } + ServerRequest::ApplyPatchApproval { request_id, params } => { + reject_server_request( + client, + request_id, + &method, + format!( + "apply_patch approval is not supported in exec mode for thread `{}`", + params.conversation_id + ), + ) + .await + } + ServerRequest::ExecCommandApproval { request_id, params } => { + reject_server_request( + client, + request_id, + &method, + format!( + "exec command approval is not supported in exec mode for thread `{}`", + params.conversation_id + ), + ) + .await + } + ServerRequest::PermissionsRequestApproval { request_id, params } => { + reject_server_request( + client, + request_id, + &method, + format!( + "permissions approval is not supported in exec mode for thread `{}`", + params.thread_id + ), + ) + .await + } + }; + + if let Err(err) = handle_result { + *error_seen = true; + warn!("{err}"); + } +} + +fn load_output_schema(path: Option) -> Option { + let path = path?; + + let schema_str = match std::fs::read_to_string(&path) { + Ok(contents) => contents, + Err(err) => { + eprintln!( + "Failed to read output schema file {}: {err}", + path.display() + ); + std::process::exit(1); + } + }; + + match serde_json::from_str::(&schema_str) { + Ok(value) => Some(value), + Err(err) => { + eprintln!( + "Output schema file {} is not valid JSON: {err}", + path.display() + ); + std::process::exit(1); + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum PromptDecodeError { + InvalidUtf8 { valid_up_to: usize }, + InvalidUtf16 { encoding: &'static str }, + UnsupportedBom { encoding: &'static str }, +} + +impl std::fmt::Display for PromptDecodeError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + PromptDecodeError::InvalidUtf8 { valid_up_to } => write!( + f, + "input is not valid UTF-8 (invalid byte at offset {valid_up_to}). Convert it to UTF-8 and retry (e.g., `iconv -f -t UTF-8 prompt.txt`)." + ), + PromptDecodeError::InvalidUtf16 { encoding } => write!( + f, + "input looked like {encoding} but could not be decoded. Convert it to UTF-8 and retry." + ), + PromptDecodeError::UnsupportedBom { encoding } => write!( + f, + "input appears to be {encoding}. Convert it to UTF-8 and retry." + ), + } + } +} + +fn decode_prompt_bytes(input: &[u8]) -> Result { + let input = input.strip_prefix(&[0xEF, 0xBB, 0xBF]).unwrap_or(input); + + if input.starts_with(&[0xFF, 0xFE, 0x00, 0x00]) { + return Err(PromptDecodeError::UnsupportedBom { + encoding: "UTF-32LE", + }); + } + + if input.starts_with(&[0x00, 0x00, 0xFE, 0xFF]) { + return Err(PromptDecodeError::UnsupportedBom { + encoding: "UTF-32BE", + }); + } + + if let Some(rest) = input.strip_prefix(&[0xFF, 0xFE]) { + return decode_utf16(rest, "UTF-16LE", u16::from_le_bytes); + } + + if let Some(rest) = input.strip_prefix(&[0xFE, 0xFF]) { + return decode_utf16(rest, "UTF-16BE", u16::from_be_bytes); + } + + std::str::from_utf8(input) + .map(str::to_string) + .map_err(|e| PromptDecodeError::InvalidUtf8 { + valid_up_to: e.valid_up_to(), + }) +} + +fn decode_utf16( + input: &[u8], + encoding: &'static str, + decode_unit: fn([u8; 2]) -> u16, +) -> Result { + if !input.len().is_multiple_of(2) { + return Err(PromptDecodeError::InvalidUtf16 { encoding }); + } + + let units: Vec = input + .chunks_exact(2) + .map(|chunk| decode_unit([chunk[0], chunk[1]])) + .collect(); + + String::from_utf16(&units).map_err(|_| PromptDecodeError::InvalidUtf16 { encoding }) +} + +fn read_prompt_from_stdin(behavior: StdinPromptBehavior) -> Option { + let stdin_is_terminal = std::io::stdin().is_terminal(); + + match behavior { + StdinPromptBehavior::RequiredIfPiped if stdin_is_terminal => { + eprintln!( + "No prompt provided. Either specify one as an argument or pipe the prompt into stdin." + ); + std::process::exit(1); + } + StdinPromptBehavior::RequiredIfPiped => { + eprintln!("Reading prompt from stdin..."); + } + StdinPromptBehavior::Forced => {} + StdinPromptBehavior::OptionalAppend if stdin_is_terminal => return None, + StdinPromptBehavior::OptionalAppend => { + eprintln!("Reading additional input from stdin..."); + } + } + + let mut bytes = Vec::new(); + if let Err(e) = std::io::stdin().read_to_end(&mut bytes) { + eprintln!("Failed to read prompt from stdin: {e}"); + std::process::exit(1); + } + + let buffer = match decode_prompt_bytes(&bytes) { + Ok(s) => s, + Err(e) => { + eprintln!("Failed to read prompt from stdin: {e}"); + std::process::exit(1); + } + }; + + if buffer.trim().is_empty() { + match behavior { + StdinPromptBehavior::OptionalAppend => None, + StdinPromptBehavior::RequiredIfPiped | StdinPromptBehavior::Forced => { + eprintln!("No prompt provided via stdin."); + std::process::exit(1); + } + } + } else { + Some(buffer) + } +} + +fn prompt_with_stdin_context(prompt: &str, stdin_text: &str) -> String { + let mut combined = format!("{prompt}\n\n\n{stdin_text}"); + if !stdin_text.ends_with('\n') { + combined.push('\n'); + } + combined.push_str(""); + combined +} + +fn resolve_prompt(prompt_arg: Option) -> String { + match prompt_arg { + Some(p) if p != "-" => p, + maybe_dash => { + let behavior = if matches!(maybe_dash.as_deref(), Some("-")) { + StdinPromptBehavior::Forced + } else { + StdinPromptBehavior::RequiredIfPiped + }; + let Some(prompt) = read_prompt_from_stdin(behavior) else { + unreachable!("required stdin prompt should produce content"); + }; + prompt + } + } +} + +fn resolve_root_prompt(prompt_arg: Option) -> String { + match prompt_arg { + Some(prompt) if prompt != "-" => { + if let Some(stdin_text) = read_prompt_from_stdin(StdinPromptBehavior::OptionalAppend) { + prompt_with_stdin_context(&prompt, &stdin_text) + } else { + prompt + } + } + maybe_dash => resolve_prompt(maybe_dash), + } +} + +fn build_review_request(args: &ReviewArgs) -> anyhow::Result { + let target = if args.uncommitted { + ReviewTarget::UncommittedChanges + } else if let Some(branch) = args.base.clone() { + ReviewTarget::BaseBranch { branch } + } else if let Some(sha) = args.commit.clone() { + ReviewTarget::Commit { + sha, + title: args.commit_title.clone(), + } + } else if let Some(prompt_arg) = args.prompt.clone() { + let prompt = resolve_prompt(Some(prompt_arg)).trim().to_string(); + if prompt.is_empty() { + anyhow::bail!("Review prompt cannot be empty"); + } + ReviewTarget::Custom { + instructions: prompt, + } + } else { + anyhow::bail!( + "Specify --uncommitted, --base, --commit, or provide custom review instructions" + ); + }; + + Ok(ReviewRequest { + target, + user_facing_hint: None, + }) +} + +#[cfg(test)] +#[path = "lib_tests.rs"] +mod tests; diff --git a/codex-rs/exec/src/lib_tests.rs b/codex-rs/exec/src/lib_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..fb79e7b83978ed3a7ed95475c3e016015670ee4f --- /dev/null +++ b/codex-rs/exec/src/lib_tests.rs @@ -0,0 +1,899 @@ +use super::*; +use codex_app_server_protocol::AuthRecoveryNotification; +use codex_otel::set_parent_from_w3c_trace_context; +use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::models::ActivePermissionProfile; +use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_WORKSPACE; +use codex_utils_absolute_path::test_support::PathBufExt; +use codex_utils_absolute_path::test_support::test_path_buf; +use opentelemetry::trace::TraceContextExt; +use opentelemetry::trace::TraceId; +use opentelemetry::trace::TracerProvider as _; +use opentelemetry_sdk::trace::SdkTracerProvider; +use pretty_assertions::assert_eq; +use std::io; +use std::io::Write; +use std::sync::Arc; +use std::sync::Mutex; +use tempfile::tempdir; +use tracing_opentelemetry::OpenTelemetrySpanExt; + +fn test_tracing_subscriber() -> impl tracing::Subscriber + Send + Sync { + let provider = SdkTracerProvider::builder().build(); + let tracer = provider.tracer("codex-exec-tests"); + tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer)) +} + +#[derive(Clone)] +struct TestLogWriter { + buffer: Arc>>, +} + +struct TestLogSink { + buffer: Arc>>, +} + +impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for TestLogWriter { + type Writer = TestLogSink; + + fn make_writer(&'a self) -> Self::Writer { + TestLogSink { + buffer: Arc::clone(&self.buffer), + } + } +} + +impl Write for TestLogSink { + fn write(&mut self, buf: &[u8]) -> io::Result { + self.buffer.lock().expect("log buffer lock").extend(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} + +#[test] +fn exec_default_stderr_filter_suppresses_otel_self_diagnostics() { + let buffer = Arc::new(Mutex::new(Vec::new())); + let writer = TestLogWriter { + buffer: Arc::clone(&buffer), + }; + let subscriber = tracing_subscriber::registry().with( + tracing_subscriber::fmt::layer() + .with_ansi(false) + .with_writer(writer) + .with_filter(EnvFilter::try_new(EXEC_DEFAULT_LOG_FILTER).expect("default filter")), + ); + + tracing::subscriber::with_default(subscriber, || { + tracing::error!(target: "opentelemetry_sdk", "telemetry export failed"); + tracing::error!(target: "opentelemetry_otlp", "telemetry request failed"); + tracing::error!(target: "codex_exec_test", "real exec error"); + }); + + let logs = String::from_utf8(buffer.lock().expect("log buffer lock").clone()).expect("utf8"); + assert!(!logs.contains("telemetry export failed")); + assert!(!logs.contains("telemetry request failed")); + assert!(logs.contains("real exec error")); +} + +#[test] +fn exec_root_span_can_be_parented_from_trace_context() { + let subscriber = test_tracing_subscriber(); + let _guard = tracing::subscriber::set_default(subscriber); + + let parent = codex_protocol::protocol::W3cTraceContext { + traceparent: Some("00-00000000000000000000000000000077-0000000000000088-01".into()), + tracestate: Some("vendor=value".into()), + }; + let exec_span = exec_root_span(); + assert!(set_parent_from_w3c_trace_context(&exec_span, &parent)); + + let trace_id = exec_span.context().span().span_context().trace_id(); + assert_eq!( + trace_id, + TraceId::from_hex("00000000000000000000000000000077").expect("trace id") + ); +} + +#[test] +fn builds_uncommitted_review_request() { + let args = ReviewArgs { + uncommitted: true, + base: None, + commit: None, + commit_title: None, + prompt: None, + }; + let request = build_review_request(&args).expect("builds uncommitted review request"); + + let expected = ReviewRequest { + target: ReviewTarget::UncommittedChanges, + user_facing_hint: None, + }; + + assert_eq!(request, expected); +} + +#[test] +fn builds_commit_review_request_with_title() { + let args = ReviewArgs { + uncommitted: false, + base: None, + commit: Some("123456789".to_string()), + commit_title: Some("Add review command".to_string()), + prompt: None, + }; + let request = build_review_request(&args).expect("builds commit review request"); + + let expected = ReviewRequest { + target: ReviewTarget::Commit { + sha: "123456789".to_string(), + title: Some("Add review command".to_string()), + }, + user_facing_hint: None, + }; + + assert_eq!(request, expected); +} + +#[test] +fn builds_custom_review_request_trims_prompt() { + let args = ReviewArgs { + uncommitted: false, + base: None, + commit: None, + commit_title: None, + prompt: Some(" custom review instructions ".to_string()), + }; + let request = build_review_request(&args).expect("builds custom review request"); + + let expected = ReviewRequest { + target: ReviewTarget::Custom { + instructions: "custom review instructions".to_string(), + }, + user_facing_hint: None, + }; + + assert_eq!(request, expected); +} + +#[test] +fn decode_prompt_bytes_strips_utf8_bom() { + let input = [0xEF, 0xBB, 0xBF, b'h', b'i', b'\n']; + + let out = decode_prompt_bytes(&input).expect("decode utf-8 with BOM"); + + assert_eq!(out, "hi\n"); +} + +#[test] +fn decode_prompt_bytes_decodes_utf16le_bom() { + // UTF-16LE BOM + "hi\n" + let input = [0xFF, 0xFE, b'h', 0x00, b'i', 0x00, b'\n', 0x00]; + + let out = decode_prompt_bytes(&input).expect("decode utf-16le with BOM"); + + assert_eq!(out, "hi\n"); +} + +#[test] +fn decode_prompt_bytes_decodes_utf16be_bom() { + // UTF-16BE BOM + "hi\n" + let input = [0xFE, 0xFF, 0x00, b'h', 0x00, b'i', 0x00, b'\n']; + + let out = decode_prompt_bytes(&input).expect("decode utf-16be with BOM"); + + assert_eq!(out, "hi\n"); +} + +#[test] +fn decode_prompt_bytes_rejects_utf32le_bom() { + // UTF-32LE BOM + "hi\n" + let input = [ + 0xFF, 0xFE, 0x00, 0x00, b'h', 0x00, 0x00, 0x00, b'i', 0x00, 0x00, 0x00, b'\n', 0x00, 0x00, + 0x00, + ]; + + let err = decode_prompt_bytes(&input).expect_err("utf-32le should be rejected"); + + assert_eq!( + err, + PromptDecodeError::UnsupportedBom { + encoding: "UTF-32LE" + } + ); +} + +#[test] +fn decode_prompt_bytes_rejects_utf32be_bom() { + // UTF-32BE BOM + "hi\n" + let input = [ + 0x00, 0x00, 0xFE, 0xFF, 0x00, 0x00, 0x00, b'h', 0x00, 0x00, 0x00, b'i', 0x00, 0x00, 0x00, + b'\n', + ]; + + let err = decode_prompt_bytes(&input).expect_err("utf-32be should be rejected"); + + assert_eq!( + err, + PromptDecodeError::UnsupportedBom { + encoding: "UTF-32BE" + } + ); +} + +#[test] +fn decode_prompt_bytes_rejects_invalid_utf8() { + // Invalid UTF-8 sequence: 0xC3 0x28 + let input = [0xC3, 0x28]; + + let err = decode_prompt_bytes(&input).expect_err("invalid utf-8 should fail"); + + assert_eq!(err, PromptDecodeError::InvalidUtf8 { valid_up_to: 0 }); +} + +#[test] +fn prompt_with_stdin_context_wraps_stdin_block() { + let combined = prompt_with_stdin_context("Summarize this concisely", "my output"); + + assert_eq!( + combined, + "Summarize this concisely\n\n\nmy output\n" + ); +} + +#[test] +fn prompt_with_stdin_context_preserves_trailing_newline() { + let combined = prompt_with_stdin_context("Summarize this concisely", "my output\n"); + + assert_eq!( + combined, + "Summarize this concisely\n\n\nmy output\n" + ); +} + +#[test] +fn lagged_event_warning_message_is_explicit() { + assert_eq!( + lagged_event_warning_message(/*skipped*/ 7), + "in-process app-server event stream lagged; dropped 7 events".to_string() + ); +} + +#[test] +fn runtime_warnings_are_filtered_to_the_primary_thread() { + let primary_thread_id = "thread-1"; + let turn_id = "turn-1"; + let outcomes = [ + codex_app_server_protocol::WarningNotification { + thread_id: None, + message: "global warning".to_string(), + }, + codex_app_server_protocol::WarningNotification { + thread_id: Some(primary_thread_id.to_string()), + message: "primary warning".to_string(), + }, + codex_app_server_protocol::WarningNotification { + thread_id: Some("thread-2".to_string()), + message: "other warning".to_string(), + }, + ] + .map(|warning| { + should_process_notification( + &ServerNotification::Warning(warning), + primary_thread_id, + turn_id, + ) + }); + + assert_eq!(outcomes, [true, true, false]); + + let recovery = AuthRecoveryNotification { + thread_id: primary_thread_id.to_string(), + turn_id: turn_id.to_string(), + provider: "example".to_string(), + message: "Refresh authentication".to_string(), + }; + let outcomes = [ + ServerNotification::AuthRecoveryStarted(recovery.clone()), + ServerNotification::AuthRecoveryCompleted(recovery.clone()), + ServerNotification::AuthRecoveryStarted(AuthRecoveryNotification { + thread_id: "thread-2".to_string(), + ..recovery.clone() + }), + ServerNotification::AuthRecoveryCompleted(AuthRecoveryNotification { + turn_id: "turn-2".to_string(), + ..recovery + }), + ] + .map(|notification| should_process_notification(¬ification, primary_thread_id, turn_id)); + + assert_eq!(outcomes, [true, true, false, false]); +} + +#[tokio::test] +async fn resume_lookup_model_providers_filters_only_last_lookup() { + let codex_home = tempdir().expect("create temp codex home"); + let cwd = tempdir().expect("create temp cwd"); + let mut config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .fallback_cwd(Some(cwd.path().to_path_buf())) + .build() + .await + .expect("build default config"); + config.model_provider_id = "test-provider".to_string(); + + let last_args = crate::cli::ResumeArgs { + session_id: None, + last: true, + all: false, + images: vec![], + prompt: None, + }; + let named_args = crate::cli::ResumeArgs { + session_id: Some("named-session".to_string()), + last: false, + all: false, + images: vec![], + prompt: None, + }; + + assert_eq!( + resume_lookup_model_providers(&config, &last_args), + Some(vec!["test-provider".to_string()]) + ); + assert_eq!(resume_lookup_model_providers(&config, &named_args), None); +} + +#[test] +fn turn_items_for_thread_returns_matching_turn_items() { + let thread = AppServerThread { + environments: None, + id: "thread-1".to_string(), + extra: None, + session_id: "thread-1".to_string(), + forked_from_id: None, + parent_thread_id: None, + preview: String::new(), + ephemeral: false, + section: None, + section_entered_at: None, + project_id: None, + daybreak_enabled: None, + history_mode: Default::default(), + model_provider: "openai".to_string(), + model: None, + reasoning_effort: None, + created_at: 0, + updated_at: 0, + recency_at: Some(0), + status: codex_app_server_protocol::ThreadStatus::Idle, + path: None, + cwd: test_path_buf("/tmp/project").abs(), + cli_version: "0.0.0-test".to_string(), + originator: None, + source: codex_app_server_protocol::SessionSource::Exec, + can_accept_direct_input: None, + thread_source: None, + agent_nickname: None, + agent_role: None, + git_info: None, + name: None, + turns: vec![ + codex_app_server_protocol::Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: vec![AppServerThreadItem::AgentMessage { + id: "msg-1".to_string(), + text: "hello".to_string(), + phase: None, + memory_citation: None, + delivery: None, + questions: None, + }], + status: codex_app_server_protocol::TurnStatus::Completed, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + }, + codex_app_server_protocol::Turn { + id: "turn-2".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: vec![AppServerThreadItem::Plan { + id: "plan-1".to_string(), + text: "ship it".to_string(), + }], + status: codex_app_server_protocol::TurnStatus::Completed, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + }, + ], + }; + + assert_eq!( + turn_items_for_thread(&thread, "turn-1"), + Some(vec![AppServerThreadItem::AgentMessage { + id: "msg-1".to_string(), + text: "hello".to_string(), + phase: None, + memory_citation: None, + delivery: None, + questions: None, + }]) + ); + assert_eq!(turn_items_for_thread(&thread, "missing-turn"), None); +} + +#[test] +fn should_backfill_turn_completed_items_backfills_persisted_summaries_only() { + let notification = + ServerNotification::TurnCompleted(codex_app_server_protocol::TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: codex_app_server_protocol::Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Summary, + items: Vec::new(), + status: codex_app_server_protocol::TurnStatus::Completed, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + }, + }); + + assert!(!should_backfill_turn_completed_items( + /*thread_ephemeral*/ true, + ¬ification + )); + assert!(should_backfill_turn_completed_items( + /*thread_ephemeral*/ false, + ¬ification + )); +} + +#[test] +fn canceled_mcp_server_elicitation_response_uses_cancel_action() { + let value = canceled_mcp_server_elicitation_response() + .expect("mcp elicitation cancel response should serialize"); + let response: McpServerElicitationRequestResponse = + serde_json::from_value(value).expect("cancel response should deserialize"); + + assert_eq!( + response, + McpServerElicitationRequestResponse { + action: McpServerElicitationAction::Cancel, + content: None, + meta: None, + } + ); +} + +#[tokio::test] +async fn thread_start_params_include_review_policy_when_review_policy_is_manual_only() { + let codex_home = tempdir().expect("create temp codex home"); + let cwd = tempdir().expect("create temp cwd"); + let config = ConfigBuilder::default() + .loader_overrides(LoaderOverrides::without_managed_config_for_tests()) + .codex_home(codex_home.path().to_path_buf()) + .harness_overrides(ConfigOverrides { + approvals_reviewer: Some(ApprovalsReviewer::User), + ..Default::default() + }) + .fallback_cwd(Some(cwd.path().to_path_buf())) + .build() + .await + .expect("build config with manual-only review policy"); + + let params = thread_start_params_from_config(&config, &ThreadSource::User); + + assert_eq!( + params.approvals_reviewer, + Some(codex_app_server_protocol::ApprovalsReviewer::User) + ); + assert_eq!(params.sandbox, None); + assert_eq!( + params.permissions, + permissions_selection_from_config(&config) + ); +} + +#[tokio::test] +async fn thread_start_params_include_review_policy_when_auto_review_is_enabled() { + let codex_home = tempdir().expect("create temp codex home"); + let cwd = tempdir().expect("create temp cwd"); + let config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .harness_overrides(ConfigOverrides { + approvals_reviewer: Some(ApprovalsReviewer::AutoReview), + ..Default::default() + }) + .fallback_cwd(Some(cwd.path().to_path_buf())) + .build() + .await + .expect("build config with guardian review policy"); + + let params = thread_start_params_from_config(&config, &ThreadSource::User); + + assert_eq!( + params.approvals_reviewer, + Some(codex_app_server_protocol::ApprovalsReviewer::AutoReview) + ); +} + +#[tokio::test] +async fn thread_resume_params_only_include_explicit_review_policy_override() { + let codex_home = tempdir().expect("create temp codex home"); + let cwd = tempdir().expect("create temp cwd"); + let config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .harness_overrides(ConfigOverrides { + approvals_reviewer: Some(ApprovalsReviewer::AutoReview), + ..Default::default() + }) + .fallback_cwd(Some(cwd.path().to_path_buf())) + .build() + .await + .expect("build config with guardian review policy"); + + let params_without_override = thread_resume_params_from_config( + &config, + "thread-id".to_string(), + /*approvals_reviewer_override*/ None, + ); + let params_with_override = thread_resume_params_from_config( + &config, + "thread-id".to_string(), + Some(codex_app_server_protocol::ApprovalsReviewer::AutoReview), + ); + + assert_eq!(params_without_override.approvals_reviewer, None); + assert_eq!( + params_with_override.approvals_reviewer, + Some(codex_app_server_protocol::ApprovalsReviewer::AutoReview) + ); +} + +#[tokio::test] +async fn build_exec_config_retries_without_invalid_headless_policy_for_auto_review() { + let codex_home = tempdir().expect("create temp codex home"); + let cwd = tempdir().expect("create temp cwd"); + std::fs::write( + codex_home.path().join("config.toml"), + r#" +approval_policy = "on-request" +approvals_reviewer = "auto_review" +"#, + ) + .expect("write config"); + let requirements_path = codex_home.path().join("requirements.toml"); + std::fs::write( + &requirements_path, + r#" +allowed_approval_policies = ["never", "on-request"] +allowed_sandbox_modes = ["read-only", "workspace-write"] +"#, + ) + .expect("write requirements"); + let mut loader_overrides = LoaderOverrides::without_managed_config_for_tests(); + loader_overrides.system_requirements_path = Some(requirements_path); + let overrides = ConfigOverrides { + cwd: Some(cwd.path().to_path_buf()), + approval_policy: Some(AskForApproval::Never), + sandbox_mode: Some(SandboxMode::DangerFullAccess), + ..Default::default() + }; + let build_config = |overrides| { + ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .loader_overrides(loader_overrides.clone()) + .harness_overrides(overrides) + .build() + }; + + let error = build_config(overrides.clone()) + .await + .expect_err("synthetic headless approval policy should fail"); + assert!( + error + .to_string() + .contains("`approval_policy = \"never\"` cannot be used") + ); + + let config = build_exec_config( + overrides, + /*preserve_headless_approval_policy*/ false, + build_config, + ) + .await + .expect("auto-review config should retry without the synthetic approval policy"); + + assert_eq!( + config.permissions.approval_policy.value(), + AskForApproval::OnRequest + ); + assert_eq!(config.approvals_reviewer, ApprovalsReviewer::AutoReview); +} + +#[tokio::test] +async fn build_exec_config_preserves_headless_error_when_retry_fails() { + let overrides = ConfigOverrides { + approval_policy: Some(AskForApproval::Never), + ..Default::default() + }; + + let error = build_exec_config( + overrides, + /*preserve_headless_approval_policy*/ false, + |overrides| async move { + let message = if overrides.approval_policy == Some(AskForApproval::Never) { + "headless error" + } else { + "retry error" + }; + Err(std::io::Error::other(message)) + }, + ) + .await + .expect_err("failed speculative retry should preserve the original error"); + + assert_eq!(error.to_string(), "headless error"); +} + +#[tokio::test] +async fn thread_start_params_match_history_to_persistence() { + let codex_home = tempdir().expect("create temp codex home"); + let cwd = tempdir().expect("create temp cwd"); + let mut config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .fallback_cwd(Some(cwd.path().to_path_buf())) + .build() + .await + .expect("build config"); + + let params = thread_start_params_from_config(&config, &ThreadSource::User); + + assert_eq!( + params.thread_source, + Some(codex_app_server_protocol::ThreadSource::User) + ); + assert_eq!(params.history_mode, Some(ThreadHistoryMode::Paginated)); + + let thread_source = ThreadSource::Feature("automated_review".to_string()); + let params = thread_start_params_from_config(&config, &thread_source); + assert_eq!(params.thread_source, Some(thread_source)); + + config.ephemeral = true; + let params = thread_start_params_from_config(&config, &ThreadSource::User); + + assert_eq!(params.ephemeral, Some(true)); + assert_eq!(params.history_mode, None); +} + +#[tokio::test] +async fn thread_lifecycle_params_preserve_hook_trust_bypass() { + let codex_home = tempdir().expect("create temp codex home"); + let cwd = tempdir().expect("create temp cwd"); + let config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .harness_overrides(ConfigOverrides { + bypass_hook_trust: Some(true), + ..Default::default() + }) + .fallback_cwd(Some(cwd.path().to_path_buf())) + .build() + .await + .expect("build config with hook trust bypass"); + let expected_config = Some(HashMap::from([( + "bypass_hook_trust".to_string(), + serde_json::Value::Bool(true), + )])); + + let start_params = thread_start_params_from_config(&config, &ThreadSource::User); + let resume_params = thread_resume_params_from_config( + &config, + "thread-id".to_string(), + /*approvals_reviewer_override*/ None, + ); + + assert_eq!(start_params.config, expected_config); + assert_eq!(resume_params.config, expected_config); +} + +#[test] +fn active_profile_selection_uses_profile_id_only() { + let selection = permission_profile_id_from_active_profile(ActivePermissionProfile::new( + BUILT_IN_PERMISSION_PROFILE_WORKSPACE, + )); + + assert_eq!(selection, BUILT_IN_PERMISSION_PROFILE_WORKSPACE.to_string()); +} + +#[tokio::test] +async fn thread_lifecycle_params_include_legacy_sandbox_when_no_active_profile() { + let codex_home = tempdir().expect("create temp codex home"); + let cwd = tempdir().expect("create temp cwd"); + let config = ConfigBuilder::default() + .loader_overrides(LoaderOverrides::without_managed_config_for_tests()) + .codex_home(codex_home.path().to_path_buf()) + .harness_overrides(ConfigOverrides { + sandbox_mode: Some(SandboxMode::DangerFullAccess), + ..Default::default() + }) + .fallback_cwd(Some(cwd.path().to_path_buf())) + .build() + .await + .expect("build config with legacy sandbox override"); + + let start_params = thread_start_params_from_config(&config, &ThreadSource::User); + let resume_params = thread_resume_params_from_config( + &config, + "thread-id".to_string(), + /*approvals_reviewer_override*/ None, + ); + + assert_eq!(config.permissions.active_permission_profile(), None); + assert_eq!( + start_params.sandbox, + Some(codex_app_server_protocol::SandboxMode::DangerFullAccess) + ); + assert_eq!(start_params.permissions, None); + assert_eq!( + resume_params.sandbox, + Some(codex_app_server_protocol::SandboxMode::DangerFullAccess) + ); + assert_eq!(resume_params.permissions, None); +} + +#[tokio::test] +async fn session_configured_from_thread_response_uses_review_policy_from_response() { + let codex_home = tempdir().expect("create temp codex home"); + let cwd = tempdir().expect("create temp cwd"); + let config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .fallback_cwd(Some(cwd.path().to_path_buf())) + .build() + .await + .expect("build config"); + let response = sample_thread_start_response(); + + let event = session_configured_from_thread_start_response(&response, &config) + .expect("build bootstrap session configured event"); + + assert_eq!( + event.session_id.to_string(), + "67e55044-10b1-426f-9247-bb680e5fe0c7" + ); + assert_eq!( + event.thread_id.to_string(), + "67e55044-10b1-426f-9247-bb680e5fe0c8" + ); + assert_eq!(event.approvals_reviewer, ApprovalsReviewer::AutoReview); +} + +#[tokio::test] +async fn session_configured_from_thread_response_uses_permission_profile_from_config() { + let codex_home = tempdir().expect("create temp codex home"); + let cwd = tempdir().expect("create temp cwd"); + let config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .fallback_cwd(Some(cwd.path().to_path_buf())) + .build() + .await + .expect("build config"); + let response = sample_thread_start_response(); + + let event = session_configured_from_thread_start_response(&response, &config) + .expect("build bootstrap session configured event"); + + assert_eq!( + event.permission_profile, + config.permissions.effective_permission_profile() + ); +} + +#[tokio::test] +async fn session_configured_from_thread_response_preserves_thread_source() { + let codex_home = tempdir().expect("create temp codex home"); + let cwd = tempdir().expect("create temp cwd"); + let config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .fallback_cwd(Some(cwd.path().to_path_buf())) + .build() + .await + .expect("build config"); + let response = sample_thread_start_response(); + + let event = session_configured_from_thread_start_response(&response, &config) + .expect("build bootstrap session configured event"); + + assert_eq!( + event.thread_source, + Some(codex_protocol::protocol::ThreadSource::User) + ); +} + +#[tokio::test] +async fn session_configured_from_thread_response_preserves_parent_thread_id() { + let codex_home = tempdir().expect("create temp codex home"); + let cwd = tempdir().expect("create temp cwd"); + let config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .fallback_cwd(Some(cwd.path().to_path_buf())) + .build() + .await + .expect("build config"); + let parent_thread_id = ThreadId::new(); + let forked_from_id = ThreadId::new(); + let mut response = sample_thread_start_response(); + response.thread.parent_thread_id = Some(parent_thread_id.to_string()); + response.thread.forked_from_id = Some(forked_from_id.to_string()); + + let event = session_configured_from_thread_start_response(&response, &config) + .expect("build bootstrap session configured event"); + + assert_eq!(event.parent_thread_id, Some(parent_thread_id)); + assert_eq!(event.forked_from_id, Some(forked_from_id)); +} + +fn sample_thread_start_response() -> ThreadStartResponse { + ThreadStartResponse { + disabled_plugin_ids: Vec::new(), + thread: codex_app_server_protocol::Thread { + originator: None, + environments: None, + id: "67e55044-10b1-426f-9247-bb680e5fe0c8".to_string(), + extra: None, + session_id: "67e55044-10b1-426f-9247-bb680e5fe0c7".to_string(), + forked_from_id: None, + parent_thread_id: None, + preview: String::new(), + ephemeral: false, + section: None, + section_entered_at: None, + project_id: None, + daybreak_enabled: None, + history_mode: Default::default(), + model_provider: "openai".to_string(), + model: None, + reasoning_effort: None, + created_at: 0, + updated_at: 0, + recency_at: Some(0), + status: codex_app_server_protocol::ThreadStatus::Idle, + path: Some(PathBuf::from("/tmp/rollout.jsonl")), + cwd: test_path_buf("/tmp").abs(), + cli_version: "0.0.0".to_string(), + source: codex_app_server_protocol::SessionSource::Cli, + can_accept_direct_input: None, + thread_source: Some(codex_app_server_protocol::ThreadSource::User), + agent_nickname: None, + agent_role: None, + git_info: None, + name: Some("thread".to_string()), + turns: vec![], + }, + model: "gpt-5.4".to_string(), + model_provider: "openai".to_string(), + service_tier: None, + cwd: test_path_buf("/tmp").abs(), + runtime_workspace_roots: Vec::new(), + instruction_sources: Vec::new(), + approval_policy: codex_app_server_protocol::AskForApproval::OnRequest, + approvals_reviewer: codex_app_server_protocol::ApprovalsReviewer::AutoReview, + sandbox: codex_app_server_protocol::SandboxPolicy::WorkspaceWrite { + writable_roots: vec![], + network_access: false, + exclude_tmpdir_env_var: false, + exclude_slash_tmp: false, + }, + active_permission_profile: None, + reasoning_effort: None, + multi_agent_mode: Default::default(), + } +} diff --git a/codex-rs/exec/src/main.rs b/codex-rs/exec/src/main.rs new file mode 100644 index 0000000000000000000000000000000000000000..61eaecdd0a8526c49e798649f3810eec343200e4 --- /dev/null +++ b/codex-rs/exec/src/main.rs @@ -0,0 +1,44 @@ +//! Entry-point for the `codex-exec` binary. +//! +//! When this CLI is invoked normally, it parses the standard `codex-exec` CLI +//! options and launches the non-interactive Codex agent. However, if it is +//! invoked with arg0 as `codex-linux-sandbox`, we instead treat the invocation +//! as a request to run the logic for the standalone `codex-linux-sandbox` +//! executable (i.e., parse any -s args and then run a *sandboxed* command under +//! Landlock + seccomp. +//! +//! This allows us to ship a completely separate set of functionality as part +//! of the `codex-exec` binary. +use clap::Parser; +use codex_arg0::Arg0DispatchPaths; +use codex_arg0::arg0_dispatch_or_else; +use codex_exec::Cli; +use codex_exec::run_main; +use codex_utils_cli::CliConfigOverrides; + +#[derive(Parser, Debug)] +struct TopCli { + #[clap(flatten)] + config_overrides: CliConfigOverrides, + + #[clap(flatten)] + inner: Cli, +} + +fn main() -> anyhow::Result<()> { + arg0_dispatch_or_else(|arg0_paths: Arg0DispatchPaths| async move { + let top_cli = TopCli::parse(); + // Merge root-level overrides into inner CLI struct so downstream logic remains unchanged. + let mut inner = top_cli.inner; + inner + .config_overrides + .prepend_root_overrides(top_cli.config_overrides); + + run_main(inner, arg0_paths).await?; + Ok(()) + }) +} + +#[cfg(test)] +#[path = "main_tests.rs"] +mod tests; diff --git a/codex-rs/exec/src/main_tests.rs b/codex-rs/exec/src/main_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..5c0a8a3bfceb1f34be07c894527e301682384894 --- /dev/null +++ b/codex-rs/exec/src/main_tests.rs @@ -0,0 +1,43 @@ +use super::*; +use pretty_assertions::assert_eq; + +#[test] +fn top_cli_parses_resume_prompt_after_config_flag() { + const PROMPT: &str = "echo resume-with-global-flags-after-subcommand"; + let cli = TopCli::parse_from([ + "codex-exec", + "resume", + "--strict-config", + "--last", + "--json", + "--model", + "gpt-5.2-codex", + "--config", + "reasoning_level=xhigh", + "--dangerously-bypass-approvals-and-sandbox", + "--skip-git-repo-check", + PROMPT, + ]); + let mut inner = cli.inner; + inner + .config_overrides + .prepend_root_overrides(cli.config_overrides); + + let Some(codex_exec::Command::Resume(args)) = inner.command.as_ref() else { + panic!("expected resume command"); + }; + let effective_prompt = args.prompt.clone().or_else(|| { + if args.last { + args.session_id.clone() + } else { + None + } + }); + assert_eq!(effective_prompt.as_deref(), Some(PROMPT)); + assert_eq!(inner.config_overrides.raw_overrides.len(), 1); + assert_eq!( + inner.config_overrides.raw_overrides[0], + "reasoning_level=xhigh" + ); + assert!(inner.strict_config); +} diff --git a/codex-rs/exec/src/worktree.rs b/codex-rs/exec/src/worktree.rs new file mode 100644 index 0000000000000000000000000000000000000000..d73e3ad639f7260c63a55bdc6f2eee5fe301febd --- /dev/null +++ b/codex-rs/exec/src/worktree.rs @@ -0,0 +1,82 @@ +//! Resolves managed fork sources through the existing name/ID lookup before final configuration. +//! The temporary app-server never starts a thread or executes a turn. + +use super::*; + +pub(super) async fn fork_source( + args: &mut crate::cli::ForkArgs, + config: &Config, + arg0_paths: &Arg0DispatchPaths, + cli_overrides: &[(String, codex_config::TomlValue)], + loader_overrides: &LoaderOverrides, + cloud_config_bundle: CloudConfigBundleLoader, + strict_config: bool, +) -> anyhow::Result { + let state_db = codex_core::init_state_db(config).await; + let environment_manager = EnvironmentManager::from_codex_home( + config.codex_home.clone(), + Some(ExecServerRuntimePaths::from_optional_paths( + arg0_paths.codex_self_exe.clone(), + arg0_paths.codex_linux_sandbox_exe.clone(), + )?), + config.http_client_factory(), + ) + .await?; + let client = InProcessAppServerClient::start(InProcessClientStartArgs { + arg0_paths: arg0_paths.clone(), + config: std::sync::Arc::new(config.clone()), + cli_overrides: cli_overrides.to_vec(), + loader_overrides: LoaderOverrides { + ignore_project_config: true, + ..loader_overrides.clone() + }, + strict_config, + cloud_config_bundle, + feedback: CodexFeedback::new(), + log_db: None, + state_db: state_db.clone(), + environment_manager: std::sync::Arc::new(environment_manager), + config_warnings: Vec::new(), + session_source: SessionSource::Exec, + enable_codex_api_key_env: true, + client_name: "codex_exec".to_string(), + client_version: env!("CARGO_PKG_VERSION").to_string(), + experimental_api: true, + mcp_server_openai_form_elicitation: false, + opt_out_notification_methods: Vec::new(), + channel_capacity: DEFAULT_IN_PROCESS_CHANNEL_CAPACITY, + }) + .await?; + let result = async { + let lookup = crate::cli::ResumeArgs { + session_id: Some(args.session_id.clone()), + last: false, + all: true, + images: Vec::new(), + prompt: None, + }; + let thread_id = resolve_resume_thread_id(&client, config, state_db.as_ref(), &lookup) + .await? + .ok_or_else(|| anyhow::anyhow!("Session not found: {}", args.session_id))?; + let source: ThreadReadResponse = send_request_with_response( + &client, + ClientRequest::ThreadRead { + request_id: RequestId::Integer(1), + params: ThreadReadParams { + thread_id: thread_id.clone(), + include_turns: false, + }, + }, + "thread/read", + ) + .await + .map_err(anyhow::Error::msg)?; + args.session_id = thread_id; + Ok::<_, anyhow::Error>(latest_thread_cwd(&source.thread).await) + } + .await; + let shutdown = client.shutdown().await; + let source_cwd = result?; + shutdown?; + Ok(source_cwd) +} diff --git a/codex-rs/exec/tests/all.rs b/codex-rs/exec/tests/all.rs new file mode 100644 index 0000000000000000000000000000000000000000..6582e12b2312bce8cfa7eea5d230e54258d37f42 --- /dev/null +++ b/codex-rs/exec/tests/all.rs @@ -0,0 +1,7 @@ +#![allow(clippy::expect_used)] + +// Single integration test binary that aggregates all test modules. +// The submodules live in `tests/suite/`. +mod suite; + +mod event_processor_with_json_output; diff --git a/codex-rs/exec/tests/event_processor_with_json_output.rs b/codex-rs/exec/tests/event_processor_with_json_output.rs new file mode 100644 index 0000000000000000000000000000000000000000..1586c9303daaab9b12d4cd833f2eb04820fa4824 --- /dev/null +++ b/codex-rs/exec/tests/event_processor_with_json_output.rs @@ -0,0 +1,1711 @@ +use codex_app_server_protocol::CollabAgentState as ApiCollabAgentState; +use codex_app_server_protocol::CollabAgentStatus as ApiCollabAgentStatus; +use codex_app_server_protocol::CollabAgentTool; +use codex_app_server_protocol::CollabAgentToolCallStatus as ApiCollabAgentToolCallStatus; +use codex_app_server_protocol::CommandAction; +use codex_app_server_protocol::CommandExecutionSource; +use codex_app_server_protocol::CommandExecutionStatus as ApiCommandExecutionStatus; +use codex_app_server_protocol::ErrorNotification; +use codex_app_server_protocol::FileUpdateChange as ApiFileUpdateChange; +use codex_app_server_protocol::ItemCompletedNotification; +use codex_app_server_protocol::ItemStartedNotification; +use codex_app_server_protocol::McpToolCallError; +use codex_app_server_protocol::McpToolCallResult; +use codex_app_server_protocol::McpToolCallStatus as ApiMcpToolCallStatus; +use codex_app_server_protocol::PatchApplyStatus as ApiPatchApplyStatus; +use codex_app_server_protocol::PatchChangeKind as ApiPatchChangeKind; +use codex_app_server_protocol::ServerNotification; +use codex_app_server_protocol::ThreadItem; +use codex_app_server_protocol::ThreadTokenUsage; +use codex_app_server_protocol::TokenUsageBreakdown; +use codex_app_server_protocol::Turn; +use codex_app_server_protocol::TurnCompletedNotification; +use codex_app_server_protocol::TurnError; +use codex_app_server_protocol::TurnPlanStep; +use codex_app_server_protocol::TurnPlanStepStatus; +use codex_app_server_protocol::TurnPlanUpdatedNotification; +use codex_app_server_protocol::TurnStartedNotification; +use codex_app_server_protocol::TurnStatus; +use codex_app_server_protocol::WebSearchAction as ApiWebSearchAction; +use codex_app_server_protocol::WebSearchItem as ApiWebSearchItem; +use codex_protocol::SessionId; +use codex_protocol::ThreadId; +use codex_protocol::models::PermissionProfile; +use codex_protocol::models::WebSearchAction; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::SessionConfiguredEvent; +use codex_utils_absolute_path::test_support::PathBufExt; +use codex_utils_absolute_path::test_support::test_path_buf; +use pretty_assertions::assert_eq; +use serde_json::json; + +use codex_exec::AgentMessageItem; +use codex_exec::CodexStatus; +use codex_exec::CollabAgentState; +use codex_exec::CollabAgentStatus; +use codex_exec::CollabTool; +use codex_exec::CollabToolCallItem; +use codex_exec::CollabToolCallStatus; +use codex_exec::CollectedThreadEvents; +use codex_exec::CommandExecutionItem; +use codex_exec::CommandExecutionStatus; +use codex_exec::ErrorItem; +use codex_exec::EventProcessorWithJsonOutput; +use codex_exec::ExecThreadItem; +use codex_exec::FileChangeItem; +use codex_exec::FileUpdateChange as ExecFileUpdateChange; +use codex_exec::ItemCompletedEvent; +use codex_exec::ItemStartedEvent; +use codex_exec::ItemUpdatedEvent; +use codex_exec::McpToolCallItem; +use codex_exec::McpToolCallItemError; +use codex_exec::McpToolCallItemResult; +use codex_exec::McpToolCallStatus; +use codex_exec::PatchApplyStatus; +use codex_exec::PatchChangeKind; +use codex_exec::ReasoningItem; +use codex_exec::ThreadErrorEvent; +use codex_exec::ThreadEvent; +use codex_exec::ThreadItemDetails; +use codex_exec::ThreadStartedEvent; +use codex_exec::TodoItem; +use codex_exec::TodoListItem; +use codex_exec::TurnCompletedEvent; +use codex_exec::TurnFailedEvent; +use codex_exec::TurnStartedEvent; +use codex_exec::Usage; +use codex_exec::WebSearchItem; + +#[test] +fn map_todo_items_preserves_text_and_completion_state() { + let items = EventProcessorWithJsonOutput::map_todo_items(&[ + TurnPlanStep { + step: "inspect bootstrap".to_string(), + status: TurnPlanStepStatus::InProgress, + }, + TurnPlanStep { + step: "drop legacy notifications".to_string(), + status: TurnPlanStepStatus::Completed, + }, + ]); + + assert_eq!( + items, + vec![ + TodoItem { + text: "inspect bootstrap".to_string(), + completed: false, + }, + TodoItem { + text: "drop legacy notifications".to_string(), + completed: true, + }, + ] + ); +} + +#[test] +fn session_configured_produces_thread_started_event() { + let thread_id = ThreadId::from_string("67e55044-10b1-426f-9247-bb680e5fe0c8") + .expect("thread id should parse"); + let session_configured = SessionConfiguredEvent { + session_id: SessionId::from(thread_id), + thread_id, + forked_from_id: None, + parent_thread_id: None, + thread_source: None, + thread_name: None, + model: "codex-mini-latest".to_string(), + model_provider_id: "test-provider".to_string(), + service_tier: None, + approval_policy: AskForApproval::Never, + approvals_reviewer: codex_protocol::config_types::ApprovalsReviewer::User, + permission_profile: PermissionProfile::read_only(), + active_permission_profile: None, + cwd: test_path_buf("/tmp/project").abs(), + reasoning_effort: None, + initial_messages: None, + network_proxy: None, + rollout_path: None, + }; + + assert_eq!( + EventProcessorWithJsonOutput::thread_started_event(&session_configured), + ThreadEvent::ThreadStarted(ThreadStartedEvent { + thread_id: "67e55044-10b1-426f-9247-bb680e5fe0c8".to_string(), + }) + ); +} + +#[test] +fn turn_started_emits_turn_started_event() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let collected = + processor.collect_thread_events(ServerNotification::TurnStarted(TurnStartedNotification { + thread_id: "thread-1".to_string(), + turn: Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: Vec::new(), + status: TurnStatus::InProgress, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + }, + })); + + assert_eq!( + collected, + CollectedThreadEvents { + events: vec![ThreadEvent::TurnStarted(TurnStartedEvent {})], + status: CodexStatus::Running, + } + ); +} + +#[test] +fn command_execution_started_and_completed_translate_to_thread_events() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + let command_item = ThreadItem::CommandExecution { + model_context: None, + id: "cmd-1".to_string(), + command: "ls".to_string(), + cwd: test_path_buf("/tmp/project").abs().into(), + process_id: Some("123".to_string()), + plugin_id: None, + script_path: None, + source: CommandExecutionSource::UserShell, + status: ApiCommandExecutionStatus::InProgress, + command_actions: Vec::::new(), + aggregated_output: None, + exit_code: None, + duration_ms: None, + }; + + let started = + processor.collect_thread_events(ServerNotification::ItemStarted(ItemStartedNotification { + item: command_item, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + started_at_ms: 0, + })); + assert_eq!( + started, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemStarted(ItemStartedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::CommandExecution(CommandExecutionItem { + command: "ls".to_string(), + aggregated_output: String::new(), + exit_code: None, + status: CommandExecutionStatus::InProgress, + }), + }, + })], + status: CodexStatus::Running, + } + ); + + let completed = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::CommandExecution { + model_context: None, + id: "cmd-1".to_string(), + command: "ls".to_string(), + cwd: test_path_buf("/tmp/project").abs().into(), + process_id: Some("123".to_string()), + plugin_id: None, + script_path: None, + source: CommandExecutionSource::UserShell, + status: ApiCommandExecutionStatus::Completed, + command_actions: Vec::::new(), + aggregated_output: Some("a.txt\n".to_string()), + exit_code: Some(0), + duration_ms: Some(3), + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + assert_eq!( + completed, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::CommandExecution(CommandExecutionItem { + command: "ls".to_string(), + aggregated_output: "a.txt\n".to_string(), + exit_code: Some(0), + status: CommandExecutionStatus::Completed, + }), + }, + })], + status: CodexStatus::Running, + } + ); +} + +#[test] +fn empty_reasoning_items_are_ignored() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let collected = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::Reasoning { + id: "reasoning-1".to_string(), + summary: Vec::new(), + content: vec!["raw reasoning".to_string()], + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!( + collected, + CollectedThreadEvents { + events: Vec::new(), + status: CodexStatus::Running, + } + ); +} + +#[test] +fn unsupported_items_do_not_consume_synthetic_ids() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let ignored = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::Plan { + id: "plan-1".to_string(), + text: "ignored plan".to_string(), + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!( + ignored, + CollectedThreadEvents { + events: Vec::new(), + status: CodexStatus::Running, + } + ); + + let collected = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::AgentMessage { + id: "message-1".to_string(), + text: "hello".to_string(), + phase: None, + memory_citation: None, + delivery: None, + questions: None, + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!( + collected, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::AgentMessage(AgentMessageItem { + text: "hello".to_string(), + }), + }, + })], + status: CodexStatus::Running, + } + ); +} + +#[test] +fn reasoning_items_emit_summary_not_raw_content() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let collected = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::Reasoning { + id: "reasoning-1".to_string(), + summary: vec!["safe summary".to_string()], + content: vec!["raw reasoning".to_string()], + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!( + collected, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::Reasoning(ReasoningItem { + text: "safe summary".to_string(), + }), + }, + })], + status: CodexStatus::Running, + } + ); +} + +#[test] +fn web_search_completion_preserves_query_and_action() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let collected = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::WebSearch(ApiWebSearchItem { + id: "search-1".to_string(), + query: "rust async await".to_string(), + action: Some(ApiWebSearchAction::Search { + query: Some("rust async await".to_string()), + queries: None, + }), + results: None, + }), + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!( + collected, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::WebSearch(WebSearchItem { + id: "search-1".to_string(), + query: "rust async await".to_string(), + action: WebSearchAction::Search { + query: Some("rust async await".to_string()), + queries: None, + }, + }), + }, + })], + status: CodexStatus::Running, + } + ); +} + +#[test] +fn web_search_start_and_completion_reuse_item_id() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let started = + processor.collect_thread_events(ServerNotification::ItemStarted(ItemStartedNotification { + item: ThreadItem::WebSearch(ApiWebSearchItem { + id: "search-1".to_string(), + query: String::new(), + action: None, + results: None, + }), + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + started_at_ms: 0, + })); + + let completed = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::WebSearch(ApiWebSearchItem { + id: "search-1".to_string(), + query: "rust async await".to_string(), + action: Some(ApiWebSearchAction::Search { + query: Some("rust async await".to_string()), + queries: None, + }), + results: None, + }), + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!( + started, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemStarted(ItemStartedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::WebSearch(WebSearchItem { + id: "search-1".to_string(), + query: String::new(), + action: WebSearchAction::Other, + }), + }, + })], + status: CodexStatus::Running, + } + ); + assert_eq!( + completed, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::WebSearch(WebSearchItem { + id: "search-1".to_string(), + query: "rust async await".to_string(), + action: WebSearchAction::Search { + query: Some("rust async await".to_string()), + queries: None, + }, + }), + }, + })], + status: CodexStatus::Running, + } + ); +} + +#[test] +fn mcp_tool_call_begin_and_end_emit_item_events() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let started = + processor.collect_thread_events(ServerNotification::ItemStarted(ItemStartedNotification { + item: ThreadItem::McpToolCall { + id: "mcp-1".to_string(), + server: "server_a".to_string(), + tool: "tool_x".to_string(), + status: ApiMcpToolCallStatus::InProgress, + arguments: json!({ "key": "value" }), + app_context: None, + mcp_app_resource_uri: None, + mcp_app_ui: None, + plugin_id: None, + read_only_hint: None, + result: None, + error: None, + duration_ms: None, + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + started_at_ms: 0, + })); + let completed = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::McpToolCall { + id: "mcp-1".to_string(), + server: "server_a".to_string(), + tool: "tool_x".to_string(), + status: ApiMcpToolCallStatus::Completed, + arguments: json!({ "key": "value" }), + app_context: None, + mcp_app_resource_uri: None, + mcp_app_ui: None, + plugin_id: None, + read_only_hint: None, + result: Some(Box::new(McpToolCallResult { + content: Vec::new(), + structured_content: None, + meta: None, + })), + error: None, + duration_ms: Some(1_000), + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!( + started, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemStarted(ItemStartedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::McpToolCall(McpToolCallItem { + server: "server_a".to_string(), + tool: "tool_x".to_string(), + arguments: json!({ "key": "value" }), + result: None, + error: None, + status: McpToolCallStatus::InProgress, + }), + }, + })], + status: CodexStatus::Running, + } + ); + assert_eq!( + completed, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::McpToolCall(McpToolCallItem { + server: "server_a".to_string(), + tool: "tool_x".to_string(), + arguments: json!({ "key": "value" }), + result: Some(McpToolCallItemResult { + content: Vec::new(), + meta: None, + structured_content: None, + }), + error: None, + status: McpToolCallStatus::Completed, + }), + }, + })], + status: CodexStatus::Running, + } + ); +} + +#[test] +fn mcp_tool_call_failure_sets_failed_status() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let collected = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::McpToolCall { + id: "mcp-2".to_string(), + server: "server_b".to_string(), + tool: "tool_y".to_string(), + status: ApiMcpToolCallStatus::Failed, + arguments: json!({ "param": 42 }), + app_context: None, + mcp_app_resource_uri: None, + mcp_app_ui: None, + plugin_id: None, + read_only_hint: None, + result: None, + error: Some(McpToolCallError { + message: "tool exploded".to_string(), + }), + duration_ms: Some(5), + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!( + collected, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::McpToolCall(McpToolCallItem { + server: "server_b".to_string(), + tool: "tool_y".to_string(), + arguments: json!({ "param": 42 }), + result: None, + error: Some(McpToolCallItemError { + message: "tool exploded".to_string(), + }), + status: McpToolCallStatus::Failed, + }), + }, + })], + status: CodexStatus::Running, + } + ); +} + +#[test] +fn mcp_tool_call_defaults_arguments_and_preserves_structured_content() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let started = + processor.collect_thread_events(ServerNotification::ItemStarted(ItemStartedNotification { + item: ThreadItem::McpToolCall { + id: "mcp-3".to_string(), + server: "server_c".to_string(), + tool: "tool_z".to_string(), + status: ApiMcpToolCallStatus::InProgress, + arguments: serde_json::Value::Null, + app_context: None, + mcp_app_resource_uri: None, + mcp_app_ui: None, + plugin_id: None, + read_only_hint: None, + result: None, + error: None, + duration_ms: None, + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + started_at_ms: 0, + })); + let completed = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::McpToolCall { + id: "mcp-3".to_string(), + server: "server_c".to_string(), + tool: "tool_z".to_string(), + status: ApiMcpToolCallStatus::Completed, + arguments: serde_json::Value::Null, + app_context: None, + mcp_app_resource_uri: None, + mcp_app_ui: None, + plugin_id: None, + read_only_hint: None, + result: Some(Box::new(McpToolCallResult { + content: vec![json!({ + "type": "text", + "text": "done", + })], + structured_content: Some(json!({ "status": "ok" })), + meta: None, + })), + error: None, + duration_ms: Some(10), + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!( + started, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemStarted(ItemStartedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::McpToolCall(McpToolCallItem { + server: "server_c".to_string(), + tool: "tool_z".to_string(), + arguments: serde_json::Value::Null, + result: None, + error: None, + status: McpToolCallStatus::InProgress, + }), + }, + })], + status: CodexStatus::Running, + } + ); + assert_eq!( + completed, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::McpToolCall(McpToolCallItem { + server: "server_c".to_string(), + tool: "tool_z".to_string(), + arguments: serde_json::Value::Null, + result: Some(McpToolCallItemResult { + content: vec![json!({ + "type": "text", + "text": "done", + })], + meta: None, + structured_content: Some(json!({ "status": "ok" })), + }), + error: None, + status: McpToolCallStatus::Completed, + }), + }, + })], + status: CodexStatus::Running, + } + ); +} + +#[test] +fn collab_spawn_begin_and_end_emit_item_events() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let started = + processor.collect_thread_events(ServerNotification::ItemStarted(ItemStartedNotification { + item: ThreadItem::CollabAgentToolCall { + id: "collab-1".to_string(), + tool: CollabAgentTool::SpawnAgent, + status: ApiCollabAgentToolCallStatus::InProgress, + sender_thread_id: "thread-parent".to_string(), + receiver_thread_ids: Vec::new(), + prompt: Some("draft a plan".to_string()), + model: Some("gpt-5".to_string()), + reasoning_effort: None, + agents_states: std::collections::HashMap::new(), + }, + thread_id: "thread-parent".to_string(), + turn_id: "turn-1".to_string(), + started_at_ms: 0, + })); + let completed = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::CollabAgentToolCall { + id: "collab-1".to_string(), + tool: CollabAgentTool::SpawnAgent, + status: ApiCollabAgentToolCallStatus::Completed, + sender_thread_id: "thread-parent".to_string(), + receiver_thread_ids: vec!["thread-child".to_string()], + prompt: Some("draft a plan".to_string()), + model: Some("gpt-5".to_string()), + reasoning_effort: None, + agents_states: std::collections::HashMap::from([( + "thread-child".to_string(), + ApiCollabAgentState { + status: ApiCollabAgentStatus::Running, + message: None, + }, + )]), + }, + thread_id: "thread-parent".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!( + started, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemStarted(ItemStartedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::CollabToolCall(CollabToolCallItem { + tool: CollabTool::SpawnAgent, + sender_thread_id: "thread-parent".to_string(), + receiver_thread_ids: Vec::new(), + prompt: Some("draft a plan".to_string()), + agents_states: std::collections::HashMap::new(), + status: CollabToolCallStatus::InProgress, + },), + }, + })], + status: CodexStatus::Running, + } + ); + assert_eq!( + completed, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::CollabToolCall(CollabToolCallItem { + tool: CollabTool::SpawnAgent, + sender_thread_id: "thread-parent".to_string(), + receiver_thread_ids: vec!["thread-child".to_string()], + prompt: Some("draft a plan".to_string()), + agents_states: std::collections::HashMap::from([( + "thread-child".to_string(), + CollabAgentState { + status: CollabAgentStatus::Running, + message: None, + }, + )]), + status: CollabToolCallStatus::Completed, + },), + }, + })], + status: CodexStatus::Running, + } + ); +} + +#[test] +fn file_change_completion_maps_change_kinds() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let collected = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::FileChange { + id: "patch-1".to_string(), + changes: vec![ + ApiFileUpdateChange { + path: "a/added.txt".to_string(), + kind: ApiPatchChangeKind::Add, + diff: String::new(), + }, + ApiFileUpdateChange { + path: "b/deleted.txt".to_string(), + kind: ApiPatchChangeKind::Delete, + diff: String::new(), + }, + ApiFileUpdateChange { + path: "c/modified.txt".to_string(), + kind: ApiPatchChangeKind::Update { move_path: None }, + diff: "@@ -1 +1 @@".to_string(), + }, + ], + status: ApiPatchApplyStatus::Completed, + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!( + collected, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::FileChange(FileChangeItem { + changes: vec![ + ExecFileUpdateChange { + path: "a/added.txt".to_string(), + kind: PatchChangeKind::Add, + }, + ExecFileUpdateChange { + path: "b/deleted.txt".to_string(), + kind: PatchChangeKind::Delete, + }, + ExecFileUpdateChange { + path: "c/modified.txt".to_string(), + kind: PatchChangeKind::Update, + }, + ], + status: PatchApplyStatus::Completed, + }), + }, + })], + status: CodexStatus::Running, + } + ); +} + +#[test] +fn file_change_declined_maps_to_failed_status() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let collected = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::FileChange { + id: "patch-2".to_string(), + changes: vec![ApiFileUpdateChange { + path: "file.txt".to_string(), + kind: ApiPatchChangeKind::Update { move_path: None }, + diff: "@@ -1 +1 @@".to_string(), + }], + status: ApiPatchApplyStatus::Declined, + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!( + collected, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::FileChange(FileChangeItem { + changes: vec![ExecFileUpdateChange { + path: "file.txt".to_string(), + kind: PatchChangeKind::Update, + }], + status: PatchApplyStatus::Failed, + }), + }, + })], + status: CodexStatus::Running, + } + ); +} + +#[test] +fn agent_message_item_updates_final_message() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let collected = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::AgentMessage { + id: "msg-1".to_string(), + text: "hello".to_string(), + phase: None, + memory_citation: None, + delivery: None, + questions: None, + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!( + collected, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::AgentMessage(AgentMessageItem { + text: "hello".to_string(), + }), + }, + })], + status: CodexStatus::Running, + } + ); + assert_eq!(processor.final_message(), Some("hello")); +} + +#[test] +fn agent_message_item_started_is_ignored() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let collected = + processor.collect_thread_events(ServerNotification::ItemStarted(ItemStartedNotification { + item: ThreadItem::AgentMessage { + id: "msg-1".to_string(), + text: "hello".to_string(), + phase: None, + memory_citation: None, + delivery: None, + questions: None, + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + started_at_ms: 0, + })); + + assert_eq!( + collected, + CollectedThreadEvents { + events: Vec::new(), + status: CodexStatus::Running, + } + ); +} + +#[test] +fn reasoning_item_completed_uses_synthetic_id() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let collected = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::Reasoning { + id: "rs-1".to_string(), + summary: vec!["thinking...".to_string()], + content: vec!["raw".to_string()], + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!( + collected, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::Reasoning(ReasoningItem { + text: "thinking...".to_string(), + }), + }, + })], + status: CodexStatus::Running, + } + ); +} + +#[test] +fn warning_event_produces_error_item() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let collected = processor.collect_warning( + "Heads up: Long conversations and multiple compactions can cause the model to be less accurate. Start a new conversation when possible to keep conversations small and targeted.".to_string(), + ); + + assert_eq!( + collected, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::Error(ErrorItem { + message: "Heads up: Long conversations and multiple compactions can cause the model to be less accurate. Start a new conversation when possible to keep conversations small and targeted.".to_string(), + }), + }, + })], + status: CodexStatus::Running, + } + ); +} + +#[test] +fn plan_update_emits_started_then_updated_then_completed() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let started = processor.collect_thread_events(ServerNotification::TurnPlanUpdated( + TurnPlanUpdatedNotification { + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + explanation: None, + plan: vec![ + TurnPlanStep { + step: "step one".to_string(), + status: TurnPlanStepStatus::Pending, + }, + TurnPlanStep { + step: "step two".to_string(), + status: TurnPlanStepStatus::InProgress, + }, + ], + }, + )); + assert_eq!( + started, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemStarted(ItemStartedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::TodoList(TodoListItem { + items: vec![ + TodoItem { + text: "step one".to_string(), + completed: false, + }, + TodoItem { + text: "step two".to_string(), + completed: false, + }, + ], + }), + }, + })], + status: CodexStatus::Running, + } + ); + + let updated = processor.collect_thread_events(ServerNotification::TurnPlanUpdated( + TurnPlanUpdatedNotification { + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + explanation: None, + plan: vec![ + TurnPlanStep { + step: "step one".to_string(), + status: TurnPlanStepStatus::Completed, + }, + TurnPlanStep { + step: "step two".to_string(), + status: TurnPlanStepStatus::InProgress, + }, + ], + }, + )); + assert_eq!( + updated, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemUpdated(ItemUpdatedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::TodoList(TodoListItem { + items: vec![ + TodoItem { + text: "step one".to_string(), + completed: true, + }, + TodoItem { + text: "step two".to_string(), + completed: false, + }, + ], + }), + }, + })], + status: CodexStatus::Running, + } + ); + + let completed = processor.collect_thread_events(ServerNotification::TurnCompleted( + TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: Vec::new(), + status: TurnStatus::Completed, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + }, + }, + )); + assert_eq!( + completed, + CollectedThreadEvents { + events: vec![ + ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::TodoList(TodoListItem { + items: vec![ + TodoItem { + text: "step one".to_string(), + completed: true, + }, + TodoItem { + text: "step two".to_string(), + completed: false, + }, + ], + }), + }, + }), + ThreadEvent::TurnCompleted(TurnCompletedEvent { + usage: Usage::default(), + }), + ], + status: CodexStatus::InitiateShutdown, + } + ); +} + +#[test] +fn plan_update_after_completion_starts_new_todo_list_with_new_id() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let _ = processor.collect_thread_events(ServerNotification::TurnPlanUpdated( + TurnPlanUpdatedNotification { + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + explanation: None, + plan: vec![TurnPlanStep { + step: "only".to_string(), + status: TurnPlanStepStatus::Pending, + }], + }, + )); + let _ = processor.collect_thread_events(ServerNotification::TurnCompleted( + TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: Vec::new(), + status: TurnStatus::Completed, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + }, + }, + )); + + let restarted = processor.collect_thread_events(ServerNotification::TurnPlanUpdated( + TurnPlanUpdatedNotification { + thread_id: "thread-1".to_string(), + turn_id: "turn-2".to_string(), + explanation: None, + plan: vec![TurnPlanStep { + step: "again".to_string(), + status: TurnPlanStepStatus::Pending, + }], + }, + )); + + assert_eq!( + restarted, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemStarted(ItemStartedEvent { + item: ExecThreadItem { + id: "item_1".to_string(), + details: ThreadItemDetails::TodoList(TodoListItem { + items: vec![TodoItem { + text: "again".to_string(), + completed: false, + }], + }), + }, + })], + status: CodexStatus::Running, + } + ); +} + +#[test] +fn token_usage_update_is_emitted_on_turn_completion() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let usage_update = + processor.collect_thread_events(ServerNotification::ThreadTokenUsageUpdated( + codex_app_server_protocol::ThreadTokenUsageUpdatedNotification { + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + token_usage: ThreadTokenUsage { + total: TokenUsageBreakdown { + total_tokens: 42, + input_tokens: 10, + cached_input_tokens: 3, + cache_write_input_tokens: 4, + output_tokens: 29, + reasoning_output_tokens: 7, + }, + last: TokenUsageBreakdown { + total_tokens: 42, + input_tokens: 10, + cached_input_tokens: 3, + cache_write_input_tokens: 4, + output_tokens: 29, + reasoning_output_tokens: 7, + }, + model_context_window: Some(128_000), + }, + }, + )); + assert_eq!( + usage_update, + CollectedThreadEvents { + events: Vec::new(), + status: CodexStatus::Running, + } + ); + + let completed = processor.collect_thread_events(ServerNotification::TurnCompleted( + TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: Vec::new(), + status: TurnStatus::Completed, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + }, + }, + )); + assert_eq!( + completed, + CollectedThreadEvents { + events: vec![ThreadEvent::TurnCompleted(TurnCompletedEvent { + usage: Usage { + input_tokens: 10, + cached_input_tokens: 3, + cache_write_input_tokens: 4, + output_tokens: 29, + reasoning_output_tokens: 7, + }, + })], + status: CodexStatus::InitiateShutdown, + } + ); +} + +#[test] +fn turn_completion_recovers_final_message_from_turn_items() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let completed = processor.collect_thread_events(ServerNotification::TurnCompleted( + TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: vec![ThreadItem::AgentMessage { + id: "msg-1".to_string(), + text: "final answer".to_string(), + phase: None, + memory_citation: None, + delivery: None, + questions: None, + }], + status: TurnStatus::Completed, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + }, + }, + )); + + assert_eq!( + completed, + CollectedThreadEvents { + events: vec![ThreadEvent::TurnCompleted(TurnCompletedEvent { + usage: Usage::default(), + })], + status: CodexStatus::InitiateShutdown, + } + ); + assert_eq!(processor.final_message(), Some("final answer")); +} + +#[test] +fn turn_completion_reconciles_started_items_from_turn_items() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let started = + processor.collect_thread_events(ServerNotification::ItemStarted(ItemStartedNotification { + item: ThreadItem::CommandExecution { + model_context: None, + id: "cmd-1".to_string(), + command: "ls".to_string(), + cwd: test_path_buf("/tmp/project").abs().into(), + process_id: Some("123".to_string()), + plugin_id: None, + script_path: None, + source: CommandExecutionSource::UserShell, + status: ApiCommandExecutionStatus::InProgress, + command_actions: Vec::::new(), + aggregated_output: None, + exit_code: None, + duration_ms: None, + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + started_at_ms: 0, + })); + assert_eq!( + started, + CollectedThreadEvents { + events: vec![ThreadEvent::ItemStarted(ItemStartedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::CommandExecution(CommandExecutionItem { + command: "ls".to_string(), + aggregated_output: String::new(), + exit_code: None, + status: CommandExecutionStatus::InProgress, + }), + }, + })], + status: CodexStatus::Running, + } + ); + + let completed = processor.collect_thread_events(ServerNotification::TurnCompleted( + TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: vec![ThreadItem::CommandExecution { + model_context: None, + id: "cmd-1".to_string(), + command: "ls".to_string(), + cwd: test_path_buf("/tmp/project").abs().into(), + process_id: Some("123".to_string()), + plugin_id: None, + script_path: None, + source: CommandExecutionSource::UserShell, + status: ApiCommandExecutionStatus::Completed, + command_actions: Vec::::new(), + aggregated_output: Some("a.txt\n".to_string()), + exit_code: Some(0), + duration_ms: Some(3), + }], + status: TurnStatus::Completed, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + }, + }, + )); + + assert_eq!( + completed, + CollectedThreadEvents { + events: vec![ + ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: ExecThreadItem { + id: "item_0".to_string(), + details: ThreadItemDetails::CommandExecution(CommandExecutionItem { + command: "ls".to_string(), + aggregated_output: "a.txt\n".to_string(), + exit_code: Some(0), + status: CommandExecutionStatus::Completed, + }), + }, + }), + ThreadEvent::TurnCompleted(TurnCompletedEvent { + usage: Usage::default(), + }), + ], + status: CodexStatus::InitiateShutdown, + } + ); +} + +#[test] +fn turn_completion_overwrites_stale_final_message_from_turn_items() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + let _ = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::AgentMessage { + id: "msg-stale".to_string(), + text: "stale answer".to_string(), + phase: None, + memory_citation: None, + delivery: None, + questions: None, + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + let completed = processor.collect_thread_events(ServerNotification::TurnCompleted( + TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: vec![ThreadItem::AgentMessage { + id: "msg-1".to_string(), + text: "final answer".to_string(), + phase: None, + memory_citation: None, + delivery: None, + questions: None, + }], + status: TurnStatus::Completed, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + }, + }, + )); + + assert_eq!( + completed, + CollectedThreadEvents { + events: vec![ThreadEvent::TurnCompleted(TurnCompletedEvent { + usage: Usage::default(), + })], + status: CodexStatus::InitiateShutdown, + } + ); + assert_eq!(processor.final_message(), Some("final answer")); +} + +#[test] +fn turn_completion_preserves_streamed_final_message_when_turn_items_are_empty() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + let _ = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::AgentMessage { + id: "msg-streamed".to_string(), + text: "streamed answer".to_string(), + phase: None, + memory_citation: None, + delivery: None, + questions: None, + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + let completed = processor.collect_thread_events(ServerNotification::TurnCompleted( + TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: Vec::new(), + status: TurnStatus::Completed, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + }, + }, + )); + + assert_eq!( + completed, + CollectedThreadEvents { + events: vec![ThreadEvent::TurnCompleted(TurnCompletedEvent { + usage: Usage::default(), + })], + status: CodexStatus::InitiateShutdown, + } + ); + assert_eq!(processor.final_message(), Some("streamed answer")); +} + +#[test] +fn failed_turn_clears_stale_final_message() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let collected = processor.collect_thread_events(ServerNotification::ItemCompleted( + ItemCompletedNotification { + item: ThreadItem::AgentMessage { + id: "msg-1".to_string(), + text: "partial answer".to_string(), + phase: None, + memory_citation: None, + delivery: None, + questions: None, + }, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + completed_at_ms: 0, + }, + )); + + assert_eq!(collected.status, CodexStatus::Running); + assert_eq!(processor.final_message(), Some("partial answer")); + + let collected = processor.collect_thread_events(ServerNotification::TurnCompleted( + TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: Vec::new(), + status: TurnStatus::Failed, + error: Some(TurnError { + misalignment: None, + message: "turn failed".to_string(), + additional_details: None, + codex_error_info: None, + }), + started_at: None, + completed_at: None, + duration_ms: None, + }, + }, + )); + + assert_eq!(collected.status, CodexStatus::InitiateShutdown); + assert_eq!(processor.final_message(), None); +} + +#[test] +fn turn_completion_falls_back_to_final_plan_text() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let completed = processor.collect_thread_events(ServerNotification::TurnCompleted( + TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: vec![ThreadItem::Plan { + id: "plan-1".to_string(), + text: "ship the typed adapter".to_string(), + }], + status: TurnStatus::Completed, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + }, + }, + )); + + assert_eq!( + completed, + CollectedThreadEvents { + events: vec![ThreadEvent::TurnCompleted(TurnCompletedEvent { + usage: Usage::default(), + })], + status: CodexStatus::InitiateShutdown, + } + ); + assert_eq!(processor.final_message(), Some("ship the typed adapter")); +} + +#[test] +fn turn_failure_prefers_structured_error_message() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let error = processor.collect_thread_events(ServerNotification::Error(ErrorNotification { + error: TurnError { + misalignment: None, + message: "backend failed".to_string(), + codex_error_info: None, + additional_details: Some("request id abc".to_string()), + }, + will_retry: false, + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + })); + assert_eq!( + error, + CollectedThreadEvents { + events: vec![ThreadEvent::Error(ThreadErrorEvent { + message: "backend failed (request id abc)".to_string(), + })], + status: CodexStatus::Running, + } + ); + + let failed = processor.collect_thread_events(ServerNotification::TurnCompleted( + TurnCompletedNotification { + thread_id: "thread-1".to_string(), + turn: Turn { + id: "turn-1".to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: Vec::new(), + status: TurnStatus::Failed, + error: None, + started_at: None, + completed_at: None, + duration_ms: None, + }, + }, + )); + assert_eq!( + failed, + CollectedThreadEvents { + events: vec![ThreadEvent::TurnFailed(TurnFailedEvent { + error: ThreadErrorEvent { + message: "backend failed (request id abc)".to_string(), + }, + })], + status: CodexStatus::InitiateShutdown, + } + ); +} + +#[test] +fn model_reroute_surfaces_as_error_item() { + let mut processor = EventProcessorWithJsonOutput::new(/*last_message_path*/ None); + + let collected = processor.collect_thread_events(ServerNotification::ModelRerouted( + codex_app_server_protocol::ModelReroutedNotification { + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + from_model: "gpt-5".to_string(), + to_model: "gpt-5-mini".to_string(), + reason: codex_app_server_protocol::ModelRerouteReason::HighRiskCyberActivity, + }, + )); + + assert_eq!(collected.status, CodexStatus::Running); + assert_eq!(collected.events.len(), 1); + let ThreadEvent::ItemCompleted(ItemCompletedEvent { item }) = &collected.events[0] else { + panic!("expected ItemCompleted"); + }; + assert_eq!(item.id, "item_0"); + assert_eq!( + item.details, + ThreadItemDetails::Error(ErrorItem { + message: "model rerouted: gpt-5 -> gpt-5-mini (HighRiskCyberActivity)".to_string(), + }) + ); +} diff --git a/codex-rs/exec/tests/fixtures/apply_patch_freeform_final.txt b/codex-rs/exec/tests/fixtures/apply_patch_freeform_final.txt new file mode 100644 index 0000000000000000000000000000000000000000..b01923e0e536015143ad914c46a5808daa647f4b --- /dev/null +++ b/codex-rs/exec/tests/fixtures/apply_patch_freeform_final.txt @@ -0,0 +1,4 @@ +class BaseClass: + def method(): + + return True diff --git a/codex-rs/exec/tests/suite/add_dir.rs b/codex-rs/exec/tests/suite/add_dir.rs new file mode 100644 index 0000000000000000000000000000000000000000..f95433d1fbcf341511e2703a3b00c1ab35236620 --- /dev/null +++ b/codex-rs/exec/tests/suite/add_dir.rs @@ -0,0 +1,72 @@ +#![cfg(not(target_os = "windows"))] +#![allow(clippy::unwrap_used)] + +use core_test_support::responses; +use core_test_support::test_codex_exec::test_codex_exec; + +/// Verify that the --add-dir flag is accepted and the command runs successfully. +/// This test confirms the CLI argument is properly wired up. +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn accepts_add_dir_flag() -> anyhow::Result<()> { + let test = test_codex_exec(); + + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("response_1"), + responses::ev_assistant_message("response_1", "Task completed"), + responses::ev_completed("response_1"), + ]); + responses::mount_sse_once(&server, body).await; + + // Create temporary directories to use with --add-dir + let temp_dir1 = tempfile::tempdir()?; + let temp_dir2 = tempfile::tempdir()?; + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("--sandbox") + .arg("workspace-write") + .arg("--add-dir") + .arg(temp_dir1.path()) + .arg("--add-dir") + .arg(temp_dir2.path()) + .arg("test with additional directories") + .assert() + .code(0); + + Ok(()) +} + +/// Verify that multiple --add-dir flags can be specified. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn accepts_multiple_add_dir_flags() -> anyhow::Result<()> { + let test = test_codex_exec(); + + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("response_1"), + responses::ev_assistant_message("response_1", "Multiple directories accepted"), + responses::ev_completed("response_1"), + ]); + responses::mount_sse_once(&server, body).await; + + let temp_dir1 = tempfile::tempdir()?; + let temp_dir2 = tempfile::tempdir()?; + let temp_dir3 = tempfile::tempdir()?; + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("--sandbox") + .arg("workspace-write") + .arg("--add-dir") + .arg(temp_dir1.path()) + .arg("--add-dir") + .arg(temp_dir2.path()) + .arg("--add-dir") + .arg(temp_dir3.path()) + .arg("test with three directories") + .assert() + .code(0); + + Ok(()) +} diff --git a/codex-rs/exec/tests/suite/agents_md.rs b/codex-rs/exec/tests/suite/agents_md.rs new file mode 100644 index 0000000000000000000000000000000000000000..3cc47f095e86945930d7e319885521f38cf70bea --- /dev/null +++ b/codex-rs/exec/tests/suite/agents_md.rs @@ -0,0 +1,74 @@ +#![allow(clippy::unwrap_used)] + +use core_test_support::responses; +use core_test_support::test_codex_exec::test_codex_exec; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_includes_workspace_agents_md_in_request() -> anyhow::Result<()> { + let test = test_codex_exec(); + std::fs::write(test.cwd_path().join("AGENTS.md"), "workspace instructions")?; + + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("resp1"), + responses::ev_assistant_message("m1", "fixture hello"), + responses::ev_completed("resp1"), + ]); + let response_mock = responses::mount_sse_once(&server, body).await; + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("tell me something") + .assert() + .success(); + + let user_messages = response_mock.single_request().message_input_texts("user"); + assert!( + user_messages + .iter() + .any(|text| text.contains("workspace instructions")), + "request should include workspace AGENTS.md instructions: {user_messages:?}" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_prefers_workspace_agents_override_md() -> anyhow::Result<()> { + let test = test_codex_exec(); + std::fs::write(test.cwd_path().join("AGENTS.md"), "base instructions")?; + std::fs::write( + test.cwd_path().join("AGENTS.override.md"), + "override instructions", + )?; + + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("resp1"), + responses::ev_assistant_message("m1", "fixture hello"), + responses::ev_completed("resp1"), + ]); + let response_mock = responses::mount_sse_once(&server, body).await; + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("tell me something") + .assert() + .success(); + + let user_messages = response_mock.single_request().message_input_texts("user"); + assert!( + user_messages + .iter() + .any(|text| text.contains("override instructions")), + "request should include AGENTS.override.md instructions: {user_messages:?}" + ); + assert!( + user_messages + .iter() + .all(|text| !text.contains("base instructions")), + "request should exclude shadowed AGENTS.md instructions: {user_messages:?}" + ); + + Ok(()) +} diff --git a/codex-rs/exec/tests/suite/apply_patch.rs b/codex-rs/exec/tests/suite/apply_patch.rs new file mode 100644 index 0000000000000000000000000000000000000000..8376b2e76504be8819848ca2785d21b66550b2bd --- /dev/null +++ b/codex-rs/exec/tests/suite/apply_patch.rs @@ -0,0 +1,313 @@ +#![allow(clippy::unwrap_used, unused_imports)] + +use anyhow::Context; +use assert_cmd::prelude::*; +use codex_apply_patch::CODEX_CORE_APPLY_PATCH_ARG1; +use core_test_support::responses::ev_apply_patch_custom_tool_call; +use core_test_support::responses::ev_completed; +use core_test_support::responses::mount_sse_sequence; +use core_test_support::responses::sse; +use core_test_support::responses::start_mock_server; +use std::fs; +use std::process::Command; +use tempfile::tempdir; + +/// While we may add an `apply-patch` subcommand to the `codex` CLI multitool +/// at some point, we must ensure that the smaller `codex-exec` CLI can still +/// emulate the `apply_patch` CLI. +#[test] +fn test_standalone_exec_cli_can_use_apply_patch() -> anyhow::Result<()> { + let tmp = tempdir()?; + let relative_path = "source.txt"; + let absolute_path = tmp.path().join(relative_path); + fs::write(&absolute_path, "original content\n")?; + + Command::new(codex_utils_cargo_bin::cargo_bin("codex-exec")?) + .arg(CODEX_CORE_APPLY_PATCH_ARG1) + .arg( + r#"*** Begin Patch +*** Update File: source.txt +@@ +-original content ++modified by apply_patch +*** End Patch"#, + ) + .current_dir(tmp.path()) + .assert() + .success() + .stdout("Success. Updated the following files:\nM source.txt\n") + .stderr(predicates::str::is_empty()); + assert_eq!( + fs::read_to_string(absolute_path)?, + "modified by apply_patch\n" + ); + Ok(()) +} + +#[cfg(not(target_os = "windows"))] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_apply_patch_tool() -> anyhow::Result<()> { + use core_test_support::skip_if_no_network; + use core_test_support::test_codex_exec::test_codex_exec; + + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let tmp_path = test.cwd_path().to_path_buf(); + let add_patch = r#"*** Begin Patch +*** Add File: test.md ++Hello world +*** End Patch"#; + let update_patch = r#"*** Begin Patch +*** Update File: test.md +@@ +-Hello world ++Final text +*** End Patch"#; + let response_streams = vec![ + sse(vec![ + ev_apply_patch_custom_tool_call("request_0", add_patch), + ev_completed("request_0"), + ]), + sse(vec![ + ev_apply_patch_custom_tool_call("request_1", update_patch), + ev_completed("request_1"), + ]), + sse(vec![ev_completed("request_2")]), + ]; + let server = start_mock_server().await; + mount_sse_sequence(&server, response_streams).await; + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-s") + .arg("danger-full-access") + .arg("foo") + .assert() + .success(); + + let final_path = tmp_path.join("test.md"); + let contents = std::fs::read_to_string(&final_path).expect("final file should be readable"); + assert_eq!(contents, "Final text\n"); + Ok(()) +} + +#[cfg(not(target_os = "windows"))] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_apply_patch_freeform_tool() -> anyhow::Result<()> { + use core_test_support::skip_if_no_network; + use core_test_support::test_codex_exec::test_codex_exec; + + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let freeform_add_patch = r#"*** Begin Patch +*** Add File: app.py ++class BaseClass: ++ def method(): ++ return False +*** End Patch"#; + let freeform_update_patch = r#"*** Begin Patch +*** Update File: app.py +@@ def method(): +- return False ++ ++ return True +*** End Patch"#; + let response_streams = vec![ + sse(vec![ + ev_apply_patch_custom_tool_call("request_0", freeform_add_patch), + ev_completed("request_0"), + ]), + sse(vec![ + ev_apply_patch_custom_tool_call("request_1", freeform_update_patch), + ev_completed("request_1"), + ]), + sse(vec![ev_completed("request_2")]), + ]; + let server = start_mock_server().await; + mount_sse_sequence(&server, response_streams).await; + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-s") + .arg("danger-full-access") + .arg("foo") + .assert() + .success(); + + // Verify final file contents + let final_path = test.cwd_path().join("app.py"); + let contents = std::fs::read_to_string(&final_path).expect("final file should be readable"); + assert_eq!( + contents, + include_str!("../fixtures/apply_patch_freeform_final.txt") + ); + Ok(()) +} + +#[cfg(not(target_os = "windows"))] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn shutdown_flushes_completed_turn_and_file_diff() -> anyhow::Result<()> { + use codex_login::AuthDotJson; + use codex_login::AuthKeyringBackendKind; + use codex_login::TokenData; + use codex_login::save_auth; + use codex_login::token_data::IdTokenInfo; + use codex_protocol::auth::AuthMode; + use core_test_support::skip_if_no_network; + use core_test_support::test_codex_exec::test_codex_exec; + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use std::time::Duration; + use wiremock::Mock; + use wiremock::ResponseTemplate; + use wiremock::matchers::method; + + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let run_dir = test.cwd_path(); + assert!( + std::process::Command::new("git") + .args(["init", "--quiet"]) + .current_dir(run_dir) + .status() + .context("initialize test git repository")? + .success() + ); + assert!( + std::process::Command::new("git") + .args([ + "remote", + "add", + "origin", + "https://github.com/test-org/turn-profile-test.git", + ]) + .current_dir(run_dir) + .status() + .context("add test git remote")? + .success() + ); + save_auth( + test.home_path(), + &AuthDotJson { + auth_mode: Some(AuthMode::Chatgpt), + openai_api_key: None, + tokens: Some(TokenData { + id_token: IdTokenInfo { + raw_jwt: "header.eyJhbGciOiJub25lIn0.eyJzdWIiOiJ1c2VyLTEyMyJ9.".to_string(), + ..Default::default() + }, + access_token: "test-access-token".to_string(), + refresh_token: "test-refresh-token".to_string(), + account_id: None, + }), + last_refresh: None, + agent_identity: None, + personal_access_token: None, + bedrock_api_key: None, + bedrock_access_keys: None, + }, + codex_login::AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + + let server = start_mock_server().await; + let deliveries = Arc::new(AtomicUsize::new(0)); + let completed_deliveries = Arc::clone(&deliveries); + Mock::given(method("POST")) + .and(|request: &wiremock::Request| request.url.path().ends_with("events")) + .respond_with(move |_request: &wiremock::Request| { + // A request is counted only once its delayed handler has completed. + std::thread::sleep(Duration::from_millis(750)); + completed_deliveries.fetch_add(1, Ordering::Release); + ResponseTemplate::new(200) + }) + .expect(3..) + .mount(&server) + .await; + mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_apply_patch_custom_tool_call( + "request_0", + "*** Begin Patch\n*** Add File: flushed.md\n+persist this line\n*** End Patch", + ), + ev_completed("request_0"), + ]), + sse(vec![ev_completed("request_1")]), + ], + ) + .await; + + test.cmd_with_server(&server) + .env_remove("CODEX_API_KEY") + .env_remove("CODEX_ANALYTICS_EVENTS_CAPTURE_FILE") + .arg("--skip-git-repo-check") + .arg("-s") + .arg("danger-full-access") + .arg("-c") + .arg("model_provider=\"test\"") + .arg("-c") + .arg(format!( + "model_providers.test={{name=\"test\",base_url={:?},wire_api=\"responses\",requires_openai_auth=false,supports_websockets=false}}", + format!("{}/v1", server.uri()) + )) + .arg("-c") + .arg("analytics.enabled=true") + .arg("-c") + .arg(format!("chatgpt_base_url={}", server.uri())) + .arg("write the file") + .assert() + .success(); + + let persisted = std::fs::read_to_string(run_dir.join("flushed.md")) + .context("read file written by apply_patch")?; + assert_eq!(persisted, "persist this line\n"); + + let requests = server + .received_requests() + .await + .context("read analytics requests")?; + let request_paths = requests + .iter() + .map(|request| request.url.path().to_string()) + .collect::>(); + let event_types = requests + .iter() + .filter(|request| request.url.path().ends_with("events")) + .flat_map(|request| { + serde_json::from_slice::(&request.body) + .ok() + .and_then(|body| { + body.get("events") + .and_then(serde_json::Value::as_array) + .cloned() + }) + .unwrap_or_default() + .into_iter() + .filter_map(|event| event["event_type"].as_str().map(str::to_owned)) + .collect::>() + }) + .collect::>(); + assert!( + deliveries.load(Ordering::Acquire) >= 3, + "the runtime exited before the delayed analytics deliveries completed (completed={}, events={event_types:?}, paths={request_paths:?})", + deliveries.load(Ordering::Acquire) + ); + assert!( + event_types.iter().any(|event| event == "codex_turn_event"), + "no completed event was delivered; request paths: {request_paths:?}" + ); + assert!( + event_types + .iter() + .any(|event| event == "codex_accepted_line_fingerprints"), + "no accepted-line event was delivered; request paths: {request_paths:?}" + ); + + Ok(()) +} diff --git a/codex-rs/exec/tests/suite/approval_policy.rs b/codex-rs/exec/tests/suite/approval_policy.rs new file mode 100644 index 0000000000000000000000000000000000000000..5efb6d0bdad5a2bd9fdf4f5f3e14587001ad0b16 --- /dev/null +++ b/codex-rs/exec/tests/suite/approval_policy.rs @@ -0,0 +1,78 @@ +#![cfg(not(target_os = "windows"))] +#![allow(clippy::unwrap_used)] + +use core_test_support::responses; +use core_test_support::test_codex_exec::test_codex_exec; + +async fn run_exec_with_config(config_toml: &str, extra_args: &[&str]) -> anyhow::Result { + let test = test_codex_exec(); + std::fs::write(test.home_path().join("config.toml"), config_toml)?; + + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("response_1"), + responses::ev_assistant_message("response_1", "done"), + responses::ev_completed("response_1"), + ]); + responses::mount_sse_once(&server, body).await; + + let mut cmd = test.cmd_with_server(&server); + let output = cmd + .arg("--skip-git-repo-check") + .args(extra_args) + .arg("check approval mode") + .output()?; + + assert!(output.status.success(), "exec run failed: {output:?}"); + + Ok(String::from_utf8(output.stderr)?) +} + +async fn run_exec_with_auto_review_config(extra_args: &[&str]) -> anyhow::Result { + run_exec_with_config( + r#" +approval_policy = "on-request" +approvals_reviewer = "auto_review" +"#, + extra_args, + ) + .await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_preserves_on_request_for_auto_review_config() -> anyhow::Result<()> { + let stderr = run_exec_with_auto_review_config(&[]).await?; + assert!( + stderr.contains("approval: on-request"), + "stderr missing preserved auto-review approval mode: {stderr}" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_approve_for_me_flag_sets_approval_mode_and_sandbox() -> anyhow::Result<()> { + let stderr = run_exec_with_config("", &["--approve-for-me"]).await?; + assert!( + stderr.contains("approval: on-request"), + "stderr missing --approve-for-me approval mode: {stderr}" + ); + assert!( + stderr.contains("sandbox: workspace-write"), + "stderr missing --approve-for-me sandbox mode: {stderr}" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_bypass_preserves_never_for_auto_review_config() -> anyhow::Result<()> { + let stderr = + run_exec_with_auto_review_config(&["--dangerously-bypass-approvals-and-sandbox"]).await?; + assert!( + stderr.contains("approval: never"), + "stderr missing bypass approval mode: {stderr}" + ); + + Ok(()) +} diff --git a/codex-rs/exec/tests/suite/auth_env.rs b/codex-rs/exec/tests/suite/auth_env.rs new file mode 100644 index 0000000000000000000000000000000000000000..179f40fbd40dd2f2a7a88b1736f19335c382a387 --- /dev/null +++ b/codex-rs/exec/tests/suite/auth_env.rs @@ -0,0 +1,31 @@ +#![allow(clippy::unwrap_used)] +use core_test_support::responses::ev_completed; +use core_test_support::responses::mount_sse_once_match; +use core_test_support::responses::sse; +use core_test_support::responses::start_mock_server; +use core_test_support::test_codex_exec::test_codex_exec; +use wiremock::matchers::header; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_uses_codex_api_key_env_var() -> anyhow::Result<()> { + let test = test_codex_exec(); + let server = start_mock_server().await; + let repo_root = codex_utils_cargo_bin::repo_root()?; + + mount_sse_once_match( + &server, + header("Authorization", "Bearer dummy"), + sse(vec![ev_completed("request_0")]), + ) + .await; + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg("echo testing codex api key") + .assert() + .success(); + + Ok(()) +} diff --git a/codex-rs/exec/tests/suite/completion_backfill_tests.rs b/codex-rs/exec/tests/suite/completion_backfill_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..939fb1c69505a87804d0d2434fcbb37724a893bc --- /dev/null +++ b/codex-rs/exec/tests/suite/completion_backfill_tests.rs @@ -0,0 +1,162 @@ +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex_exec::test_codex_exec; +use pretty_assertions::assert_eq; +use serde_json::json; + +const PARENT_PROMPT: &str = "spawn a child and wait for it"; +const CHILD_PROMPT: &str = "child: finish first"; +const SPAWN_CALL_ID: &str = "spawn-call"; +const WAIT_CALL_ID: &str = "wait-call"; + +fn body_contains(request: &wiremock::Request, text: &str) -> bool { + std::str::from_utf8(&request.body).is_ok_and(|body| body.contains(text)) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn ignores_unrelated_turn_completion_before_backfilling_primary_turn() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let spawn_args = json!({ + "message": CHILD_PROMPT, + "task_name": "worker", + "fork_turns": "none", + }) + .to_string(); + let _parent_turn = responses::mount_sse_once_match( + &server, + |request: &wiremock::Request| body_contains(request, PARENT_PROMPT), + responses::sse(vec![ + responses::ev_response_created("resp-parent-1"), + responses::ev_function_call_with_namespace( + SPAWN_CALL_ID, + "collaboration", + "spawn_agent", + &spawn_args, + ), + responses::ev_completed("resp-parent-1"), + ]), + ) + .await; + let _child_turn = responses::mount_sse_once_match( + &server, + |request: &wiremock::Request| { + body_contains(request, CHILD_PROMPT) && !body_contains(request, SPAWN_CALL_ID) + }, + responses::sse(vec![ + responses::ev_response_created("resp-child"), + responses::ev_assistant_message("msg-child", "child done"), + responses::ev_completed("resp-child"), + ]), + ) + .await; + let _parent_wait = responses::mount_sse_once_match( + &server, + |request: &wiremock::Request| { + body_contains(request, SPAWN_CALL_ID) && !body_contains(request, WAIT_CALL_ID) + }, + responses::sse(vec![ + responses::ev_response_created("resp-parent-2"), + responses::ev_function_call_with_namespace( + WAIT_CALL_ID, + "collaboration", + "wait_agent", + "{}", + ), + responses::ev_completed("resp-parent-2"), + ]), + ) + .await; + let _parent_completion = responses::mount_sse_once_match( + &server, + |request: &wiremock::Request| body_contains(request, WAIT_CALL_ID), + responses::sse(vec![ + responses::ev_response_created("resp-parent-3"), + responses::ev_assistant_message("msg-parent", "parent done"), + responses::ev_completed("resp-parent-3"), + ]), + ) + .await; + + let test = test_codex_exec(); + let mock_provider = format!( + "model_providers.mock_provider={{name=\"Mock provider for test\",base_url=\"{}/v1\",wire_api=\"responses\",supports_websockets=false}}", + server.uri() + ); + let output = test + .cmd() + .env( + "RUST_LOG", + "codex_app_server::message_processor=trace,codex_app_server::outgoing_message=trace", + ) + .arg("--skip-git-repo-check") + .arg("--json") + .arg("-c") + .arg(mock_provider) + .args([ + "-c", + "model_provider=\"mock_provider\"", + "-c", + "features.multi_agent=true", + "-c", + "features.multi_agent_v2=true", + "-c", + "features.enable_request_compression=false", + ]) + .arg(PARENT_PROMPT) + .output()?; + assert!(output.status.success(), "exec run failed: {output:?}"); + + let stdout = String::from_utf8(output.stdout)?; + assert!( + stdout.contains("parent done"), + "primary completion was not processed: {stdout}" + ); + let stderr = String::from_utf8(output.stderr)?; + let lines = stderr.lines().collect::>(); + let turn_completions = lines + .iter() + .enumerate() + .filter_map(|(index, line)| { + line.contains("app-server event: turn/completed") + .then_some(index) + }) + .collect::>(); + assert_eq!( + turn_completions.len(), + 2, + "expected the child completion before the primary completion: {stderr}" + ); + + let [child_completion, primary_completion] = turn_completions.as_slice() else { + unreachable!("checked turn/completed count") + }; + assert_eq!( + lines[*child_completion + 1..*primary_completion] + .iter() + .filter(|line| line.contains("app-server typed request")) + .count(), + 0, + "the unrelated completion must not issue thread/read: {stderr}" + ); + assert_eq!( + lines[*primary_completion + 1..] + .iter() + .filter(|line| line.contains("app-server typed request")) + .count(), + 2, + "the primary completion should issue thread/read and thread/unsubscribe: {stderr}" + ); + + assert_eq!( + lines + .iter() + .filter(|line| line.contains("app-server typed request")) + .count(), + 5, + "only the primary completion should issue an extra request: {stderr}" + ); + + Ok(()) +} diff --git a/codex-rs/exec/tests/suite/ephemeral.rs b/codex-rs/exec/tests/suite/ephemeral.rs new file mode 100644 index 0000000000000000000000000000000000000000..83522e4bd1572503c361ebb8f13fd04ca8812d0b --- /dev/null +++ b/codex-rs/exec/tests/suite/ephemeral.rs @@ -0,0 +1,67 @@ +#![cfg(not(target_os = "windows"))] +#![allow(clippy::unwrap_used)] + +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex_exec::test_codex_exec; +use walkdir::WalkDir; +use wiremock::MockServer; + +fn exec_sse_response() -> String { + responses::sse(vec![ + responses::ev_response_created("resp-ephemeral"), + responses::ev_assistant_message("msg-ephemeral", "ephemeral response"), + responses::ev_completed("resp-ephemeral"), + ]) +} + +fn session_rollout_count(home_path: &std::path::Path) -> usize { + let sessions_dir = home_path.join("sessions"); + if !sessions_dir.exists() { + return 0; + } + + WalkDir::new(sessions_dir) + .into_iter() + .filter_map(Result::ok) + .filter(|entry| entry.file_type().is_file()) + .filter(|entry| entry.file_name().to_string_lossy().ends_with(".jsonl")) + .count() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn persists_rollout_file_by_default() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let server = MockServer::start().await; + let _response_mock = responses::mount_sse_once(&server, exec_sse_response()).await; + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("default persistence behavior") + .assert() + .code(0); + + assert_eq!(session_rollout_count(test.home_path()), 1); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn does_not_persist_rollout_file_in_ephemeral_mode() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let server = MockServer::start().await; + let _response_mock = responses::mount_sse_once(&server, exec_sse_response()).await; + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("--ephemeral") + .arg("ephemeral behavior") + .assert() + .code(0); + + assert_eq!(session_rollout_count(test.home_path()), 0); + Ok(()) +} diff --git a/codex-rs/exec/tests/suite/hooks.rs b/codex-rs/exec/tests/suite/hooks.rs new file mode 100644 index 0000000000000000000000000000000000000000..7be9480d11db64bf720999a52de70492370fec8b --- /dev/null +++ b/codex-rs/exec/tests/suite/hooks.rs @@ -0,0 +1,43 @@ +#![cfg(not(target_os = "windows"))] + +use core_test_support::responses; +use core_test_support::test_codex_exec::test_codex_exec; +use serde_json::json; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_hook_trust_bypass_runs_session_start_hook() -> anyhow::Result<()> { + let test = test_codex_exec(); + let marker_path = test.home_path().join("session-start-ran"); + let command = format!("touch {}", marker_path.display()); + std::fs::write( + test.home_path().join("hooks.json"), + serde_json::to_vec_pretty(&json!({ + "hooks": { + "SessionStart": [{ + "hooks": [{ + "type": "command", + "command": command, + }], + }], + }, + }))?, + )?; + + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("response_1"), + responses::ev_assistant_message("response_1", "done"), + responses::ev_completed("response_1"), + ]); + responses::mount_sse_once(&server, body).await; + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("--dangerously-bypass-hook-trust") + .arg("run the session start hook") + .assert() + .success(); + + assert!(marker_path.exists(), "session start hook did not run"); + Ok(()) +} diff --git a/codex-rs/exec/tests/suite/mcp_required_exit.rs b/codex-rs/exec/tests/suite/mcp_required_exit.rs new file mode 100644 index 0000000000000000000000000000000000000000..738dc3a5d4c471545c86e646631558b54db216d6 --- /dev/null +++ b/codex-rs/exec/tests/suite/mcp_required_exit.rs @@ -0,0 +1,38 @@ +#![cfg(not(target_os = "windows"))] +#![allow(clippy::unwrap_used)] + +use core_test_support::responses; +use core_test_support::test_codex_exec::test_codex_exec; +use predicates::str::contains; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exits_non_zero_when_required_mcp_server_fails_to_initialize() -> anyhow::Result<()> { + let test = test_codex_exec(); + + let config_toml = r#" + [mcp_servers.required_broken] + command = "codex-definitely-not-a-real-binary" + required = true + "#; + std::fs::write(test.home_path().join("config.toml"), config_toml)?; + + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("resp_1"), + responses::ev_assistant_message("msg_1", "hello"), + responses::ev_completed("resp_1"), + ]); + responses::mount_sse_once(&server, body).await; + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("--experimental-json") + .arg("tell me something") + .assert() + .code(1) + .stderr(contains( + "required MCP servers failed to initialize: required_broken", + )); + + Ok(()) +} diff --git a/codex-rs/exec/tests/suite/mod.rs b/codex-rs/exec/tests/suite/mod.rs new file mode 100644 index 0000000000000000000000000000000000000000..4f7a3ab89fc91ad3c974f58ea8cd84c7ef74baee --- /dev/null +++ b/codex-rs/exec/tests/suite/mod.rs @@ -0,0 +1,20 @@ +// Aggregates all former standalone integration tests as modules. +mod add_dir; +mod agents_md; +mod apply_patch; +mod approval_policy; +mod auth_env; +#[path = "completion_backfill_tests.rs"] +mod completion_backfill; +mod ephemeral; +mod hooks; +mod mcp_required_exit; +mod originator; +mod output_schema; +mod prompt_stdin; +mod resume; +mod sandbox; +#[cfg(target_os = "macos")] +mod seatbelt; +mod server_error_exit; +mod worktree; diff --git a/codex-rs/exec/tests/suite/originator.rs b/codex-rs/exec/tests/suite/originator.rs new file mode 100644 index 0000000000000000000000000000000000000000..d2a1f829f5ce5b1967b5717409980cc5e9bee917 --- /dev/null +++ b/codex-rs/exec/tests/suite/originator.rs @@ -0,0 +1,54 @@ +#![cfg(not(target_os = "windows"))] +#![allow(clippy::unwrap_used)] + +use codex_login::default_client::CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR; +use core_test_support::responses; +use core_test_support::test_codex_exec::test_codex_exec; +use wiremock::matchers::header; + +/// Verify that when the server reports an error, `codex-exec` exits with a +/// non-zero status code so automation can detect failures. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn send_codex_exec_originator() -> anyhow::Result<()> { + let test = test_codex_exec(); + + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("response_1"), + responses::ev_assistant_message("response_1", "Hello, world!"), + responses::ev_completed("response_1"), + ]); + responses::mount_sse_once_match(&server, header("Originator", "codex_exec"), body).await; + + test.cmd_with_server(&server) + .env_remove(CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR) + .arg("--skip-git-repo-check") + .arg("tell me something") + .assert() + .code(0); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn supports_originator_override() -> anyhow::Result<()> { + let test = test_codex_exec(); + + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("response_1"), + responses::ev_assistant_message("response_1", "Hello, world!"), + responses::ev_completed("response_1"), + ]); + responses::mount_sse_once_match(&server, header("Originator", "codex_exec_override"), body) + .await; + + test.cmd_with_server(&server) + .env("CODEX_INTERNAL_ORIGINATOR_OVERRIDE", "codex_exec_override") + .arg("--skip-git-repo-check") + .arg("tell me something") + .assert() + .code(0); + + Ok(()) +} diff --git a/codex-rs/exec/tests/suite/output_schema.rs b/codex-rs/exec/tests/suite/output_schema.rs new file mode 100644 index 0000000000000000000000000000000000000000..bfc9e8e589d5c7a271e2230b4732bdc94ccd4d34 --- /dev/null +++ b/codex-rs/exec/tests/suite/output_schema.rs @@ -0,0 +1,62 @@ +#![cfg(not(target_os = "windows"))] +#![allow(clippy::unwrap_used)] + +use core_test_support::responses; +use core_test_support::test_codex_exec::test_codex_exec; +use serde_json::Value; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_includes_output_schema_in_request() -> anyhow::Result<()> { + let test = test_codex_exec(); + + let schema_contents = serde_json::json!({ + "type": "object", + "properties": { + "answer": { "type": "string" } + }, + "required": ["answer"], + "additionalProperties": false + }); + let schema_path = test.cwd_path().join("schema.json"); + std::fs::write(&schema_path, serde_json::to_vec_pretty(&schema_contents)?)?; + let expected_schema: Value = schema_contents; + + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("resp1"), + responses::ev_assistant_message("m1", "fixture hello"), + responses::ev_completed("resp1"), + ]); + let response_mock = responses::mount_sse_once(&server, body).await; + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + // keep using -C in the test to exercise the flag as well + .arg("-C") + .arg(test.cwd_path()) + .arg("--output-schema") + .arg(&schema_path) + .arg("-m") + .arg("gpt-5.1") + .arg("tell me a joke") + .assert() + .success(); + + let request = response_mock.single_request(); + let payload: Value = request.body_json(); + let text = payload.get("text").expect("request missing text field"); + let format = text + .get("format") + .expect("request missing text.format field"); + assert_eq!( + format, + &serde_json::json!({ + "name": "codex_output_schema", + "type": "json_schema", + "strict": true, + "schema": expected_schema, + }) + ); + + Ok(()) +} diff --git a/codex-rs/exec/tests/suite/prompt_stdin.rs b/codex-rs/exec/tests/suite/prompt_stdin.rs new file mode 100644 index 0000000000000000000000000000000000000000..4789fa42aec41d7c808bdd0f286affa33a0f108d --- /dev/null +++ b/codex-rs/exec/tests/suite/prompt_stdin.rs @@ -0,0 +1,171 @@ +#![cfg(not(target_os = "windows"))] +#![allow(clippy::unwrap_used)] + +use core_test_support::responses; +use core_test_support::test_codex_exec::test_codex_exec; +use predicates::str::contains; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_appends_piped_stdin_to_prompt_argument() -> anyhow::Result<()> { + let test = test_codex_exec(); + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("resp1"), + responses::ev_assistant_message("m1", "fixture hello"), + responses::ev_completed("resp1"), + ]); + let response_mock = responses::mount_sse_once(&server, body).await; + + // echo "my output" | codex exec --skip-git-repo-check -C -m gpt-5.1 "Summarize this concisely" + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(test.cwd_path()) + .arg("-m") + .arg("gpt-5.1") + .arg("Summarize this concisely") + .write_stdin("my output\n") + .assert() + .success(); + + let request = response_mock.single_request(); + assert!( + request.has_message_with_input_texts("user", |texts| { + texts == ["Summarize this concisely\n\n\nmy output\n".to_string()] + }), + "request should include a user message with the prompt plus piped stdin context" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_ignores_empty_piped_stdin_when_prompt_argument_is_present() -> anyhow::Result<()> { + let test = test_codex_exec(); + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("resp1"), + responses::ev_assistant_message("m1", "fixture hello"), + responses::ev_completed("resp1"), + ]); + let response_mock = responses::mount_sse_once(&server, body).await; + + // printf "" | codex exec --skip-git-repo-check -C -m gpt-5.1 "Summarize this concisely" + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(test.cwd_path()) + .arg("-m") + .arg("gpt-5.1") + .arg("Summarize this concisely") + .write_stdin("") + .assert() + .success(); + + let request = response_mock.single_request(); + assert!( + request.has_message_with_input_texts("user", |texts| texts + == ["Summarize this concisely".to_string()]), + "request should preserve the prompt when stdin is empty" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_dash_prompt_reads_stdin_as_the_prompt() -> anyhow::Result<()> { + let test = test_codex_exec(); + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("resp1"), + responses::ev_assistant_message("m1", "fixture hello"), + responses::ev_completed("resp1"), + ]); + let response_mock = responses::mount_sse_once(&server, body).await; + + // echo "prompt from stdin" | codex exec --skip-git-repo-check -C -m gpt-5.1 - + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(test.cwd_path()) + .arg("-m") + .arg("gpt-5.1") + .arg("-") + .write_stdin("prompt from stdin\n") + .assert() + .success(); + + let request = response_mock.single_request(); + assert!( + request.has_message_with_input_texts("user", |texts| { + texts == ["prompt from stdin\n".to_string()] + }), + "dash prompt should preserve the existing forced-stdin behavior" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_without_prompt_argument_reads_piped_stdin_as_the_prompt() -> anyhow::Result<()> { + let test = test_codex_exec(); + let server = responses::start_mock_server().await; + let body = responses::sse(vec![ + responses::ev_response_created("resp1"), + responses::ev_assistant_message("m1", "fixture hello"), + responses::ev_completed("resp1"), + ]); + let response_mock = responses::mount_sse_once(&server, body).await; + + // echo "prompt from stdin" | codex exec --skip-git-repo-check -C -m gpt-5.1 + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(test.cwd_path()) + .arg("-m") + .arg("gpt-5.1") + .write_stdin("prompt from stdin\n") + .assert() + .success(); + + let request = response_mock.single_request(); + assert!( + request.has_message_with_input_texts("user", |texts| { + texts == ["prompt from stdin\n".to_string()] + }), + "missing prompt argument should preserve the existing piped-stdin prompt behavior" + ); + + Ok(()) +} + +#[test] +fn exec_without_prompt_argument_rejects_empty_piped_stdin() { + let test = test_codex_exec(); + + // printf "" | codex exec --skip-git-repo-check -C + test.cmd() + .arg("--skip-git-repo-check") + .arg("-C") + .arg(test.cwd_path()) + .write_stdin("") + .assert() + .code(1) + .stderr(contains("No prompt provided via stdin.")); +} + +#[test] +fn exec_dash_prompt_rejects_empty_piped_stdin() { + let test = test_codex_exec(); + + // printf "" | codex exec --skip-git-repo-check -C - + test.cmd() + .arg("--skip-git-repo-check") + .arg("-C") + .arg(test.cwd_path()) + .arg("-") + .write_stdin("") + .assert() + .code(1) + .stderr(contains("No prompt provided via stdin.")); +} diff --git a/codex-rs/exec/tests/suite/resume.rs b/codex-rs/exec/tests/suite/resume.rs new file mode 100644 index 0000000000000000000000000000000000000000..d318bcf73792b6e85af7852022e9ec87746467df --- /dev/null +++ b/codex-rs/exec/tests/suite/resume.rs @@ -0,0 +1,1054 @@ +#![allow(clippy::unwrap_used)] +use anyhow::Context; +use codex_core::config::ConfigBuilder; +use codex_core::init_state_db; +use codex_protocol::ThreadId; +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex_exec::test_codex_exec; +use pretty_assertions::assert_eq; +use serde_json::Value; +use std::process::Stdio; +use std::string::ToString; +use std::time::Duration; +use tempfile::TempDir; +use uuid::Uuid; +use walkdir::WalkDir; +use wiremock::MockServer; + +/// Utility: scan the sessions dir for a rollout file that contains `marker` +/// in any response_item.message.content entry. Returns the absolute path. +fn find_session_file_containing_marker( + sessions_dir: &std::path::Path, + marker: &str, +) -> Option { + for entry in WalkDir::new(sessions_dir) { + let entry = match entry { + Ok(e) => e, + Err(_) => continue, + }; + if !entry.file_type().is_file() { + continue; + } + if !entry.file_name().to_string_lossy().ends_with(".jsonl") { + continue; + } + let path = entry.path(); + let Ok(content) = std::fs::read_to_string(path) else { + continue; + }; + // Skip the first meta line and scan remaining JSONL entries. + let mut lines = content.lines(); + if lines.next().is_none() { + continue; + } + for line in lines { + if line.trim().is_empty() { + continue; + } + let Ok(item): Result = serde_json::from_str(line) else { + continue; + }; + if item.get("type").and_then(|t| t.as_str()) == Some("response_item") + && let Some(payload) = item.get("payload") + && payload.get("type").and_then(|t| t.as_str()) == Some("message") + && payload + .get("content") + .map(ToString::to_string) + .unwrap_or_default() + .contains(marker) + { + return Some(path.to_path_buf()); + } + } + } + None +} + +/// Extract the conversation UUID from the first SessionMeta line in the rollout file. +fn extract_conversation_id(path: &std::path::Path) -> String { + let content = std::fs::read_to_string(path).unwrap(); + let mut lines = content.lines(); + let meta_line = lines.next().expect("missing meta line"); + let meta: Value = serde_json::from_str(meta_line).expect("invalid meta json"); + meta.get("payload") + .and_then(|p| p.get("id")) + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string() +} + +fn last_user_image_count(path: &std::path::Path) -> usize { + let content = std::fs::read_to_string(path).unwrap_or_default(); + let mut last_count = 0; + for line in content.lines() { + if line.trim().is_empty() { + continue; + } + let Ok(item): Result = serde_json::from_str(line) else { + continue; + }; + if item.get("type").and_then(|t| t.as_str()) != Some("response_item") { + continue; + } + let Some(payload) = item.get("payload") else { + continue; + }; + if payload.get("type").and_then(|t| t.as_str()) != Some("message") { + continue; + } + if payload.get("role").and_then(|r| r.as_str()) != Some("user") { + continue; + } + let Some(content_items) = payload.get("content").and_then(|v| v.as_array()) else { + continue; + }; + last_count = content_items + .iter() + .filter(|entry| entry.get("type").and_then(|t| t.as_str()) == Some("input_image")) + .count(); + } + last_count +} + +fn exec_repo_root() -> anyhow::Result { + Ok(codex_utils_cargo_bin::repo_root()?) +} + +fn exec_sse_response(index: usize) -> String { + let response_id = format!("resp-exec-{index}"); + let message_id = format!("msg-exec-{index}"); + responses::sse(vec![ + responses::ev_response_created(&response_id), + responses::ev_assistant_message(&message_id, "exec response"), + responses::ev_completed(&response_id), + ]) +} + +async fn mount_exec_responses( + server: &MockServer, + count: usize, +) -> core_test_support::responses::ResponseMock { + responses::mount_sse_sequence(server, (0..count).map(exec_sse_response).collect()).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_falls_back_to_legacy_history_when_thread_store_cannot_paginate() -> anyhow::Result<()> +{ + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let server = MockServer::start().await; + let _response_mock = mount_exec_responses(&server, /*count*/ 1).await; + let store_id = Uuid::new_v4(); + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-c") + .arg(format!( + "experimental_thread_store={{type=\"in_memory\",id=\"{store_id}\"}}" + )) + .arg("continue without paginated history") + .assert() + .success(); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_resume_last_appends_to_existing_file() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let server = MockServer::start().await; + let response_mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_response_created("resp-exec-0"), + responses::ev_assistant_message("msg-exec-0", "exec response"), + responses::ev_completed_with_tokens("resp-exec-0", /*total_tokens*/ 7), + ]), + exec_sse_response(/*index*/ 1), + ], + ) + .await; + let repo_root = exec_repo_root()?; + + // 1) First run: create a session with a unique marker in the content. + let marker = format!("resume-last-{}", Uuid::new_v4()); + let prompt = format!("echo {marker}"); + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg(&prompt) + .assert() + .success(); + + // Find the created session file containing the marker. + let sessions_dir = test.home_path().join("sessions"); + let path = find_session_file_containing_marker(&sessions_dir, &marker) + .expect("no session file found after first run"); + let content = std::fs::read_to_string(&path)?; + let meta: Value = serde_json::from_str( + content + .lines() + .next() + .expect("rollout should contain session metadata"), + )?; + assert_eq!(meta["payload"]["history_mode"], "paginated"); + assert_eq!(meta["payload"]["thread_source"], "user"); + + // 2) Second run: resume the most recent file with a new marker. + let marker2 = format!("resume-last-2-{}", Uuid::new_v4()); + let prompt2 = format!("echo {marker2}"); + + let output = test + .cmd_with_server(&server) + .env("RUST_LOG", "codex_app_server::outgoing_message=trace") + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg(&prompt2) + .arg("resume") + .arg("--last") + .output() + .context("resume run should succeed")?; + let stderr = String::from_utf8_lossy(&output.stderr); + assert!(output.status.success(), "resume failed: {stderr}"); + assert_eq!( + stderr + .matches("app-server event: thread/tokenUsage/updated") + .count(), + 2, + "paginated resume should replay restored token usage before the new turn: {stderr}" + ); + + // Ensure the same file was updated and contains both markers. + let resumed_path = find_session_file_containing_marker(&sessions_dir, &marker2) + .expect("no resumed session file containing marker2"); + assert_eq!( + resumed_path, path, + "resume --last should append to existing file" + ); + let content = std::fs::read_to_string(&resumed_path)?; + assert!(content.contains(&marker)); + assert!(content.contains(&marker2)); + let requests = response_mock.requests(); + assert_eq!(requests.len(), 2); + let resumed_request = requests[1].body_json().to_string(); + assert!(resumed_request.contains(&marker)); + assert!(resumed_request.contains(&marker2)); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_resume_last_repairs_rollout_missing_from_state_db() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let server = MockServer::start().await; + let _response_mock = mount_exec_responses(&server, /*count*/ 2).await; + let repo_root = exec_repo_root()?; + + let marker = format!("resume-last-repair-{}", Uuid::new_v4()); + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg(format!("echo {marker}")) + .assert() + .success(); + + let sessions_dir = test.home_path().join("sessions"); + let path = find_session_file_containing_marker(&sessions_dir, &marker) + .expect("no session file found after first run"); + let thread_id = ThreadId::from_string(&extract_conversation_id(&path))?; + let config = ConfigBuilder::default() + .codex_home(test.home_path().to_path_buf()) + .build() + .await?; + let state_db = init_state_db(&config) + .await + .expect("state DB should initialize"); + assert_eq!(state_db.delete_thread(thread_id).await?, 1); + state_db + .mark_backfill_complete(/*last_watermark*/ None) + .await?; + + let resumed_marker = format!("resume-last-repaired-{}", Uuid::new_v4()); + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg("resume") + .arg("--last") + .arg(format!("echo {resumed_marker}")) + .assert() + .success(); + + let resumed_path = find_session_file_containing_marker(&sessions_dir, &resumed_marker) + .expect("no resumed session file after SQLite repair"); + assert_eq!(resumed_path, path); + assert!(state_db.get_thread(thread_id).await?.is_some()); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_resume_last_trusts_usable_state_db_candidate() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let server = MockServer::start().await; + let _response_mock = mount_exec_responses(&server, /*count*/ 3).await; + let repo_root = exec_repo_root()?; + let sessions_dir = test.home_path().join("sessions"); + + let older_marker = format!("resume-last-indexed-{}", Uuid::new_v4()); + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg(format!("echo {older_marker}")) + .assert() + .success(); + let older_path = find_session_file_containing_marker(&sessions_dir, &older_marker) + .expect("no indexed session file after first run"); + + let newer_marker = format!("resume-last-unindexed-{}", Uuid::new_v4()); + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg(format!("echo {newer_marker}")) + .assert() + .success(); + let newer_path = find_session_file_containing_marker(&sessions_dir, &newer_marker) + .expect("no unindexed session file after second run"); + let newer_thread_id = ThreadId::from_string(&extract_conversation_id(&newer_path))?; + + let config = ConfigBuilder::default() + .codex_home(test.home_path().to_path_buf()) + .build() + .await?; + let state_db = init_state_db(&config) + .await + .expect("state DB should initialize"); + assert_eq!(state_db.delete_thread(newer_thread_id).await?, 1); + state_db + .mark_backfill_complete(/*last_watermark*/ None) + .await?; + + let resumed_marker = format!("resume-last-authoritative-{}", Uuid::new_v4()); + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg("resume") + .arg("--last") + .arg(format!("echo {resumed_marker}")) + .assert() + .success(); + + let resumed_path = find_session_file_containing_marker(&sessions_dir, &resumed_marker) + .expect("no resumed session file after SQLite lookup"); + assert_eq!( + ( + resumed_path, + state_db + .get_thread(newer_thread_id) + .await? + .map(|metadata| metadata.id), + ), + (older_path, None), + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_resume_last_skips_mismatched_state_db_candidate() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let server = MockServer::start().await; + let _response_mock = mount_exec_responses(&server, /*count*/ 3).await; + let repo_root = exec_repo_root()?; + let sessions_dir = test.home_path().join("sessions"); + + let older_marker = format!("resume-last-valid-{}", Uuid::new_v4()); + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg(format!("echo {older_marker}")) + .assert() + .success(); + let older_path = find_session_file_containing_marker(&sessions_dir, &older_marker) + .expect("no valid session file after first run"); + + let newer_marker = format!("resume-last-mismatched-{}", Uuid::new_v4()); + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg(format!("echo {newer_marker}")) + .assert() + .success(); + let newer_path = find_session_file_containing_marker(&sessions_dir, &newer_marker) + .expect("no mismatched session file after second run"); + let newer_thread_id = ThreadId::from_string(&extract_conversation_id(&newer_path))?; + + let config = ConfigBuilder::default() + .codex_home(test.home_path().to_path_buf()) + .build() + .await?; + let state_db = init_state_db(&config) + .await + .expect("state DB should initialize"); + let mut mismatched = state_db + .get_thread(newer_thread_id) + .await? + .expect("newer thread should be indexed"); + mismatched.rollout_path = older_path.clone(); + state_db.upsert_thread(&mismatched).await?; + + let resumed_marker = format!("resume-last-valid-resumed-{}", Uuid::new_v4()); + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg("resume") + .arg("--last") + .arg(format!("echo {resumed_marker}")) + .assert() + .success(); + + let resumed_path = find_session_file_containing_marker(&sessions_dir, &resumed_marker) + .expect("no resumed session file after skipping mismatched SQLite row"); + assert_eq!(resumed_path, older_path); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_resume_last_accepts_prompt_after_flag_in_json_mode() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let server = MockServer::start().await; + let _response_mock = mount_exec_responses(&server, /*count*/ 2).await; + let repo_root = exec_repo_root()?; + + // 1) First run: create a session with a unique marker in the content. + let marker = format!("resume-last-json-{}", Uuid::new_v4()); + let prompt = format!("echo {marker}"); + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg(&prompt) + .assert() + .success(); + + // Find the created session file containing the marker. + let sessions_dir = test.home_path().join("sessions"); + let path = find_session_file_containing_marker(&sessions_dir, &marker) + .expect("no session file found after first run"); + + // 2) Second run: resume the most recent file and pass the prompt after --last. + let marker2 = format!("resume-last-json-2-{}", Uuid::new_v4()); + let prompt2 = format!("echo {marker2}"); + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg("--json") + .arg("resume") + .arg("--last") + .arg(&prompt2) + .assert() + .success(); + + let resumed_path = find_session_file_containing_marker(&sessions_dir, &marker2) + .expect("no resumed session file containing marker2"); + assert_eq!( + resumed_path, path, + "resume --last should append to existing file" + ); + let content = std::fs::read_to_string(&resumed_path)?; + assert!(content.contains(&marker)); + assert!(content.contains(&marker2)); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_resume_last_respects_cwd_filter_and_all_flag() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let server = MockServer::start().await; + let _response_mock = mount_exec_responses(&server, /*count*/ 5).await; + + let dir_a = TempDir::new()?; + let dir_b = TempDir::new()?; + + let marker_a = format!("resume-cwd-a-{}", Uuid::new_v4()); + let prompt_a = format!("echo {marker_a}"); + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(dir_a.path()) + .arg(&prompt_a) + .assert() + .success(); + + let marker_b = format!("resume-cwd-b-{}", Uuid::new_v4()); + let prompt_b = format!("echo {marker_b}"); + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(dir_b.path()) + .arg(&prompt_b) + .assert() + .success(); + + let sessions_dir = test.home_path().join("sessions"); + find_session_file_containing_marker(&sessions_dir, &marker_a) + .expect("no session file found for marker_a"); + let path_b = find_session_file_containing_marker(&sessions_dir, &marker_b) + .expect("no session file found for marker_b"); + + // `updated_at` is second-granularity, so ensure the touch lands in a later second + // than the initial session creation on fast CI (especially Windows). + std::thread::sleep(std::time::Duration::from_millis(1100)); + + // Make thread B deterministically newest according to rollout metadata. + let session_id_b = extract_conversation_id(&path_b); + let marker_b_touch = format!("resume-cwd-b-touch-{}", Uuid::new_v4()); + let prompt_b_touch = format!("echo {marker_b_touch}"); + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(dir_b.path()) + .arg("resume") + .arg(&session_id_b) + .arg(&prompt_b_touch) + .assert() + .success(); + + // `resume --last` sorts by `updated_at`, which is second-granularity. Sleep so + // the upcoming `resume --last --all` write lands in a later second and becomes + // deterministically newest (instead of tying and falling back to UUID order). + std::thread::sleep(std::time::Duration::from_millis(1100)); + + let marker_b2 = format!("resume-cwd-b-2-{}", Uuid::new_v4()); + let prompt_b2 = format!("echo {marker_b2}"); + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(dir_a.path()) + .arg("resume") + .arg("--last") + .arg("--all") + .arg(&prompt_b2) + .assert() + .success(); + + let resumed_path_all = find_session_file_containing_marker(&sessions_dir, &marker_b2) + .expect("no resumed session file containing marker_b2"); + assert_eq!( + resumed_path_all, path_b, + "resume --last --all should pick newest session" + ); + + // Selection must still use the latest turn's cwd when only the compressed rollout exists. + zstd::stream::copy_encode( + std::fs::File::open(&path_b)?, + std::fs::File::create(path_b.with_extension("jsonl.zst"))?, + /*level*/ 3, + )?; + std::fs::remove_file(&path_b)?; + + let marker_a2 = format!("resume-cwd-a-2-{}", Uuid::new_v4()); + let prompt_a2 = format!("echo {marker_a2}"); + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(dir_a.path()) + .arg("resume") + .arg("--last") + .arg(&prompt_a2) + .assert() + .success(); + + let resumed_path_cwd = find_session_file_containing_marker(&sessions_dir, &marker_a2) + .expect("no resumed session file containing marker_a2"); + // The `--all` resume above appends a new turn to `path_b` while running from `dir_a`, so the + // session's latest cwd now matches `dir_a`. A subsequent `resume --last` should therefore pick + // the newest matching session (`path_b`). + assert_eq!( + resumed_path_cwd, path_b, + "resume --last should prefer sessions whose latest turn context matches the current cwd" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_resume_accepts_global_flags_after_subcommand() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let server = MockServer::start().await; + let _response_mock = mount_exec_responses(&server, /*count*/ 2).await; + + // Seed a session. + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("echo seed-resume-session") + .assert() + .success(); + + // Resume while passing global flags after the subcommand to ensure clap accepts them. + let base = format!("{}/v1", server.uri()); + let base_config = format!("openai_base_url={}", serde_json::to_string(&base)?); + test.cmd() + .arg("resume") + .arg("--last") + .arg("--config") + .arg(base_config) + .arg("--json") + .arg("--model") + .arg("gpt-5.2-codex") + .arg("--config") + .arg("reasoning_level=xhigh") + .arg("--dangerously-bypass-approvals-and-sandbox") + .arg("--skip-git-repo-check") + .arg("echo resume-with-global-flags-after-subcommand") + .assert() + .success(); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_resume_includes_output_schema_in_request() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let server = MockServer::start().await; + let response_mock = mount_exec_responses(&server, /*count*/ 2).await; + + let schema_contents = serde_json::json!({ + "type": "object", + "properties": { + "answer": { "type": "string" } + }, + "required": ["answer"], + "additionalProperties": false + }); + let schema_path = test.cwd_path().join("schema.json"); + std::fs::write(&schema_path, serde_json::to_vec_pretty(&schema_contents)?)?; + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("echo seed-resume-session") + .assert() + .success(); + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("resume") + .arg("--last") + .arg("--json") + .arg("--output-schema") + .arg(&schema_path) + .arg("echo resume-with-schema") + .assert() + .success(); + + let requests = response_mock.requests(); + assert_eq!(requests.len(), 2); + let payload: Value = requests[1].body_json(); + let text = payload.get("text").expect("request missing text field"); + let format = text + .get("format") + .expect("request missing text.format field"); + assert_eq!( + format, + &serde_json::json!({ + "name": "codex_output_schema", + "type": "json_schema", + "strict": true, + "schema": schema_contents, + }) + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_resume_by_id_appends_to_existing_file() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let server = MockServer::start().await; + let _response_mock = mount_exec_responses(&server, /*count*/ 2).await; + let repo_root = exec_repo_root()?; + + // 1) First run: create a session + let marker = format!("resume-by-id-{}", Uuid::new_v4()); + let prompt = format!("echo {marker}"); + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg(&prompt) + .assert() + .success(); + + let sessions_dir = test.home_path().join("sessions"); + let path = find_session_file_containing_marker(&sessions_dir, &marker) + .expect("no session file found after first run"); + let session_id = extract_conversation_id(&path); + assert!( + !session_id.is_empty(), + "missing conversation id in meta line" + ); + + // 2) Resume by id + let marker2 = format!("resume-by-id-2-{}", Uuid::new_v4()); + let prompt2 = format!("echo {marker2}"); + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg(&prompt2) + .arg("resume") + .arg(&session_id) + .assert() + .success(); + + let resumed_path = find_session_file_containing_marker(&sessions_dir, &marker2) + .expect("no resumed session file containing marker2"); + assert_eq!( + resumed_path, path, + "resume by id should append to existing file" + ); + let content = std::fs::read_to_string(&resumed_path)?; + assert!(content.contains(&marker)); + assert!(content.contains(&marker2)); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_resume_preserves_cli_configuration_overrides() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let server = MockServer::start().await; + let _response_mock = mount_exec_responses(&server, /*count*/ 2).await; + let repo_root = exec_repo_root()?; + + let marker = format!("resume-config-{}", Uuid::new_v4()); + let prompt = format!("echo {marker}"); + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("--sandbox") + .arg("workspace-write") + .arg("--model") + .arg("gpt-5.1") + .arg("-C") + .arg(&repo_root) + .arg(&prompt) + .assert() + .success(); + + let sessions_dir = test.home_path().join("sessions"); + let path = find_session_file_containing_marker(&sessions_dir, &marker) + .expect("no session file found after first run"); + + let marker2 = format!("resume-config-2-{}", Uuid::new_v4()); + let prompt2 = format!("echo {marker2}"); + + let output = test + .cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("--sandbox") + .arg("workspace-write") + .arg("--model") + .arg("gpt-5.1-high") + .arg("-C") + .arg(&repo_root) + .arg(&prompt2) + .arg("resume") + .arg("--last") + .output() + .context("resume run should succeed")?; + + assert!(output.status.success(), "resume run failed: {output:?}"); + + let stderr = String::from_utf8(output.stderr)?; + assert!( + stderr.contains("model: gpt-5.1-high"), + "stderr missing model override: {stderr}" + ); + if cfg!(target_os = "windows") { + assert!( + stderr.contains("sandbox: read-only"), + "stderr missing downgraded sandbox note: {stderr}" + ); + } else { + assert!( + stderr.contains("sandbox: workspace-write"), + "stderr missing sandbox override: {stderr}" + ); + } + + let resumed_path = find_session_file_containing_marker(&sessions_dir, &marker2) + .expect("no resumed session file containing marker2"); + assert_eq!(resumed_path, path, "resume should append to same file"); + + let content = std::fs::read_to_string(&resumed_path)?; + assert!(content.contains(&marker)); + assert!(content.contains(&marker2)); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_resume_accepts_images_after_subcommand() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let server = MockServer::start().await; + let _response_mock = mount_exec_responses(&server, /*count*/ 2).await; + let repo_root = exec_repo_root()?; + + let marker = format!("resume-image-{}", Uuid::new_v4()); + let prompt = format!("echo {marker}"); + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg(&prompt) + .assert() + .success(); + + let image_path = test.cwd_path().join("resume_image.png"); + let image_path_2 = test.cwd_path().join("resume_image_2.png"); + let image_bytes: &[u8] = &[ + 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, + 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x06, 0x00, 0x00, 0x00, 0x1F, + 0x15, 0xC4, 0x89, 0x00, 0x00, 0x00, 0x0A, 0x49, 0x44, 0x41, 0x54, 0x78, 0x9C, 0x63, 0x00, + 0x01, 0x00, 0x00, 0x05, 0x00, 0x01, 0x0D, 0x0A, 0x2D, 0xB4, 0x00, 0x00, 0x00, 0x00, 0x49, + 0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82, + ]; + std::fs::write(&image_path, image_bytes)?; + std::fs::write(&image_path_2, image_bytes)?; + + let marker2 = format!("resume-image-2-{}", Uuid::new_v4()); + let prompt2 = format!("echo {marker2}"); + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(&repo_root) + .arg("resume") + .arg("--last") + .arg("--image") + .arg(&image_path) + .arg("--image") + .arg(&image_path_2) + .arg(&prompt2) + .assert() + .success(); + + let sessions_dir = test.home_path().join("sessions"); + let resumed_path = find_session_file_containing_marker(&sessions_dir, &marker2) + .expect("no session file found after resume with images"); + let image_count = last_user_image_count(&resumed_path); + assert_eq!( + image_count, 2, + "resume prompt should include both attached images" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_fork_creates_distinct_threads_with_and_without_a_prompt() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let test = test_codex_exec(); + let server = MockServer::start().await; + let response_mock = mount_exec_responses(&server, /*count*/ 2).await; + let source_marker = format!("fork-source-{}", Uuid::new_v4()); + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("--thread-source") + .arg("source_feature") + .arg(format!("echo {source_marker}")) + .assert() + .success(); + + let sessions_dir = test.home_path().join("sessions"); + let source_path = find_session_file_containing_marker(&sessions_dir, &source_marker) + .expect("source thread should have a rollout"); + let source_id = extract_conversation_id(&source_path); + let original_source = std::fs::read_to_string(&source_path)?; + let source_meta: Value = serde_json::from_str( + original_source + .lines() + .next() + .expect("source rollout should contain session metadata"), + )?; + assert_eq!(source_meta["payload"]["thread_source"], "source_feature"); + + for (args, expected_error) in [ + ( + vec!["--image", "unused.png"], + "Forking with images requires a prompt", + ), + ( + vec!["--output-schema", "unused.json"], + "Forking with output options requires a prompt", + ), + ( + vec!["--output-last-message", "unused.md"], + "Forking with output options requires a prompt", + ), + (vec!["--ephemeral"], "Ephemeral forks require a prompt"), + ] { + let output = test + .cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("fork") + .arg(&source_id) + .args(args) + .output()?; + assert!(!output.status.success()); + assert!( + String::from_utf8_lossy(&output.stderr).contains(expected_error), + "fork failed without the expected error: {output:?}" + ); + } + + let mut promptless_command = test.cmd_with_server(&server); + promptless_command + .arg("--skip-git-repo-check") + .arg("fork") + .arg(&source_id) + .arg("--json"); + let mut child_command = tokio::process::Command::new(promptless_command.get_program()); + child_command + .args(promptless_command.get_args()) + .envs( + promptless_command + .get_envs() + .filter_map(|(key, value)| value.map(|value| (key, value))), + ) + .current_dir(test.cwd_path()) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true); + let mut child = child_command.spawn()?; + let _open_stdin = child.stdin.take().expect("stdin should be piped"); + let promptless_output = + tokio::time::timeout(Duration::from_secs(/*secs*/ 10), child.wait_with_output()) + .await + .context("promptless fork should not wait for stdin to close")??; + assert!( + promptless_output.status.success(), + "promptless fork failed: {}", + String::from_utf8_lossy(&promptless_output.stderr) + ); + let promptless_events = String::from_utf8(promptless_output.stdout)? + .lines() + .map(serde_json::from_str::) + .collect::, _>>()?; + assert_eq!(promptless_events.len(), 1); + assert_eq!(promptless_events[0]["type"], "thread.started"); + let promptless_thread_id = promptless_events[0]["thread_id"] + .as_str() + .expect("promptless fork should emit its new thread id"); + assert_ne!(promptless_thread_id, source_id); + assert_eq!(response_mock.requests().len(), 1); + + let source_name = format!("fork-named-{}", Uuid::new_v4()); + let config = ConfigBuilder::default() + .codex_home(test.home_path().to_path_buf()) + .build() + .await?; + let state_db = init_state_db(&config) + .await + .expect("state DB should initialize"); + assert!( + state_db + .update_thread_title(ThreadId::from_string(&source_id)?, &source_name) + .await? + ); + + let fork_marker = format!("fork-prompt-{}", Uuid::new_v4()); + let fork_output = test + .cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("-C") + .arg(test.home_path()) + .arg("fork") + .arg(&source_name) + .arg("--thread-source") + .arg("fork_feature") + .arg("--json") + .arg("-") + .write_stdin(format!("echo {fork_marker}")) + .output()?; + assert!( + fork_output.status.success(), + "fork with prompt failed: {}", + String::from_utf8_lossy(&fork_output.stderr) + ); + let fork_events = String::from_utf8(fork_output.stdout)? + .lines() + .map(serde_json::from_str::) + .collect::, _>>()?; + assert_eq!(fork_events[0]["type"], "thread.started"); + let fork_thread_id = fork_events[0]["thread_id"] + .as_str() + .expect("fork should emit its new thread id"); + assert_ne!(fork_thread_id, source_id); + assert_ne!(fork_thread_id, promptless_thread_id); + + let fork_path = find_session_file_containing_marker(&sessions_dir, &fork_marker) + .expect("forked thread should have a separate rollout"); + assert_ne!(fork_path, source_path); + assert_eq!(extract_conversation_id(&fork_path), fork_thread_id); + let fork_contents = std::fs::read_to_string(&fork_path)?; + let fork_meta: Value = serde_json::from_str( + fork_contents + .lines() + .next() + .expect("fork rollout should contain session metadata"), + )?; + assert_eq!(fork_meta["payload"]["forked_from_id"], source_id); + assert_eq!(fork_meta["payload"]["thread_source"], "fork_feature"); + assert_eq!(fork_meta["payload"]["history_base"]["thread_id"], source_id); + assert!(!fork_contents.contains(&source_marker)); + assert!(fork_contents.contains(&fork_marker)); + assert_eq!(std::fs::read_to_string(&source_path)?, original_source); + + let requests = response_mock.requests(); + assert_eq!(requests.len(), 2); + let fork_request = requests[1].body_json().to_string(); + assert!(fork_request.contains(&source_marker)); + assert!(fork_request.contains(&fork_marker)); + + Ok(()) +} diff --git a/codex-rs/exec/tests/suite/sandbox.rs b/codex-rs/exec/tests/suite/sandbox.rs new file mode 100644 index 0000000000000000000000000000000000000000..1ccaee58497fdcb4ebfc7da6475d58dbd5facab0 --- /dev/null +++ b/codex-rs/exec/tests/suite/sandbox.rs @@ -0,0 +1,552 @@ +#![cfg(unix)] +use codex_core::spawn::StdioPolicy; +use codex_protocol::models::PermissionProfile; +use codex_protocol::permissions::NetworkSandboxPolicy; +use codex_utils_absolute_path::AbsolutePathBuf; +use codex_utils_absolute_path::test_support::PathBufExt; +use codex_utils_path_uri::PathUri; +use std::collections::HashMap; +use std::future::Future; +use std::io; +use std::process::ExitStatus; +use tokio::fs::create_dir_all; +use tokio::process::Child; + +pub(super) async fn spawn_command_under_sandbox( + command: Vec, + command_cwd: AbsolutePathBuf, + permission_profile: &PermissionProfile, + sandbox_cwd: &AbsolutePathBuf, + stdio_policy: StdioPolicy, + env: HashMap, +) -> std::io::Result { + use codex_core::exec::ExecCapturePolicy; + use codex_core::exec::ExecParams; + use codex_core::exec::build_exec_request; + use codex_core::sandboxing::SandboxPermissions; + use codex_protocol::config_types::WindowsSandboxLevel; + use std::process::Stdio; + + #[cfg(target_os = "linux")] + let codex_linux_sandbox_exe = Some( + core_test_support::find_codex_linux_sandbox_exe() + .map_err(|err| io::Error::new(io::ErrorKind::NotFound, err))?, + ); + #[cfg(target_os = "macos")] + let codex_linux_sandbox_exe = None; + let exec_request = build_exec_request( + ExecParams { + command, + cwd: command_cwd, + expiration: 1000.into(), + capture_policy: ExecCapturePolicy::ShellTool, + env, + network: None, + network_environment_id: None, + sandbox_permissions: SandboxPermissions::UseDefault, + windows_sandbox_level: WindowsSandboxLevel::Disabled, + windows_sandbox_private_desktop: false, + justification: None, + arg0: None, + }, + permission_profile, + sandbox_cwd, + &[PathUri::from_abs_path(sandbox_cwd)], + &codex_linux_sandbox_exe, + /*codex_self_exe*/ &None, + /*use_legacy_landlock*/ false, + ) + .map_err(|err| io::Error::other(err.to_string()))?; + + let (program, args) = exec_request + .command + .split_first() + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "command args are empty"))?; + + let mut child = tokio::process::Command::new(program); + if let Some(arg0) = exec_request.arg0.as_deref() { + child.arg0(arg0); + } + child.args(args); + // TODO(anp): Keep PathUri through the macOS sandbox process launch boundary. + let native_cwd = exec_request + .cwd + .to_abs_path() + .map_err(|err| io::Error::other(err.to_string()))?; + child.current_dir(native_cwd); + child.env_clear(); + child.envs(exec_request.env); + + match stdio_policy { + StdioPolicy::RedirectForShellTool => { + child.stdin(Stdio::null()); + child.stdout(Stdio::piped()).stderr(Stdio::piped()); + } + StdioPolicy::Inherit => { + child + .stdin(Stdio::inherit()) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()); + } + } + + child.kill_on_drop(true).spawn() +} + +#[cfg(target_os = "linux")] +/// Determines whether Linux sandbox tests can run on this host. +/// +/// These tests require an enforceable filesystem sandbox. We run a tiny command +/// under the production Landlock path and skip when enforcement is unavailable +/// (for example on kernels or container profiles where Landlock is not +/// enforced). +async fn linux_sandbox_test_env() -> Option> { + let command_cwd = AbsolutePathBuf::current_dir().ok()?; + let sandbox_cwd = command_cwd.clone(); + let permission_profile = PermissionProfile::read_only(); + + if can_apply_linux_sandbox_policy( + &permission_profile, + &command_cwd, + &sandbox_cwd, + HashMap::new(), + ) + .await + { + return Some(HashMap::new()); + } + + eprintln!("Skipping test: Landlock is not enforceable on this host."); + None +} + +#[cfg(target_os = "linux")] +/// Returns whether a minimal command can run successfully with the requested +/// Linux sandbox policy applied. +/// +/// This is used as a capability probe so sandbox behavior tests only run when +/// Landlock enforcement is actually active. +async fn can_apply_linux_sandbox_policy( + permission_profile: &PermissionProfile, + command_cwd: &AbsolutePathBuf, + sandbox_cwd: &AbsolutePathBuf, + env: HashMap, +) -> bool { + let spawn_result = spawn_command_under_sandbox( + vec!["/usr/bin/true".to_string()], + command_cwd.clone(), + permission_profile, + sandbox_cwd, + StdioPolicy::RedirectForShellTool, + env, + ) + .await; + let Ok(mut child) = spawn_result else { + return false; + }; + child + .wait() + .await + .map(|status| status.success()) + .unwrap_or(false) +} + +#[tokio::test] +async fn python_multiprocessing_lock_works_under_sandbox() { + core_test_support::skip_if_sandbox!(); + #[cfg(target_os = "linux")] + let sandbox_env = match linux_sandbox_test_env().await { + Some(env) => env, + // Skip on Linux hosts where Landlock cannot actually be enforced. + None => return, + }; + #[cfg(not(target_os = "linux"))] + let sandbox_env = HashMap::new(); + #[cfg(target_os = "macos")] + let writable_roots = Vec::::new(); + + // From https://man7.org/linux/man-pages/man7/sem_overview.7.html + // + // > On Linux, named semaphores are created in a virtual filesystem, + // > normally mounted under /dev/shm. + #[cfg(target_os = "linux")] + let writable_roots: Vec = vec!["/dev/shm".try_into().unwrap()]; + + let permission_profile = PermissionProfile::workspace_write_with( + &writable_roots, + NetworkSandboxPolicy::Restricted, + /*exclude_tmpdir_env_var*/ true, + /*exclude_slash_tmp*/ true, + ); + + let python_code = r#"import multiprocessing +from multiprocessing import Lock, Process + +def f(lock): + with lock: + print("Lock acquired in child process") + +if __name__ == '__main__': + lock = Lock() + p = Process(target=f, args=(lock,)) + p.start() + p.join() +"#; + + let command_cwd = AbsolutePathBuf::current_dir().expect("should be able to get current dir"); + let sandbox_cwd = command_cwd.clone(); + let mut child = spawn_command_under_sandbox( + vec![ + "python3".to_string(), + "-c".to_string(), + python_code.to_string(), + ], + command_cwd, + &permission_profile, + &sandbox_cwd, + StdioPolicy::Inherit, + sandbox_env, + ) + .await + .expect("should be able to spawn python under sandbox"); + + let status = child.wait().await.expect("should wait for child process"); + assert!(status.success(), "python exited with {status:?}"); +} + +#[tokio::test] +async fn python_getpwuid_works_under_sandbox() { + core_test_support::skip_if_sandbox!(); + #[cfg(target_os = "linux")] + let sandbox_env = match linux_sandbox_test_env().await { + Some(env) => env, + None => return, + }; + #[cfg(not(target_os = "linux"))] + let sandbox_env = HashMap::new(); + + if std::process::Command::new("python3") + .arg("--version") + .status() + .is_err() + { + eprintln!("python3 not found in PATH, skipping test."); + return; + } + + let permission_profile = PermissionProfile::read_only(); + let command_cwd = AbsolutePathBuf::current_dir().expect("should be able to get current dir"); + let sandbox_cwd = command_cwd.clone(); + + let mut child = spawn_command_under_sandbox( + vec![ + "python3".to_string(), + "-c".to_string(), + "import pwd, os; print(pwd.getpwuid(os.getuid()))".to_string(), + ], + command_cwd, + &permission_profile, + &sandbox_cwd, + StdioPolicy::RedirectForShellTool, + sandbox_env, + ) + .await + .expect("should be able to spawn python under sandbox"); + + let status = child + .wait() + .await + .expect("should be able to wait for child process"); + assert!(status.success(), "python exited with {status:?}"); +} + +#[tokio::test] +async fn sandbox_distinguishes_command_and_policy_cwds() { + core_test_support::skip_if_sandbox!(); + #[cfg(target_os = "linux")] + let sandbox_env = match linux_sandbox_test_env().await { + Some(env) => env, + None => return, + }; + #[cfg(not(target_os = "linux"))] + let sandbox_env = HashMap::new(); + let temp = tempfile::tempdir().expect("should be able to create temp dir"); + let sandbox_root = temp.path().join("sandbox"); + let command_root = temp.path().join("command").abs(); + create_dir_all(&sandbox_root).await.expect("mkdir"); + create_dir_all(&command_root).await.expect("mkdir"); + let canonical_sandbox_root = tokio::fs::canonicalize(&sandbox_root) + .await + .expect("canonicalize sandbox root") + .abs(); + let canonical_allowed_path = canonical_sandbox_root.join("allowed.txt"); + + let disallowed_path = command_root.join("forbidden.txt"); + + // Note writable_roots is empty: verify that `canonical_allowed_path` is + // writable only because it is under the sandbox policy cwd, not because it + // is under a writable root. + let permission_profile = PermissionProfile::workspace_write_with( + &[], + NetworkSandboxPolicy::Restricted, + /*exclude_tmpdir_env_var*/ true, + /*exclude_slash_tmp*/ true, + ); + + // Attempt to write inside the command cwd, which is outside of the sandbox policy cwd. + let mut child = spawn_command_under_sandbox( + vec![ + "bash".to_string(), + "-lc".to_string(), + "echo forbidden > forbidden.txt".to_string(), + ], + command_root.clone(), + &permission_profile, + &canonical_sandbox_root, + StdioPolicy::Inherit, + sandbox_env.clone(), + ) + .await + .expect("should spawn command writing to forbidden path"); + + let status = child + .wait() + .await + .expect("should wait for forbidden command"); + assert!( + !status.success(), + "sandbox unexpectedly allowed writing to command cwd: {status:?}" + ); + let forbidden_exists = tokio::fs::try_exists(&disallowed_path) + .await + .expect("try_exists failed"); + assert!( + !forbidden_exists, + "forbidden path should not have been created" + ); + + // Writing to the sandbox policy cwd after changing directories into it should succeed. + let mut child = spawn_command_under_sandbox( + vec![ + "/usr/bin/touch".to_string(), + canonical_allowed_path.to_string_lossy().into_owned(), + ], + command_root, + &permission_profile, + &canonical_sandbox_root, + StdioPolicy::Inherit, + sandbox_env, + ) + .await + .expect("should spawn command writing to sandbox root"); + + let status = child.wait().await.expect("should wait for allowed command"); + assert!( + status.success(), + "sandbox blocked allowed write: {status:?}" + ); + let allowed_exists = tokio::fs::try_exists(&canonical_allowed_path) + .await + .expect("try_exists allowed failed"); + assert!(allowed_exists, "allowed path should exist"); +} + +#[tokio::test] +async fn sandbox_blocks_first_time_dot_codex_creation() { + core_test_support::skip_if_sandbox!(); + #[cfg(target_os = "linux")] + let sandbox_env = match linux_sandbox_test_env().await { + Some(env) => env, + None => return, + }; + #[cfg(not(target_os = "linux"))] + let sandbox_env = HashMap::new(); + + let temp = tempfile::tempdir().expect("should be able to create temp dir"); + let repo_root = temp.path().join("repo").abs(); + create_dir_all(&repo_root).await.expect("mkdir repo"); + let dot_codex = repo_root.join(".codex"); + let config_toml = dot_codex.join("config.toml"); + let permission_profile = PermissionProfile::workspace_write_with( + &[], + NetworkSandboxPolicy::Restricted, + /*exclude_tmpdir_env_var*/ true, + /*exclude_slash_tmp*/ true, + ); + + let mut child = spawn_command_under_sandbox( + vec![ + "bash".to_string(), + "-lc".to_string(), + "mkdir -p .codex && echo 'sandbox_mode = \"danger-full-access\"' > .codex/config.toml" + .to_string(), + ], + repo_root.clone(), + &permission_profile, + &repo_root, + StdioPolicy::RedirectForShellTool, + sandbox_env, + ) + .await + .expect("should spawn command creating .codex"); + + let status = child.wait().await.expect("should wait for .codex command"); + assert!( + !status.success(), + "sandbox unexpectedly allowed first-time .codex creation: {status:?}" + ); + let dot_codex_metadata = tokio::fs::symlink_metadata(&dot_codex).await; + if let Ok(metadata) = dot_codex_metadata { + assert!( + !metadata.is_dir(), + "{} should not be creatable as a directory", + dot_codex.display() + ); + } else if let Err(err) = &dot_codex_metadata { + assert_eq!( + err.kind(), + io::ErrorKind::NotFound, + "unexpected metadata error for {}: {err}", + dot_codex.display() + ); + } + let config_toml_exists = match tokio::fs::try_exists(&config_toml).await { + Ok(exists) => exists, + Err(err) if err.kind() == io::ErrorKind::NotADirectory => false, + Err(err) => panic!("try_exists {} failed: {err}", config_toml.display()), + }; + assert!( + !config_toml_exists, + "{} should not have been created", + config_toml.display() + ); +} + +fn unix_sock_body() { + unsafe { + let mut fds = [0i32; 2]; + let r = libc::socketpair(libc::AF_UNIX, libc::SOCK_DGRAM, 0, fds.as_mut_ptr()); + assert_eq!( + r, + 0, + "socketpair(AF_UNIX, SOCK_DGRAM) failed: {}", + io::Error::last_os_error() + ); + + let msg = b"hello_unix"; + // write() from one end (generic write is allowed) + let sent = libc::write(fds[0], msg.as_ptr() as *const libc::c_void, msg.len()); + assert!(sent >= 0, "write() failed: {}", io::Error::last_os_error()); + + // recvfrom() on the other end. We don’t need the address for socketpair, + // so we pass null pointers for src address. + let mut buf = [0u8; 64]; + let recvd = libc::recvfrom( + fds[1], + buf.as_mut_ptr() as *mut libc::c_void, + buf.len(), + 0, + std::ptr::null_mut(), + std::ptr::null_mut(), + ); + assert!( + recvd >= 0, + "recvfrom() failed: {}", + io::Error::last_os_error() + ); + + let recvd_slice = &buf[..(recvd as usize)]; + assert_eq!( + recvd_slice, + &msg[..], + "payload mismatch: sent {} bytes, got {} bytes", + msg.len(), + recvd + ); + + // Also exercise AF_UNIX stream socketpair quickly to ensure AF_UNIX in general works. + let mut sfds = [0i32; 2]; + let sr = libc::socketpair(libc::AF_UNIX, libc::SOCK_STREAM, 0, sfds.as_mut_ptr()); + assert_eq!( + sr, + 0, + "socketpair(AF_UNIX, SOCK_STREAM) failed: {}", + io::Error::last_os_error() + ); + let snt2 = libc::write(sfds[0], msg.as_ptr() as *const libc::c_void, msg.len()); + assert!( + snt2 >= 0, + "write(stream) failed: {}", + io::Error::last_os_error() + ); + let mut b2 = [0u8; 64]; + let rcv2 = libc::recv(sfds[1], b2.as_mut_ptr() as *mut libc::c_void, b2.len(), 0); + assert!( + rcv2 >= 0, + "recv(stream) failed: {}", + io::Error::last_os_error() + ); + + // Clean up + let _ = libc::close(sfds[0]); + let _ = libc::close(sfds[1]); + let _ = libc::close(fds[0]); + let _ = libc::close(fds[1]); + } +} + +#[tokio::test] +async fn allow_unix_socketpair_recvfrom() { + run_code_under_sandbox( + "allow_unix_socketpair_recvfrom", + &PermissionProfile::read_only(), + || async { unix_sock_body() }, + ) + .await + .expect("should be able to reexec"); +} + +const IN_SANDBOX_ENV_VAR: &str = "IN_SANDBOX"; + +pub async fn run_code_under_sandbox( + test_selector: &str, + permission_profile: &PermissionProfile, + child_body: F, +) -> io::Result> +where + F: FnOnce() -> Fut + Send + 'static, + Fut: Future + Send + 'static, +{ + if std::env::var(IN_SANDBOX_ENV_VAR).is_err() { + let exe = std::env::current_exe()?; + let mut cmds = vec![exe.to_string_lossy().into_owned(), "--exact".into()]; + let mut stdio_policy = StdioPolicy::RedirectForShellTool; + // Allow for us to pass forward --nocapture / use the right stdio policy. + if std::env::args().any(|a| a == "--nocapture") { + cmds.push("--nocapture".into()); + stdio_policy = StdioPolicy::Inherit; + } + cmds.push(test_selector.into()); + + // Your existing launcher: + let command_cwd = + AbsolutePathBuf::current_dir().expect("should be able to get current dir"); + let sandbox_cwd = command_cwd.clone(); + let mut child = spawn_command_under_sandbox( + cmds, + command_cwd, + permission_profile, + &sandbox_cwd, + stdio_policy, + HashMap::from([("IN_SANDBOX".into(), "1".into())]), + ) + .await?; + + let status = child.wait().await?; + Ok(Some(status)) + } else { + // Child branch: run the provided body. + child_body().await; + Ok(None) + } +} diff --git a/codex-rs/exec/tests/suite/seatbelt.rs b/codex-rs/exec/tests/suite/seatbelt.rs new file mode 100644 index 0000000000000000000000000000000000000000..26b287a6259a7f42cf3247e2f3f2a911bae766c9 --- /dev/null +++ b/codex-rs/exec/tests/suite/seatbelt.rs @@ -0,0 +1,880 @@ +use super::sandbox::spawn_command_under_sandbox; +use codex_core::spawn::StdioPolicy; +use codex_protocol::models::PermissionProfile; +use codex_protocol::permissions::FileSystemAccessMode; +use codex_protocol::permissions::FileSystemPath; +use codex_protocol::permissions::FileSystemSandboxEntry; +use codex_protocol::permissions::FileSystemSandboxPolicy; +use codex_protocol::permissions::NetworkSandboxPolicy; +use codex_utils_absolute_path::AbsolutePathBuf; +use pretty_assertions::assert_eq; +use std::collections::HashMap; +use std::fs; +use std::os::unix::fs::symlink; +use std::process::Output; +use std::time::Duration; +use tempfile::TempDir; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn codex_home_symlink_opt_out_uses_loaded_user_config() -> anyhow::Result<()> { + use core_test_support::responses; + use core_test_support::test_codex_exec::test_codex_exec; + use serde_json::json; + + core_test_support::skip_if_sandbox!(Ok(())); + for (home_env, root_alias, enabled, ignore_user_config) in [ + ("user-home", Some("user-home"), true, false), + ("home-link/../.codex", Some(".codex"), false, false), + ("home-link/../.codex", None, true, false), + ("user-home", Some("user-home"), true, true), + ] { + let test = test_codex_exec(); + let home = test.home_path().join(".codex"); + let project_home = test.cwd_path().join(".codex"); + let child = test.home_path().join("child"); + for directory in [&home, &project_home, &child] { + fs::create_dir(directory)?; + } + symlink(&home, test.cwd_path().join("user-home"))?; + symlink(&child, test.cwd_path().join("home-link"))?; + fs::write( + home.join("config.toml"), + format!("allow_symlinked_codex_home = {enabled}\n"), + )?; + // The checkout must not enable the opt-out, even when `..` follows a symlink. + fs::write( + project_home.join("config.toml"), + "allow_symlinked_codex_home = true\n", + )?; + let target = TempDir::new()?; + for root in [&home, &project_home] { + symlink(target.path(), root.join("visualizations"))?; + } + let writable_home = + root_alias.map_or_else(|| home.clone(), |alias| test.cwd_path().join(alias)); + let visualizations = writable_home.join("visualizations"); + let shell_args = json!({ + "cmd": "printf shell > shell.txt", + "workdir": visualizations, + "shell": "/bin/sh", + "login": false, + }) + .to_string(); + let patch = format!( + "*** Begin Patch\n*** Add File: {}/patch.txt\n+patched\n*** End Patch", + visualizations.display() + ); + let server = responses::start_mock_server().await; + let should_run = enabled && !ignore_user_config; + let mock = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_function_call("shell", "exec_command", &shell_args), + responses::ev_apply_patch_custom_tool_call("patch", &patch), + responses::ev_completed("response-1"), + ]), + responses::sse(vec![responses::ev_completed("response-2")]), + ], + ) + .await; + let mut command = test.cmd_with_server(&server); + command + .env_remove("CODEX_API_KEY") + .env_remove("OPENAI_API_KEY") + .env_remove("CODEX_ACCESS_TOKEN") + .env("CODEX_HOME", home_env) + .args(["--skip-git-repo-check", "--sandbox", "workspace-write", "--add-dir"]) + .arg(&visualizations) + .args(["-c", "approvals_reviewer=\"user\"", "-c", "features.shell_snapshot_v2=false"]) + // Session overrides cannot authorize the host's exception either. + .args(["-c", "allow_symlinked_codex_home=true"]) + .args(["-c", "model_provider=\"test\"", "-c"]) + .arg(format!( + "model_providers.test={{name=\"test\",base_url={:?},wire_api=\"responses\",requires_openai_auth=false,supports_websockets=false}}", + format!("{}/v1", server.uri()) + )) + .arg("write the files"); + if ignore_user_config { + command.arg("--ignore-user-config"); + } + let output = command.output()?; + assert!(output.status.success(), "{output:?}"); + assert_eq!(mock.requests().len(), 2); + if should_run { + assert_eq!(fs::read(target.path().join("shell.txt"))?, b"shell"); + assert_eq!(fs::read(target.path().join("patch.txt"))?, b"patched\n"); + } else { + let shell_output = mock + .function_call_output_text("shell") + .expect("shell tool output"); + assert!( + shell_output.contains("symlinked writable roots are not supported"), + "{shell_output}" + ); + assert!(!target.path().join("shell.txt").exists()); + assert!(!target.path().join("patch.txt").exists()); + } + } + Ok(()) +} + +fn workspace_fixture() -> (TempDir, AbsolutePathBuf) { + let temp_dir = TempDir::new().expect("create temporary directory"); + let workspace = temp_dir.path().join("workspace"); + fs::create_dir(&workspace).expect("create workspace"); + let workspace = AbsolutePathBuf::from_absolute_path( + workspace.canonicalize().expect("canonicalize workspace"), + ) + .expect("workspace is absolute"); + (temp_dir, workspace) +} + +fn permission_profile( + workspace: &AbsolutePathBuf, + restricted_paths: Vec<(FileSystemPath, FileSystemAccessMode)>, +) -> PermissionProfile { + let mut file_system_policy = FileSystemSandboxPolicy::read_only(); + file_system_policy.entries.push(FileSystemSandboxEntry::new( + workspace.clone().into(), + FileSystemAccessMode::Write, + )); + file_system_policy.entries.extend( + restricted_paths + .into_iter() + .map(|(path, access)| FileSystemSandboxEntry::new(path, access)), + ); + + PermissionProfile::from_runtime_permissions( + &file_system_policy, + NetworkSandboxPolicy::Restricted, + ) +} + +fn denied_glob_profile(workspace: &AbsolutePathBuf, pattern: String) -> PermissionProfile { + permission_profile( + workspace, + vec![( + FileSystemPath::GlobPattern { pattern }, + FileSystemAccessMode::Deny, + )], + ) +} + +async fn run_sandboxed_command( + command: Vec, + workspace: &AbsolutePathBuf, + permission_profile: &PermissionProfile, +) -> Output { + spawn_command_under_sandbox( + command, + workspace.clone(), + permission_profile, + workspace, + StdioPolicy::RedirectForShellTool, + HashMap::new(), + ) + .await + .expect("spawn command under production Seatbelt policy") + .wait_with_output() + .await + .expect("wait for sandboxed command") +} + +async fn run_path_command( + executable: &str, + path: &AbsolutePathBuf, + workspace: &AbsolutePathBuf, + profile: &PermissionProfile, +) -> Output { + run_sandboxed_command( + vec![executable.to_string(), path.display().to_string()], + workspace, + profile, + ) + .await +} + +async fn assert_denied_read( + path: &AbsolutePathBuf, + workspace: &AbsolutePathBuf, + profile: &PermissionProfile, +) { + let output = run_path_command("/bin/cat", path, workspace, profile).await; + assert!( + !output.status.success(), + "Seatbelt allowed reading {}: {output:?}", + path.display() + ); + assert!(output.stdout.is_empty(), "protected contents leaked"); +} + +fn assert_file_contents(path: &AbsolutePathBuf, contents: &str) { + assert_eq!(fs::read_to_string(path).expect("read fixture"), contents); +} + +async fn run_shell_command( + script: &str, + path: &AbsolutePathBuf, + workspace: &AbsolutePathBuf, + profile: &PermissionProfile, +) -> Output { + run_sandboxed_command( + vec![ + "/bin/sh".to_string(), + "-c".to_string(), + script.to_string(), + "sh".to_string(), + path.display().to_string(), + ], + workspace, + profile, + ) + .await +} + +async fn run_after_policy_applied( + operation: &str, + target: &AbsolutePathBuf, + extra_args: Vec, + workspace: &AbsolutePathBuf, + profile: &PermissionProfile, + setup: impl FnOnce(), +) -> Output { + let label = target + .as_path() + .file_name() + .expect("filename") + .to_string_lossy(); + let ready = workspace.join(format!("{label}.ready.txt")); + let proceed = workspace.join(format!("{label}.proceed.txt")); + let script = format!( + "printf ready > \"$2\"; while [ ! -e \"$3\" ]; do /bin/sleep 0.01; done; {operation}" + ); + let mut command = vec![ + "/bin/sh".to_string(), + "-c".to_string(), + script, + "sh".to_string(), + target.display().to_string(), + ready.display().to_string(), + proceed.display().to_string(), + ]; + command.extend(extra_args); + let child = spawn_command_under_sandbox( + command, + workspace.clone(), + profile, + workspace, + StdioPolicy::RedirectForShellTool, + HashMap::new(), + ) + .await + .expect("spawn command before creating protected fixture"); + + tokio::time::timeout(Duration::from_secs(5), async { + while !ready.as_path().exists() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("sandboxed command did not signal policy application"); + setup(); + fs::write(&proceed, "proceed").expect("release sandboxed command"); + child + .wait_with_output() + .await + .expect("wait for sandboxed command") +} + +#[tokio::test] +async fn seatbelt_enforces_brace_alternation_and_escaped_deny_globs() { + core_test_support::skip_if_sandbox!(); + + let (temp_dir, workspace) = workspace_fixture(); + let dotenv = workspace.join(".env"); + let secrets = workspace.join("secrets.yml"); + let escaped = workspace.join("config?.env"); + let allowed = workspace.join("notes.txt"); + for path in [&dotenv, &secrets, &escaped] { + fs::write(path, "protected contents").expect("write protected fixture"); + } + fs::write(&allowed, "allowed contents").expect("write allowed fixture"); + + let config = format!( + r#"default_permissions = "test" +[permissions.test] +extends = ":read-only" +[permissions.test.filesystem] +"{workspace}" = "write" +"{workspace}/{{.env,secrets.yml}}*" = "deny" +'{workspace}/config\?.env' = "deny" +"#, + workspace = workspace.display() + ); + fs::write(temp_dir.path().join("config.toml"), config).expect("write permissions config"); + let profile = core_test_support::load_default_config_for_test(&temp_dir) + .await + .permissions + .effective_permission_profile(); + + for denied_path in [&dotenv, &secrets, &escaped] { + assert_denied_read(denied_path, &workspace, &profile).await; + } + + let output = run_path_command("/bin/cat", &allowed, &workspace, &profile).await; + assert!(output.status.success(), "allowed read failed: {output:?}"); + assert_eq!(output.stdout, b"allowed contents"); +} + +#[tokio::test] +async fn seatbelt_blocks_renaming_ancestors_of_explicit_protected_files() { + core_test_support::skip_if_sandbox!(); + + for access in [FileSystemAccessMode::Read, FileSystemAccessMode::Deny] { + let (_temp_dir, workspace) = workspace_fixture(); + let private = workspace.join("private"); + let protected = private.join("secret.env"); + let renamed = workspace.join("exposed"); + fs::create_dir(&private).expect("create protected directory"); + fs::write(&protected, "protected contents").expect("write protected fixture"); + + let profile = permission_profile(&workspace, vec![(protected.clone().into(), access)]); + let command = vec![ + "/bin/mv".to_string(), + private.display().to_string(), + renamed.display().to_string(), + ]; + let output = run_sandboxed_command(command.clone(), &workspace, &profile).await; + + assert!( + !output.status.success(), + "Seatbelt allowed renaming an ancestor of a {access:?} path: {output:?}" + ); + assert_file_contents(&protected, "protected contents"); + assert!(!renamed.as_path().exists(), "rename unexpectedly succeeded"); + + let unprotected = permission_profile(&workspace, Vec::new()); + let control = run_sandboxed_command(command, &workspace, &unprotected).await; + assert!( + control.status.success(), + "rename without a protected descendant should succeed: {control:?}" + ); + } +} + +#[tokio::test] +async fn seatbelt_blocks_renaming_dynamic_glob_ancestors() { + core_test_support::skip_if_sandbox!(); + + for component in [".githu?", "{.github,.gitlab}", r".githu\b", "*"] { + let (_temp_dir, workspace) = workspace_fixture(); + let github = workspace.join(".github"); + let workflows = github.join("workflows"); + let protected = workflows.join("release.yml"); + let renamed = workspace.join("exposed"); + fs::create_dir_all(&workflows).expect("create protected directory"); + fs::write(&protected, "protected workflow").expect("write protected fixture"); + + let profile = denied_glob_profile( + &workspace, + format!("{}/{component}/workflows/*.yml", workspace.display()), + ); + let command = vec![ + "/bin/mv".to_string(), + github.display().to_string(), + renamed.display().to_string(), + ]; + let output = run_sandboxed_command(command.clone(), &workspace, &profile).await; + + assert!( + !output.status.success(), + "Seatbelt allowed renaming ancestor matched by {component:?}: {output:?}" + ); + assert_file_contents(&protected, "protected workflow"); + + let unprotected = permission_profile(&workspace, Vec::new()); + let control = run_sandboxed_command(command, &workspace, &unprotected).await; + assert!( + control.status.success(), + "rename without the deny glob should succeed: {control:?}" + ); + } +} + +#[tokio::test] +async fn seatbelt_blocks_renaming_protected_workspace_into_another_writable_root() { + core_test_support::skip_if_sandbox!(); + + let (_temp_dir, workspace) = workspace_fixture(); + let destination = workspace + .parent() + .expect("workspace has a parent") + .join("destination"); + let moved_workspace = destination.join("moved"); + let protected = workspace.join("secret.env"); + fs::create_dir(&destination).expect("create writable destination"); + fs::write(&protected, "protected contents").expect("write protected fixture"); + + let profile = permission_profile( + &workspace, + vec![ + (destination.clone().into(), FileSystemAccessMode::Write), + (protected.clone().into(), FileSystemAccessMode::Read), + ], + ); + let output = run_sandboxed_command( + vec![ + "/bin/mv".to_string(), + workspace.display().to_string(), + moved_workspace.display().to_string(), + ], + &workspace, + &profile, + ) + .await; + + assert!( + !output.status.success(), + "Seatbelt allowed moving the workspace into another writable root: {output:?}" + ); + assert_file_contents(&protected, "protected contents"); + assert!( + !moved_workspace.as_path().exists(), + "workspace unexpectedly moved" + ); +} + +#[tokio::test] +async fn seatbelt_ancestor_protection_does_not_block_unlinking_regular_files() { + core_test_support::skip_if_sandbox!(); + + let (_temp_dir, workspace) = workspace_fixture(); + let regular_file = workspace.join("regular-file"); + let hypothetical_descendant = regular_file.join("secret.env"); + fs::write(®ular_file, "replaceable").expect("write regular file"); + let profile = permission_profile( + &workspace, + vec![(hypothetical_descendant.into(), FileSystemAccessMode::Read)], + ); + + let output = run_path_command("/bin/rm", ®ular_file, &workspace, &profile).await; + + assert!( + output.status.success(), + "directory-only ancestor protection blocked a regular file: {output:?}" + ); + assert!( + !regular_file.as_path().exists(), + "regular file was not removed" + ); +} + +#[tokio::test] +async fn seatbelt_deny_globs_block_writes_to_existing_and_new_files() { + core_test_support::skip_if_sandbox!(); + + let (_temp_dir, workspace) = workspace_fixture(); + let protected = workspace.join("secret.env"); + let allowed = workspace.join("notes.txt"); + fs::write(&protected, "original").expect("write protected fixture"); + let profile = denied_glob_profile(&workspace, format!("{}/*.env", workspace.display())); + + for (operation, script) in [ + ("overwrite", "printf changed >> \"$1\""), + ("truncate", ": > \"$1\""), + ] { + let output = run_shell_command(script, &protected, &workspace, &profile).await; + + assert!( + !output.status.success(), + "Seatbelt allowed {operation} on a denied file: {output:?}" + ); + assert_file_contents(&protected, "original"); + } + + let unlink = run_path_command("/bin/rm", &protected, &workspace, &profile).await; + assert!( + !unlink.status.success(), + "Seatbelt unexpectedly allowed unlinking a denied glob: {unlink:?}" + ); + assert_file_contents(&protected, "original"); + + for (operation, write_script) in [ + ("overwrite", "printf changed >> \"$1\""), + ("truncate", ": > \"$1\""), + ] { + let late_protected = workspace.join(format!("late-{operation}.env")); + let output = run_after_policy_applied( + write_script, + &late_protected, + Vec::new(), + &workspace, + &profile, + || { + fs::write(&late_protected, "created after policy generation") + .expect("create protected fixture after sandbox launch"); + }, + ) + .await; + + assert!( + !output.status.success(), + "Seatbelt allowed {operation} on a file created after policy generation: {output:?}" + ); + assert_file_contents(&late_protected, "created after policy generation"); + } + + let allowed_write = + run_shell_command("printf allowed > \"$1\"", &allowed, &workspace, &profile).await; + assert!( + allowed_write.status.success(), + "unrelated write should remain allowed: {allowed_write:?}" + ); + assert_file_contents(&allowed, "allowed"); +} + +#[tokio::test] +async fn seatbelt_standard_profiles_preserve_expected_filesystem_access() { + core_test_support::skip_if_sandbox!(); + let (_temp_dir, workspace) = workspace_fixture(); + let inside = workspace.join("notes.txt"); + let outside = workspace.parent().expect("parent").join("outside.txt"); + fs::write(&inside, "original").expect("write workspace fixture"); + let profile = permission_profile(&workspace, Vec::new()); + let script = "set -e; /bin/cat \"$1\" >/dev/null; printf changed >> \"$1\"; /bin/mkdir \"$1.dir\"; /bin/mv \"$1\" \"$1.dir/moved\"; /bin/rm \"$1.dir/moved\"; /bin/rmdir \"$1.dir\""; + let inside_output = run_shell_command(script, &inside, &workspace, &profile).await; + assert!( + inside_output.status.success(), + "ordinary workspace operations failed: {inside_output:?}" + ); + + let outside_output = + run_shell_command("printf denied > \"$1\"", &outside, &workspace, &profile).await; + assert!( + !outside_output.status.success(), + "write outside the workspace unexpectedly succeeded: {outside_output:?}" + ); + assert!(!outside.as_path().exists(), "outside file was created"); + let readable = workspace.join("readable.txt"); + fs::write(&readable, "readable contents").expect("write read-only fixture"); + let profile = PermissionProfile::read_only(); + let read = run_path_command("/bin/cat", &readable, &workspace, &profile).await; + assert!(read.status.success(), "ordinary read failed: {read:?}"); + assert_eq!(read.stdout, b"readable contents"); + + let write = run_shell_command("printf changed > \"$1\"", &readable, &workspace, &profile).await; + assert!( + !write.status.success(), + "read-only write succeeded: {write:?}" + ); + assert_file_contents(&readable, "readable contents"); +} + +#[tokio::test] +async fn seatbelt_blocks_protected_reads_and_writes_through_symlink_aliases() { + core_test_support::skip_if_sandbox!(); + + let (_temp_dir, workspace) = workspace_fixture(); + let private = workspace.join("private"); + let alias = workspace.join("alias"); + let protected = private.join("secret.env"); + let alias_path = alias.join("secret.env"); + fs::create_dir(&private).expect("create protected directory"); + fs::write(&protected, "protected contents").expect("write protected fixture"); + symlink(private.as_path(), alias.as_path()).expect("create directory symlink"); + + for (restricted_path, accessed_path) in [ + (FileSystemPath::from(protected.clone()), alias_path.clone()), + ( + FileSystemPath::GlobPattern { + pattern: format!("{}/*.env", private.display()), + }, + alias_path.clone(), + ), + ( + FileSystemPath::GlobPattern { + pattern: format!("{}/*.env", alias.display()), + }, + protected.clone(), + ), + ] { + let profile = permission_profile( + &workspace, + vec![(restricted_path, FileSystemAccessMode::Deny)], + ); + assert_denied_read(&accessed_path, &workspace, &profile).await; + + let write = run_shell_command( + "printf changed > \"$1\"", + &accessed_path, + &workspace, + &profile, + ) + .await; + assert!( + !write.status.success(), + "Seatbelt allowed writing through protected symlink alias: {write:?}" + ); + assert_file_contents(&protected, "protected contents"); + } +} + +#[tokio::test] +async fn seatbelt_globstar_protects_files_and_ancestors_at_every_depth() { + core_test_support::skip_if_sandbox!(); + + let (_temp_dir, workspace) = workspace_fixture(); + let first = workspace.join("first"); + let second = first.join("second"); + fs::create_dir_all(&second).expect("create nested fixture directories"); + let protected_paths = [ + workspace.join("root.env"), + first.join("first.env"), + second.join("second.env"), + ]; + for path in &protected_paths { + fs::write(path, "protected contents").expect("write protected fixture"); + } + let profile = denied_glob_profile(&workspace, format!("{}/**/*.env", workspace.display())); + + for path in &protected_paths { + assert_denied_read(path, &workspace, &profile).await; + } + + for ancestor in [&first, &second] { + let destination = ancestor.with_extension("renamed"); + let output = run_sandboxed_command( + vec![ + "/bin/mv".to_string(), + ancestor.display().to_string(), + destination.display().to_string(), + ], + &workspace, + &profile, + ) + .await; + assert!( + !output.status.success(), + "globstar allowed renaming protected ancestor {}: {output:?}", + ancestor.display() + ); + } +} + +#[tokio::test] +async fn seatbelt_brace_alternatives_with_slashes_protect_each_branch() { + core_test_support::skip_if_sandbox!(); + + let (_temp_dir, workspace) = workspace_fixture(); + let nested = workspace.join("private/nested"); + let other = workspace.join("other"); + let allowed = workspace.join("notes.txt"); + fs::create_dir_all(&nested).expect("create nested protected directory"); + fs::create_dir(&other).expect("create alternate protected directory"); + let protected_paths = [nested.join("secret.env"), other.join("secret.env")]; + for path in &protected_paths { + fs::write(path, "protected contents").expect("write protected fixture"); + } + fs::write(&allowed, "allowed contents").expect("write allowed fixture"); + let profile = denied_glob_profile( + &workspace, + format!("{}/{{private/nested,other}}/*.env", workspace.display()), + ); + + for path in &protected_paths { + assert_denied_read(path, &workspace, &profile).await; + } + + for ancestor in [&workspace.join("private"), &nested, &other] { + let destination = ancestor.with_extension("renamed"); + let output = run_sandboxed_command( + vec![ + "/bin/mv".to_string(), + ancestor.display().to_string(), + destination.display().to_string(), + ], + &workspace, + &profile, + ) + .await; + assert!( + !output.status.success(), + "brace alternative allowed renaming {}: {output:?}", + ancestor.display() + ); + } + + let output = run_path_command("/bin/cat", &allowed, &workspace, &profile).await; + assert!( + output.status.success(), + "brace alternatives produced an invalid or overly broad policy: {output:?}" + ); + assert_eq!(output.stdout, b"allowed contents"); +} + +#[tokio::test] +async fn seatbelt_protects_directory_trees_created_after_policy_application() { + core_test_support::skip_if_sandbox!(); + + for operation in ["read", "rename"] { + let (_temp_dir, workspace) = workspace_fixture(); + let late_directory = workspace.join("late/deep"); + let protected = late_directory.join("secret.env"); + let destination = workspace.join("exposed"); + let profile = denied_glob_profile(&workspace, format!("{}/**/*.env", workspace.display())); + let operation_script = match operation { + "read" => "/bin/cat \"$1\"", + "rename" => "/bin/mv \"$4\" \"$5\"", + _ => unreachable!("unknown operation"), + }; + let output = run_after_policy_applied( + operation_script, + &protected, + vec![ + workspace.join("late").display().to_string(), + destination.display().to_string(), + ], + &workspace, + &profile, + || { + fs::create_dir_all(&late_directory) + .expect("create protected directory after policy"); + fs::write(&protected, "protected contents").expect("write protected fixture"); + }, + ) + .await; + + assert!( + !output.status.success(), + "Seatbelt allowed {operation} after protected directory appeared: {output:?}" + ); + assert_file_contents(&protected, "protected contents"); + } +} + +#[tokio::test] +async fn seatbelt_blocks_atomic_rename_exchange_and_destination_replacement() { + core_test_support::skip_if_sandbox!(); + + for (operation, flags) in [("exchange", libc::RENAME_SWAP), ("replace", 0)] { + let (_temp_dir, workspace) = workspace_fixture(); + let private = workspace.join("private"); + let destination = workspace.join("destination"); + let protected = private.join("secret.env"); + fs::create_dir(&private).expect("create protected directory"); + fs::create_dir(&destination).expect("create replacement directory"); + fs::write(&protected, "protected contents").expect("write protected fixture"); + let profile = permission_profile( + &workspace, + vec![(protected.clone().into(), FileSystemAccessMode::Read)], + ); + let script = "import ctypes, os, sys; libc = ctypes.CDLL(None, use_errno=True); result = libc.renameatx_np(-2, os.fsencode(sys.argv[1]), -2, os.fsencode(sys.argv[2]), int(sys.argv[3])); sys.exit(0 if result == 0 else ctypes.get_errno())"; + let command = vec![ + "/usr/bin/python3".to_string(), + "-c".to_string(), + script.to_string(), + private.display().to_string(), + destination.display().to_string(), + flags.to_string(), + ]; + let output = run_sandboxed_command(command.clone(), &workspace, &profile).await; + + assert!( + !output.status.success(), + "Seatbelt allowed atomic {operation} of a protected ancestor: {output:?}" + ); + assert_file_contents(&protected, "protected contents"); + + let unprotected = permission_profile(&workspace, Vec::new()); + let control = run_sandboxed_command(command, &workspace, &unprotected).await; + assert!( + control.status.success(), + "unprotected atomic {operation} should succeed: {control:?}" + ); + } +} + +#[tokio::test] +async fn seatbelt_blocks_protected_ancestor_moves_across_writable_roots() { + core_test_support::skip_if_sandbox!(); + + for destination_is_nested in [false, true] { + for use_glob in [false, true] { + let (_temp_dir, workspace) = workspace_fixture(); + let private = workspace.join("private"); + let protected = private.join("secret.env"); + let destination = if destination_is_nested { + workspace.join("destination") + } else { + workspace + .parent() + .expect("workspace has a parent") + .join("destination") + }; + let moved = destination.join("moved"); + fs::create_dir(&private).expect("create protected directory"); + fs::create_dir(&destination).expect("create writable destination"); + fs::write(&protected, "protected contents").expect("write protected fixture"); + let protected_path = if use_glob { + FileSystemPath::GlobPattern { + pattern: format!("{}/*.env", private.display()), + } + } else { + protected.clone().into() + }; + let profile = permission_profile( + &workspace, + vec![ + (destination.clone().into(), FileSystemAccessMode::Write), + (protected_path, FileSystemAccessMode::Deny), + ], + ); + + let output = run_sandboxed_command( + vec![ + "/bin/mv".to_string(), + private.display().to_string(), + moved.display().to_string(), + ], + &workspace, + &profile, + ) + .await; + assert!( + !output.status.success(), + "protected ancestor crossed writable roots: nested={destination_is_nested}, glob={use_glob}, output={output:?}" + ); + assert_file_contents(&protected, "protected contents"); + assert!(!moved.as_path().exists(), "protected ancestor was moved"); + } + } +} + +#[tokio::test] +async fn seatbelt_blocks_creating_files_matched_by_denied_globs() { + core_test_support::skip_if_sandbox!(); + + let (_temp_dir, workspace) = workspace_fixture(); + let nested = workspace.join("nested"); + fs::create_dir(&nested).expect("create nested directory"); + let profile = denied_glob_profile(&workspace, format!("{}/**/*.env", workspace.display())); + + for path in [workspace.join("created.env"), nested.join("created.env")] { + let output = + run_shell_command("printf created > \"$1\"", &path, &workspace, &profile).await; + assert!( + !output.status.success(), + "Seatbelt allowed creation at a denied path: {output:?}" + ); + assert!(!path.as_path().exists(), "denied file was created"); + } + + let allowed = nested.join("created.txt"); + let output = run_path_command("/usr/bin/touch", &allowed, &workspace, &profile).await; + assert!( + output.status.success(), + "unrelated file creation should succeed: {output:?}" + ); + assert!(allowed.as_path().exists(), "unrelated file was not created"); +} diff --git a/codex-rs/exec/tests/suite/server_error_exit.rs b/codex-rs/exec/tests/suite/server_error_exit.rs new file mode 100644 index 0000000000000000000000000000000000000000..bd0fad698758951ea41d233f2580fc9d79a050d0 --- /dev/null +++ b/codex-rs/exec/tests/suite/server_error_exit.rs @@ -0,0 +1,33 @@ +#![cfg(not(target_os = "windows"))] +#![allow(clippy::unwrap_used)] + +use core_test_support::responses; +use core_test_support::test_codex_exec::test_codex_exec; + +/// Verify that when the server reports an error, `codex-exec` exits with a +/// non-zero status code so automation can detect failures. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exits_non_zero_when_server_reports_error() -> anyhow::Result<()> { + let test = test_codex_exec(); + + // Mock a simple Responses API SSE stream that immediately reports a + // `response.failed` event with an error message. + let server = responses::start_mock_server().await; + let body = responses::sse(vec![serde_json::json!({ + "type": "response.failed", + "response": { + "id": "resp_err_1", + "error": {"code": "rate_limit_exceeded", "message": "synthetic server error"} + } + })]); + responses::mount_sse_once(&server, body).await; + + test.cmd_with_server(&server) + .arg("--skip-git-repo-check") + .arg("tell me something") + .arg("--experimental-json") + .assert() + .code(1); + + Ok(()) +} diff --git a/codex-rs/exec/tests/suite/worktree.rs b/codex-rs/exec/tests/suite/worktree.rs new file mode 100644 index 0000000000000000000000000000000000000000..d2c7afd5d9f25837cbe085a68829ebd503d7ba1d --- /dev/null +++ b/codex-rs/exec/tests/suite/worktree.rs @@ -0,0 +1,441 @@ +//! Exercises gated shared-pool execution and preserves existing checkout ownership. +//! Observe ownership at the first model request, before the process can finish. + +use anyhow::Context; +use codex_utils_absolute_path::AbsolutePathBuf; +use codex_worktree::CreateWorktree; +use codex_worktree::WorktreeManager; +use codex_worktree::WorktreeSettings; +use core_test_support::responses; +use core_test_support::test_codex_exec::test_codex_exec; +use pretty_assertions::assert_eq; +use serde_json::Value; +use std::fs; +use std::path::Path; +use std::process::Command; +use std::process::Output; +use std::sync::Arc; +use std::sync::Mutex; + +fn git(cwd: &Path, args: &[&str]) -> anyhow::Result<()> { + let output = Command::new("git") + .current_dir(cwd) + .args([ + "-c", + "user.name=Worktree Test", + "-c", + "user.email=test@example.invalid", + ]) + .args(args) + .output()?; + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + Ok(()) +} + +fn started_thread(output: Output) -> anyhow::Result { + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let events = String::from_utf8(output.stdout)? + .lines() + .map(serde_json::from_str::) + .collect::, _>>()?; + events + .iter() + .find(|event| event["type"] == "thread.started") + .and_then(|event| event["thread_id"].as_str()) + .map(str::to_owned) + .context("exec should report its thread id") +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn worktree_start_and_fork_use_host_pool_and_preserve_legacy_resume() -> anyhow::Result<()> { + let test = test_codex_exec(); + let home = AbsolutePathBuf::from_absolute_path(test.home_path())? + .canonicalize()? + .into_path_buf(); + let source = AbsolutePathBuf::from_absolute_path(test.cwd_path())? + .canonicalize()? + .into_path_buf(); + let extra = source.join("extra"); + fs::create_dir(&extra)?; + fs::write(extra.join("tracked.txt"), "tracked")?; + fs::write(source.join("AGENTS.md"), "managed checkout instructions")?; + git(&source, &["init", "--quiet"])?; + git(&source, &["add", "."])?; + git( + &source, + &["commit", "--quiet", "--no-gpg-sign", "-m", "initial"], + )?; + let pool = home.join("shared-worktrees"); + let other_pool = serde_json::to_string(&home.join("session-only-pool"))?; + fs::write( + home.join("config.toml"), + format!( + "sandbox_mode = \"workspace-write\"\nfeatures.worktrees = true\n[desktop]\ngit-worktree-root = {}\n[projects.{}]\ntrust_level = \"trusted\"\n", + serde_json::to_string(&pool)?, + serde_json::to_string(&source)?, + ), + )?; + fs::write( + home.join("demo.config.toml"), + format!("features.worktrees = false\n[desktop]\ngit-worktree-root = {other_pool}\n"), + )?; + let mut settings = WorktreeSettings::for_cli(&home, /*desktop*/ None)?; + settings.root = pool; + let legacy_manager = WorktreeManager::new(WorktreeSettings { + root: home.join("worktrees-cli"), + ..settings.clone() + }); + let legacy = legacy_manager.create(&CreateWorktree { + source_cwd: source.clone(), + base: None, + })?; + let manager = WorktreeManager::new(settings); + let body = responses::sse(vec![ + responses::ev_response_created("response_1"), + responses::ev_assistant_message("message_1", "done"), + responses::ev_completed("response_1"), + ]); + let server = responses::start_mock_server().await; + let seed_mock = responses::mount_sse_once(&server, body.clone()).await; + let legacy_id = started_thread( + test.cmd_with_server(&server) + .current_dir(legacy.cwd.join("extra")) + .args(["--json", "legacy seed"]) + .output()?, + )?; + seed_mock.single_request(); + legacy_manager.bind_thread(&legacy.root, &legacy_id)?; + fs::write(source.join("AGENTS.md"), "uncommitted source instructions")?; + fs::create_dir(source.join(".codex"))?; + fs::write(source.join(".codex/config.toml"), "[invalid source config")?; + + let launcher = tempfile::tempdir()?; + fs::create_dir(launcher.path().join("extra"))?; + let mut known_roots = Vec::new(); + for (fork, overrides) in [ + (false, Vec::new()), + ( + true, + vec![ + "--profile".to_owned(), + "demo".to_owned(), + "-c".to_owned(), + "features.worktrees=true".to_owned(), + "-c".to_owned(), + format!("desktop.git-worktree-root={other_pool}"), + "-c".to_owned(), + format!("desktop={{git-worktree-root={other_pool}}}"), + ], + ), + (true, vec!["--cd".to_owned(), extra.display().to_string()]), + ] { + let explicit_cd = overrides.first().is_some_and(|arg| arg == "--cd"); + let server = responses::start_mock_server().await; + let observed_owner = Arc::new(Mutex::>::default()); + let captured_owner = Arc::clone(&observed_owner); + let observed_manager = manager.clone(); + let observed_source = source.clone(); + let previous_roots = known_roots.clone(); + // A matcher is needed here to inspect metadata at request time, not after exit. + let response_mock = responses::mount_sse_once_match( + &server, + move |_: &wiremock::Request| { + let owner = observed_manager + .list(&observed_source) + .expect("list shared checkouts") + .into_iter() + .find(|checkout| !previous_roots.contains(&checkout.root)) + .and_then(|checkout| { + observed_manager.owner(&checkout.root).expect("read owner") + }); + *captured_owner.lock().expect("capture request owner") = owner; + true + }, + body.clone(), + ) + .await; + let mut command = test.cmd_with_server(&server); + command + .args(["--json", "--worktree", "--add-dir", "extra"]) + .args(overrides); + if fork { + command + .current_dir(launcher.path()) + .arg("fork") + .arg(&legacy_id); + } + let thread_id = started_thread(command.arg("describe checkout").output()?)?; + assert_eq!( + *observed_owner.lock().expect("read captured owner"), + Some(thread_id) + ); + let checkout = manager + .list(&source)? + .into_iter() + .find(|checkout| !known_roots.contains(&checkout.root)) + .context("one new shared checkout")?; + let request = response_mock.single_request(); + let context = request + .message_input_texts("user") + .into_iter() + .chain(request.message_input_texts("developer")) + .collect::>() + .join("\n"); + let checkout_cwd = AbsolutePathBuf::from_absolute_path(if explicit_cd { + checkout.root.join("extra") + } else { + checkout.cwd.clone() + })? + .canonicalize()?; + assert!( + context.contains(checkout_cwd.to_string_lossy().as_ref()), + "{context}" + ); + assert!( + context.contains( + if fork { + launcher.path().join("extra") + } else { + extra.clone() + } + .to_string_lossy() + .as_ref() + ), + "{context}" + ); + assert!( + context.contains("managed checkout instructions"), + "{context}" + ); + assert!( + !context.contains("uncommitted source instructions"), + "{context}" + ); + known_roots.push(checkout.root); + } + + let server = responses::start_mock_server().await; + let resume_mock = responses::mount_sse_once(&server, body).await; + let resumed_id = started_thread( + test.cmd_with_server(&server) + .current_dir(legacy.cwd.join("extra")) + .args(["--json", "resume", &legacy_id, "continue legacy"]) + .output()?, + )?; + assert_eq!(resumed_id, legacy_id); + assert_eq!(legacy_manager.owner(&legacy.root)?, Some(legacy_id.clone())); + let legacy_cwd = + AbsolutePathBuf::from_absolute_path(legacy.cwd.join("extra"))?.canonicalize()?; + assert!( + resume_mock + .single_request() + .message_input_texts("user") + .iter() + .any(|text| text.contains(legacy_cwd.to_string_lossy().as_ref())) + ); + fs::write( + home.join("config.toml"), + format!( + "features.worktrees=true\n[projects.{}]\ntrust_level=\"untrusted\"\n", + serde_json::to_string(&source)?, + ), + )?; + for args in [vec!["prompt"], vec!["fork", legacy_id.as_str(), "prompt"]] { + let output = test.cmd().arg("--worktree").args(args).output()?; + assert!(!output.status.success()); + assert!(String::from_utf8_lossy(&output.stderr).contains("explicitly untrusted")); + } + assert!(!home.join("worktrees").exists()); + + // Synthetic enterprise claims let the cloud bundle, not a local project entry, + // supply the source's distrust decision. + fs::write( + home.join("auth.json"), + serde_json::to_vec(&serde_json::json!({ + "auth_mode": "chatgpt", + "tokens": { + "id_token": "e30.eyJodHRwczovL2FwaS5vcGVuYWkuY29tL2F1dGgiOnsiY2hhdGdwdF9wbGFuX3R5cGUiOiJlbnRlcnByaXNlIiwiY2hhdGdwdF9hY2NvdW50X2lkIjoid29ya3NwYWNlLTEyMyIsImNoYXRncHRfdXNlcl9pZCI6InVzZXItMTIzIn19.signature", + "access_token": "test-access-token", + "refresh_token": "test-refresh-token", + "account_id": "workspace-123" + } + }))?, + )?; + fs::write( + home.join("config.toml"), + format!( + "features.worktrees=true\ncli_auth_credentials_store=\"file\"\nchatgpt_base_url=\"{}/backend-api\"\n", + server.uri(), + ), + )?; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/backend-api/wham/config/bundle")) + .respond_with( + wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "config_toml": { "enterprise_managed": [{ + "id": "source-trust", "name": "Source trust", + "contents": format!( + "[projects.{}]\ntrust_level=\"untrusted\"\n[projects.{}]\ntrust_level=\"untrusted\"\n", + serde_json::to_string(&source)?, + serde_json::to_string(&legacy.cwd.join("extra"))?, + ) + }] } + })), + ) + .mount(&server) + .await; + for args in [vec!["prompt"], vec!["fork", legacy_id.as_str(), "prompt"]] { + let request_count = if args.first() == Some(&"fork") { + // The saved checkout can choose file auth even when the launcher does not. + fs::remove_file(source.join(".codex/config.toml"))?; + fs::create_dir_all(legacy.cwd.join(".codex"))?; + fs::write( + legacy.cwd.join(".codex/config.toml"), + "cli_auth_credentials_store=\"file\"\n", + )?; + git(&legacy.cwd, &["add", ".codex/config.toml"])?; + git( + &legacy.cwd, + &["commit", "--quiet", "--no-gpg-sign", "-m", "file auth"], + )?; + fs::write( + home.join("config.toml"), + format!( + "features.worktrees=true\ncli_auth_credentials_store=\"ephemeral\"\nchatgpt_base_url=\"{}/backend-api\"\n[projects.{}]\ntrust_level=\"trusted\"\n", + server.uri(), + serde_json::to_string(&source)?, + ), + )?; + let cache = home.join("cloud-config-bundle-cache.json"); + if cache.exists() { + fs::remove_file(&cache)?; + } + Some(server.received_requests().await.context("requests")?.len()) + } else { + None + }; + let output = test + .cmd() + .current_dir(if request_count.is_some() { + launcher.path() + } else { + &source + }) + .env_remove("CODEX_API_KEY") + .env_remove("OPENAI_API_KEY") + .env_remove("CODEX_ACCESS_TOKEN") + .args(["--json", "--worktree", "--strict-config"]) + .args(args) + .output()?; + assert!( + !output.status.success() + && String::from_utf8_lossy(&output.stderr).contains("explicitly untrusted"), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + if let Some(request_count) = request_count { + assert!(!String::from_utf8_lossy(&output.stdout).contains("thread.started")); + let denied_manager = + WorktreeManager::new(WorktreeSettings::for_cli(&home, /*desktop*/ None)?); + let checkout = denied_manager + .list(&legacy.cwd)? + .into_iter() + .next() + .context("denied checkout remains for manual removal")?; + assert_eq!(denied_manager.owner(&checkout.root)?, None); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + stderr.contains( + checkout + .root + .file_name() + .context("checkout name")? + .to_string_lossy() + .as_ref() + ) + ); + assert!(stderr.contains("Remove it manually with `git worktree remove`")); + assert!( + server + .received_requests() + .await + .context("requests")? + .iter() + .skip(request_count) + .any(|request| request.url.path() == "/backend-api/wham/config/bundle") + ); + } else { + assert!(!home.join("worktrees").exists()); + } + } + Ok(()) +} + +#[test] +fn worktree_rejects_remote_defaults_before_allocation() -> anyhow::Result<()> { + for route in ["environment", "file"] { + let test = test_codex_exec(); + let mut command = test.cmd(); + if route == "environment" { + command.env("CODEX_EXEC_SERVER_URL", "ws://127.0.0.1:9"); + } else { + fs::write( + test.home_path().join("environments.toml"), + "default = 'remote'\n[[environments]]\nid = 'remote'\nurl = 'ws://127.0.0.1:9'\n", + )?; + } + let output = command + .args(["--worktree", "-c", "features.worktrees=true", "prompt"]) + .output()?; + assert!(!output.status.success(), "{route}"); + assert!( + String::from_utf8_lossy(&output.stderr).contains("--worktree requires local execution"), + "{route}: {}", + String::from_utf8_lossy(&output.stderr) + ); + assert!(!test.home_path().join("worktrees").exists(), "{route}"); + } + Ok(()) +} + +#[test] +fn worktree_rejects_disabled_features_and_ignored_config_before_allocation() -> anyhow::Result<()> { + for (config, args, expected) in [ + ( + "", + vec!["--ignore-user-config", "prompt"], + "--worktree cannot be combined with --ignore-user-config", + ), + ( + "features.worktrees = false", + vec!["prompt"], + "--enable worktrees", + ), + ( + "features.worktrees = false", + vec!["fork", "missing"], + "--enable worktrees", + ), + ( + "features.worktrees = true", + vec!["-c", "features.worktrees=false", "prompt"], + "--enable worktrees", + ), + ] { + let test = test_codex_exec(); + fs::write(test.home_path().join("config.toml"), config)?; + let output = test.cmd().arg("--worktree").args(args).output()?; + assert!(!output.status.success()); + assert!(String::from_utf8_lossy(&output.stderr).contains(expected)); + assert!(!test.home_path().join("worktrees").exists()); + } + Ok(()) +} diff --git a/codex-rs/execpolicy/examples/example.codexpolicy b/codex-rs/execpolicy/examples/example.codexpolicy new file mode 100644 index 0000000000000000000000000000000000000000..4040f612092ece72676032dcc38b659ee2b4069b --- /dev/null +++ b/codex-rs/execpolicy/examples/example.codexpolicy @@ -0,0 +1,78 @@ + +# Example policy to illustrate syntax; not comprehensive and not recommended for actual use. + +prefix_rule( + pattern = ["git", "reset", "--hard"], + decision = "forbidden", + justification = "destructive operation", + match = [ + ["git", "reset", "--hard"], + ], + not_match = [ + ["git", "reset", "--keep"], + "git reset --merge", + ], +) + +prefix_rule( + pattern = ["ls"], + match = [ + ["ls"], + ["ls", "-l"], + ["ls", "-a", "."], + ], +) + +prefix_rule( + pattern = ["cat"], + match = [ + ["cat", "file.txt"], + ["cat", "-n", "README.md"], + ], +) + +prefix_rule( + pattern = ["cp"], + decision = "prompt", + match = [ + ["cp", "foo", "bar"], + "cp -r src dest", + ], +) + +prefix_rule( + pattern = ["head"], + match = [ + ["head", "README.md"], + ["head", "-n", "5", "CHANGELOG.md"], + ], + not_match = [ + ["hea", "-n", "1,5p", "CHANGELOG.md"], + ], +) + +prefix_rule( + pattern = ["printenv"], + match = [ + ["printenv"], + ["printenv", "PATH"], + ], + not_match = [ + ["print", "-0"], + ], +) + +prefix_rule( + pattern = ["pwd"], + match = [ + ["pwd"], + ], +) + +prefix_rule( + pattern = ["which"], + match = [ + ["which", "python3"], + ["which", "-a", "python3"], + ], +) diff --git a/codex-rs/execpolicy/src/amend.rs b/codex-rs/execpolicy/src/amend.rs new file mode 100644 index 0000000000000000000000000000000000000000..e25fd1bd91440fb349e3979151aba62eb5c16ed8 --- /dev/null +++ b/codex-rs/execpolicy/src/amend.rs @@ -0,0 +1,337 @@ +use std::fs::OpenOptions; +use std::io::Read; +use std::io::Seek; +use std::io::SeekFrom; +use std::io::Write; +use std::path::Path; +use std::path::PathBuf; + +use crate::decision::Decision; +use crate::rule::NetworkRuleProtocol; +use crate::rule::normalize_network_rule_host; +use thiserror::Error; + +#[derive(Debug, Error)] +pub enum AmendError { + #[error("prefix rule requires at least one token")] + EmptyPrefix, + #[error("invalid network rule: {0}")] + InvalidNetworkRule(String), + #[error("policy path has no parent: {path}")] + MissingParent { path: PathBuf }, + #[error("failed to create policy directory {dir}: {source}")] + CreatePolicyDir { + dir: PathBuf, + source: std::io::Error, + }, + #[error("failed to format prefix tokens: {source}")] + SerializePrefix { source: serde_json::Error }, + #[error("failed to serialize network rule field: {source}")] + SerializeNetworkRule { source: serde_json::Error }, + #[error("failed to open policy file {path}: {source}")] + OpenPolicyFile { + path: PathBuf, + source: std::io::Error, + }, + #[error("failed to write to policy file {path}: {source}")] + WritePolicyFile { + path: PathBuf, + source: std::io::Error, + }, + #[error("failed to lock policy file {path}: {source}")] + LockPolicyFile { + path: PathBuf, + source: std::io::Error, + }, + #[error("failed to seek policy file {path}: {source}")] + SeekPolicyFile { + path: PathBuf, + source: std::io::Error, + }, + #[error("failed to read policy file {path}: {source}")] + ReadPolicyFile { + path: PathBuf, + source: std::io::Error, + }, + #[error("failed to read metadata for policy file {path}: {source}")] + PolicyMetadata { + path: PathBuf, + source: std::io::Error, + }, +} + +/// Note this thread uses advisory file locking and performs blocking I/O, so it should be used with +/// [`tokio::task::spawn_blocking`] when called from an async context. +pub fn blocking_append_allow_prefix_rule( + policy_path: &Path, + prefix: &[String], +) -> Result<(), AmendError> { + if prefix.is_empty() { + return Err(AmendError::EmptyPrefix); + } + + let tokens = prefix + .iter() + .map(serde_json::to_string) + .collect::, _>>() + .map_err(|source| AmendError::SerializePrefix { source })?; + let pattern = format!("[{}]", tokens.join(", ")); + let rule = format!(r#"prefix_rule(pattern={pattern}, decision="allow")"#); + append_rule_line(policy_path, &rule) +} + +/// Note this function uses advisory file locking and performs blocking I/O, so it should be used +/// with [`tokio::task::spawn_blocking`] when called from an async context. +pub fn blocking_append_network_rule( + policy_path: &Path, + host: &str, + protocol: NetworkRuleProtocol, + decision: Decision, + justification: Option<&str>, +) -> Result<(), AmendError> { + let host = normalize_network_rule_host(host) + .map_err(|err| AmendError::InvalidNetworkRule(err.to_string()))?; + if let Some(raw) = justification + && raw.trim().is_empty() + { + return Err(AmendError::InvalidNetworkRule( + "justification cannot be empty".to_string(), + )); + } + + let host = serde_json::to_string(&host) + .map_err(|source| AmendError::SerializeNetworkRule { source })?; + let protocol = serde_json::to_string(protocol.as_policy_string()) + .map_err(|source| AmendError::SerializeNetworkRule { source })?; + let decision = serde_json::to_string(match decision { + Decision::Allow => "allow", + Decision::Prompt => "prompt", + Decision::Forbidden => "deny", + }) + .map_err(|source| AmendError::SerializeNetworkRule { source })?; + + let mut args = vec![ + format!("host={host}"), + format!("protocol={protocol}"), + format!("decision={decision}"), + ]; + if let Some(justification) = justification { + let justification = serde_json::to_string(justification) + .map_err(|source| AmendError::SerializeNetworkRule { source })?; + args.push(format!("justification={justification}")); + } + let rule = format!("network_rule({})", args.join(", ")); + append_rule_line(policy_path, &rule) +} + +fn append_rule_line(policy_path: &Path, rule: &str) -> Result<(), AmendError> { + let dir = policy_path + .parent() + .ok_or_else(|| AmendError::MissingParent { + path: policy_path.to_path_buf(), + })?; + match std::fs::create_dir(dir) { + Ok(()) => {} + Err(ref source) if source.kind() == std::io::ErrorKind::AlreadyExists => {} + Err(source) => { + return Err(AmendError::CreatePolicyDir { + dir: dir.to_path_buf(), + source, + }); + } + } + + append_locked_line(policy_path, rule) +} + +fn append_locked_line(policy_path: &Path, line: &str) -> Result<(), AmendError> { + let mut file = OpenOptions::new() + .create(true) + .read(true) + .append(true) + .open(policy_path) + .map_err(|source| AmendError::OpenPolicyFile { + path: policy_path.to_path_buf(), + source, + })?; + file.lock().map_err(|source| AmendError::LockPolicyFile { + path: policy_path.to_path_buf(), + source, + })?; + + file.seek(SeekFrom::Start(0)) + .map_err(|source| AmendError::SeekPolicyFile { + path: policy_path.to_path_buf(), + source, + })?; + let mut contents = String::new(); + file.read_to_string(&mut contents) + .map_err(|source| AmendError::ReadPolicyFile { + path: policy_path.to_path_buf(), + source, + })?; + + if contents.lines().any(|existing| existing == line) { + return Ok(()); + } + + if !contents.is_empty() && !contents.ends_with('\n') { + file.write_all(b"\n") + .map_err(|source| AmendError::WritePolicyFile { + path: policy_path.to_path_buf(), + source, + })?; + } + + file.write_all(format!("{line}\n").as_bytes()) + .map_err(|source| AmendError::WritePolicyFile { + path: policy_path.to_path_buf(), + source, + })?; + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + use tempfile::tempdir; + + #[test] + fn appends_rule_and_creates_directories() { + let tmp = tempdir().expect("create temp dir"); + let policy_path = tmp.path().join("rules").join("default.rules"); + + blocking_append_allow_prefix_rule( + &policy_path, + &[String::from("echo"), String::from("Hello, world!")], + ) + .expect("append rule"); + + let contents = std::fs::read_to_string(&policy_path).expect("default.rules should exist"); + assert_eq!( + contents, + r#"prefix_rule(pattern=["echo", "Hello, world!"], decision="allow") +"# + ); + } + + #[test] + fn appends_rule_without_duplicate_newline() { + let tmp = tempdir().expect("create temp dir"); + let policy_path = tmp.path().join("rules").join("default.rules"); + std::fs::create_dir_all(policy_path.parent().unwrap()).expect("create policy dir"); + std::fs::write( + &policy_path, + r#"prefix_rule(pattern=["ls"], decision="allow") +"#, + ) + .expect("write seed rule"); + + blocking_append_allow_prefix_rule( + &policy_path, + &[String::from("echo"), String::from("Hello, world!")], + ) + .expect("append rule"); + + let contents = std::fs::read_to_string(&policy_path).expect("read policy"); + assert_eq!( + contents, + r#"prefix_rule(pattern=["ls"], decision="allow") +prefix_rule(pattern=["echo", "Hello, world!"], decision="allow") +"# + ); + } + + #[test] + fn inserts_newline_when_missing_before_append() { + let tmp = tempdir().expect("create temp dir"); + let policy_path = tmp.path().join("rules").join("default.rules"); + std::fs::create_dir_all(policy_path.parent().unwrap()).expect("create policy dir"); + std::fs::write( + &policy_path, + r#"prefix_rule(pattern=["ls"], decision="allow")"#, + ) + .expect("write seed rule without newline"); + + blocking_append_allow_prefix_rule( + &policy_path, + &[String::from("echo"), String::from("Hello, world!")], + ) + .expect("append rule"); + + let contents = std::fs::read_to_string(&policy_path).expect("read policy"); + assert_eq!( + contents, + r#"prefix_rule(pattern=["ls"], decision="allow") +prefix_rule(pattern=["echo", "Hello, world!"], decision="allow") +"# + ); + } + + #[test] + fn appends_network_rule() { + let tmp = tempdir().expect("create temp dir"); + let policy_path = tmp.path().join("rules").join("default.rules"); + + blocking_append_network_rule( + &policy_path, + "Api.GitHub.com", + NetworkRuleProtocol::Https, + Decision::Allow, + Some("Allow https_connect access to api.github.com"), + ) + .expect("append network rule"); + + let contents = std::fs::read_to_string(&policy_path).expect("read policy"); + assert_eq!( + contents, + r#"network_rule(host="api.github.com", protocol="https", decision="allow", justification="Allow https_connect access to api.github.com") +"# + ); + } + + #[test] + fn appends_prefix_and_network_rules() { + let tmp = tempdir().expect("create temp dir"); + let policy_path = tmp.path().join("rules").join("default.rules"); + + blocking_append_allow_prefix_rule(&policy_path, &[String::from("curl")]) + .expect("append prefix rule"); + blocking_append_network_rule( + &policy_path, + "api.github.com", + NetworkRuleProtocol::Https, + Decision::Allow, + Some("Allow https_connect access to api.github.com"), + ) + .expect("append network rule"); + + let contents = std::fs::read_to_string(&policy_path).expect("read policy"); + assert_eq!( + contents, + r#"prefix_rule(pattern=["curl"], decision="allow") +network_rule(host="api.github.com", protocol="https", decision="allow", justification="Allow https_connect access to api.github.com") +"# + ); + } + + #[test] + fn rejects_wildcard_network_rule_host() { + let tmp = tempdir().expect("create temp dir"); + let policy_path = tmp.path().join("rules").join("default.rules"); + let err = blocking_append_network_rule( + &policy_path, + "*.example.com", + NetworkRuleProtocol::Https, + Decision::Allow, + /*justification*/ None, + ) + .expect_err("wildcards should be rejected"); + assert_eq!( + err.to_string(), + "invalid network rule: invalid rule: network_rule host must be a specific host; wildcards are not allowed" + ); + } +} diff --git a/codex-rs/execpolicy/src/decision.rs b/codex-rs/execpolicy/src/decision.rs new file mode 100644 index 0000000000000000000000000000000000000000..15dc1f711d7891742576ac000baef756a9342894 --- /dev/null +++ b/codex-rs/execpolicy/src/decision.rs @@ -0,0 +1,27 @@ +use serde::Deserialize; +use serde::Serialize; + +use crate::error::Error; +use crate::error::Result; + +#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum Decision { + /// Command may run without further approval. + Allow, + /// Request explicit user approval; rejected outright when running with `approval_policy="never"`. + Prompt, + /// Command is blocked without further consideration. + Forbidden, +} + +impl Decision { + pub fn parse(raw: &str) -> Result { + match raw { + "allow" => Ok(Self::Allow), + "prompt" => Ok(Self::Prompt), + "forbidden" => Ok(Self::Forbidden), + other => Err(Error::InvalidDecision(other.to_string())), + } + } +} diff --git a/codex-rs/execpolicy/src/error.rs b/codex-rs/execpolicy/src/error.rs new file mode 100644 index 0000000000000000000000000000000000000000..0f2be753aa65f18cec7d6a4fabeec8a0fa6d253c --- /dev/null +++ b/codex-rs/execpolicy/src/error.rs @@ -0,0 +1,101 @@ +use starlark::Error as StarlarkError; +use thiserror::Error; + +pub type Result = std::result::Result; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct TextPosition { + pub line: usize, + pub column: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct TextRange { + pub start: TextPosition, + pub end: TextPosition, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ErrorLocation { + pub path: String, + pub range: TextRange, +} + +#[derive(Debug, Error)] +pub enum Error { + #[error("invalid decision: {0}")] + InvalidDecision(String), + #[error("invalid pattern element: {0}")] + InvalidPattern(String), + #[error("invalid example: {0}")] + InvalidExample(String), + #[error("invalid rule: {0}")] + InvalidRule(String), + #[error( + "expected every example to match at least one rule. rules: {rules:?}; unmatched examples: \ + {examples:?}" + )] + ExampleDidNotMatch { + rules: Vec, + examples: Vec, + location: Option, + }, + #[error("expected example to not match rule `{rule}`: {example}")] + ExampleDidMatch { + rule: String, + example: String, + location: Option, + }, + #[error("starlark error: {0}")] + Starlark(StarlarkError), +} + +impl Error { + pub fn with_location(self, location: ErrorLocation) -> Self { + match self { + Error::ExampleDidNotMatch { + rules, + examples, + location: None, + } => Error::ExampleDidNotMatch { + rules, + examples, + location: Some(location), + }, + Error::ExampleDidMatch { + rule, + example, + location: None, + } => Error::ExampleDidMatch { + rule, + example, + location: Some(location), + }, + other => other, + } + } + + pub fn location(&self) -> Option { + match self { + Error::ExampleDidNotMatch { location, .. } + | Error::ExampleDidMatch { location, .. } => location.clone(), + Error::Starlark(err) => err.span().map(|span| { + let resolved = span.resolve_span(); + ErrorLocation { + path: span.filename().to_string(), + range: TextRange { + start: TextPosition { + line: resolved.begin.line + 1, + column: resolved.begin.column + 1, + }, + end: TextPosition { + line: resolved.end.line + 1, + column: resolved.end.column + 1, + }, + }, + } + }), + _ => None, + } + } +} diff --git a/codex-rs/execpolicy/src/execpolicycheck.rs b/codex-rs/execpolicy/src/execpolicycheck.rs new file mode 100644 index 0000000000000000000000000000000000000000..b72995d04b2ecff16f2a39af954a7facadfefe18 --- /dev/null +++ b/codex-rs/execpolicy/src/execpolicycheck.rs @@ -0,0 +1,95 @@ +use std::fs; +use std::path::PathBuf; + +use anyhow::Context; +use anyhow::Result; +use clap::Parser; +use serde::Serialize; + +use crate::Decision; +use crate::MatchOptions; +use crate::Policy; +use crate::PolicyParser; +use crate::RuleMatch; + +/// Arguments for evaluating a command against one or more execpolicy files. +#[derive(Debug, Parser, Clone)] +pub struct ExecPolicyCheckCommand { + /// Paths to execpolicy rule files to evaluate (repeatable). + #[arg(short = 'r', long = "rules", value_name = "PATH", required = true)] + pub rules: Vec, + + /// Pretty-print the JSON output. + #[arg(long)] + pub pretty: bool, + + /// Resolve absolute program paths against basename rules, gated by any + /// `host_executable()` definitions in the loaded policy files. + #[arg(long)] + pub resolve_host_executables: bool, + + /// Command tokens to check against the policy. + #[arg( + value_name = "COMMAND", + required = true, + trailing_var_arg = true, + allow_hyphen_values = true + )] + pub command: Vec, +} + +impl ExecPolicyCheckCommand { + /// Load the policies for this command, evaluate the command, and render JSON output. + pub fn run(&self) -> Result<()> { + let policy = load_policies(&self.rules)?; + let matched_rules = policy.matches_for_command_with_options( + &self.command, + /*heuristics_fallback*/ None, + &MatchOptions { + resolve_host_executables: self.resolve_host_executables, + }, + ); + + let json = format_matches_json(&matched_rules, self.pretty)?; + println!("{json}"); + + Ok(()) + } +} + +pub fn format_matches_json(matched_rules: &[RuleMatch], pretty: bool) -> Result { + let output = ExecPolicyCheckOutput { + matched_rules, + decision: matched_rules.iter().map(RuleMatch::decision).max(), + }; + + if pretty { + serde_json::to_string_pretty(&output).map_err(Into::into) + } else { + serde_json::to_string(&output).map_err(Into::into) + } +} + +pub fn load_policies(policy_paths: &[PathBuf]) -> Result { + let mut parser = PolicyParser::new(); + + for policy_path in policy_paths { + let policy_file_contents = fs::read_to_string(policy_path) + .with_context(|| format!("failed to read policy at {}", policy_path.display()))?; + let policy_identifier = policy_path.to_string_lossy().to_string(); + parser + .parse(&policy_identifier, &policy_file_contents) + .with_context(|| format!("failed to parse policy at {}", policy_path.display()))?; + } + + Ok(parser.build()) +} + +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +struct ExecPolicyCheckOutput<'a> { + #[serde(rename = "matchedRules")] + matched_rules: &'a [RuleMatch], + #[serde(skip_serializing_if = "Option::is_none")] + decision: Option, +} diff --git a/codex-rs/execpolicy/src/executable_name.rs b/codex-rs/execpolicy/src/executable_name.rs new file mode 100644 index 0000000000000000000000000000000000000000..bb3e8a536c2f246d5dc075817cb8fa737a3cb5e3 --- /dev/null +++ b/codex-rs/execpolicy/src/executable_name.rs @@ -0,0 +1,29 @@ +use std::path::Path; + +#[cfg(windows)] +const WINDOWS_EXECUTABLE_SUFFIXES: [&str; 4] = [".exe", ".cmd", ".bat", ".com"]; + +pub(crate) fn executable_lookup_key(raw: &str) -> String { + #[cfg(windows)] + { + let raw = raw.to_ascii_lowercase(); + for suffix in WINDOWS_EXECUTABLE_SUFFIXES { + if raw.ends_with(suffix) { + let stripped_len = raw.len() - suffix.len(); + return raw[..stripped_len].to_string(); + } + } + raw + } + + #[cfg(not(windows))] + { + raw.to_string() + } +} + +pub(crate) fn executable_path_lookup_key(path: &Path) -> Option { + path.file_name() + .and_then(|name| name.to_str()) + .map(executable_lookup_key) +} diff --git a/codex-rs/execpolicy/src/lib.rs b/codex-rs/execpolicy/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..293e2750589dd6046a49ac7561c20c6021799bfd --- /dev/null +++ b/codex-rs/execpolicy/src/lib.rs @@ -0,0 +1,33 @@ +pub(crate) mod amend; +pub(crate) mod decision; +pub(crate) mod error; +pub(crate) mod execpolicycheck; +mod executable_name; +pub(crate) mod parser; +pub(crate) mod policy; +pub mod rule; +mod sandbox_migration; + +pub use amend::AmendError; +pub use amend::blocking_append_allow_prefix_rule; +pub use amend::blocking_append_network_rule; +pub use decision::Decision; +pub use error::Error; +pub use error::ErrorLocation; +pub use error::Result; +pub use error::TextPosition; +pub use error::TextRange; +pub use execpolicycheck::ExecPolicyCheckCommand; +pub use parser::PolicyParser; +pub use policy::Evaluation; +pub use policy::MatchOptions; +pub use policy::Policy; +pub use policy::RequirementsExecPolicy; +pub use rule::NetworkRuleProtocol; +pub use rule::PatternToken; +pub use rule::PrefixPattern; +pub use rule::PrefixRule; +pub use rule::Rule; +pub use rule::RuleMatch; +pub use rule::RuleRef; +pub use sandbox_migration::prefix_rule_migration; diff --git a/codex-rs/execpolicy/src/main.rs b/codex-rs/execpolicy/src/main.rs new file mode 100644 index 0000000000000000000000000000000000000000..82e1bc12c4435ff6dc17526ed6426ce73f9f574a --- /dev/null +++ b/codex-rs/execpolicy/src/main.rs @@ -0,0 +1,18 @@ +use anyhow::Result; +use clap::Parser; +use codex_execpolicy::ExecPolicyCheckCommand; + +/// CLI for evaluating exec policies +#[derive(Parser)] +#[command(name = "codex-execpolicy")] +enum Cli { + /// Evaluate a command against a policy. + Check(ExecPolicyCheckCommand), +} + +fn main() -> Result<()> { + let cli = Cli::parse(); + match cli { + Cli::Check(cmd) => cmd.run(), + } +} diff --git a/codex-rs/execpolicy/src/parser.rs b/codex-rs/execpolicy/src/parser.rs new file mode 100644 index 0000000000000000000000000000000000000000..33262db1b3a8bdff7a0c5ec6d93b58b437eb8166 --- /dev/null +++ b/codex-rs/execpolicy/src/parser.rs @@ -0,0 +1,473 @@ +use codex_utils_absolute_path::AbsolutePathBuf; +use multimap::MultiMap; +use starlark::any::ProvidesStaticType; +use starlark::codemap::FileSpan; +use starlark::environment::GlobalsBuilder; +use starlark::environment::Module; +use starlark::eval::Evaluator; +use starlark::starlark_module; +use starlark::syntax::AstModule; +use starlark::syntax::Dialect; +use starlark::values::Value; +use starlark::values::list::ListRef; +use starlark::values::list::UnpackList; +use starlark::values::none::NoneType; +use std::cell::RefCell; +use std::cell::RefMut; +use std::collections::HashMap; +use std::path::Path; +use std::sync::Arc; + +use crate::decision::Decision; +use crate::error::Error; +use crate::error::ErrorLocation; +use crate::error::Result; +use crate::error::TextPosition; +use crate::error::TextRange; +use crate::executable_name::executable_lookup_key; +use crate::executable_name::executable_path_lookup_key; +use crate::rule::NetworkRule; +use crate::rule::NetworkRuleProtocol; +use crate::rule::PatternToken; +use crate::rule::PrefixPattern; +use crate::rule::PrefixRule; +use crate::rule::RuleRef; +use crate::rule::validate_match_examples; +use crate::rule::validate_not_match_examples; + +pub struct PolicyParser { + builder: RefCell, +} + +impl Default for PolicyParser { + fn default() -> Self { + Self::new() + } +} + +impl PolicyParser { + pub fn new() -> Self { + Self { + builder: RefCell::new(PolicyBuilder::new()), + } + } + + /// Parses a policy, tagging parser errors with `policy_identifier` so failures include the + /// identifier alongside line numbers. + pub fn parse(&mut self, policy_identifier: &str, policy_file_contents: &str) -> Result<()> { + let pending_validation_count = self.builder.borrow().pending_example_validations.len(); + let mut dialect = Dialect::Extended.clone(); + dialect.enable_f_strings = true; + let ast = AstModule::parse( + policy_identifier, + policy_file_contents.to_string(), + &dialect, + ) + .map_err(Error::Starlark)?; + let globals = GlobalsBuilder::standard().with(policy_builtins).build(); + Module::with_temp_heap(|module| { + let mut eval = Evaluator::new(&module); + eval.extra = Some(&self.builder); + eval.eval_module(ast, &globals) + .map(|_| ()) + .map_err(Error::Starlark) + })?; + self.builder + .borrow() + .validate_pending_examples_from(pending_validation_count)?; + Ok(()) + } + + pub fn build(self) -> crate::policy::Policy { + self.builder.into_inner().build() + } +} + +#[derive(Debug, ProvidesStaticType)] +struct PolicyBuilder { + rules_by_program: MultiMap, + network_rules: Vec, + host_executables_by_name: HashMap>, + pending_example_validations: Vec, +} + +impl PolicyBuilder { + fn new() -> Self { + Self { + rules_by_program: MultiMap::new(), + network_rules: Vec::new(), + host_executables_by_name: HashMap::new(), + pending_example_validations: Vec::new(), + } + } + + fn add_rule(&mut self, rule: RuleRef) { + self.rules_by_program + .insert(rule.program().to_string(), rule); + } + + fn add_network_rule(&mut self, rule: NetworkRule) { + self.network_rules.push(rule); + } + + fn add_host_executable(&mut self, name: String, paths: Vec) { + self.host_executables_by_name.insert(name, paths.into()); + } + + fn add_pending_example_validation( + &mut self, + rules: Vec, + matches: Vec>, + not_matches: Vec>, + location: Option, + ) { + self.pending_example_validations + .push(PendingExampleValidation { + rules, + matches, + not_matches, + location, + }); + } + + fn validate_pending_examples_from(&self, start: usize) -> Result<()> { + for validation in &self.pending_example_validations[start..] { + let mut rules_by_program = MultiMap::new(); + for rule in &validation.rules { + rules_by_program.insert(rule.program().to_string(), rule.clone()); + } + + let policy = crate::policy::Policy::from_parts( + rules_by_program, + Vec::new(), + self.host_executables_by_name.clone(), + ); + validate_not_match_examples(&policy, &validation.rules, &validation.not_matches) + .map_err(|error| attach_validation_location(error, validation.location.clone()))?; + validate_match_examples(&policy, &validation.rules, &validation.matches) + .map_err(|error| attach_validation_location(error, validation.location.clone()))?; + } + + Ok(()) + } + + fn build(self) -> crate::policy::Policy { + crate::policy::Policy::from_parts( + self.rules_by_program, + self.network_rules, + self.host_executables_by_name, + ) + } +} + +#[derive(Debug)] +struct PendingExampleValidation { + rules: Vec, + matches: Vec>, + not_matches: Vec>, + location: Option, +} + +fn parse_pattern<'v>(pattern: UnpackList>) -> Result> { + let tokens: Vec = pattern + .items + .into_iter() + .map(parse_pattern_token) + .collect::>()?; + if tokens.is_empty() { + Err(Error::InvalidPattern("pattern cannot be empty".to_string())) + } else { + Ok(tokens) + } +} + +fn parse_pattern_token<'v>(value: Value<'v>) -> Result { + if let Some(s) = value.unpack_str() { + Ok(PatternToken::Single(s.to_string())) + } else if let Some(list) = ListRef::from_value(value) { + let tokens: Vec = list + .content() + .iter() + .map(|value| { + value + .unpack_str() + .ok_or_else(|| { + Error::InvalidPattern(format!( + "pattern alternative must be a string (got {})", + value.get_type() + )) + }) + .map(str::to_string) + }) + .collect::>()?; + + match tokens.as_slice() { + [] => Err(Error::InvalidPattern( + "pattern alternatives cannot be empty".to_string(), + )), + [single] => Ok(PatternToken::Single(single.clone())), + _ => Ok(PatternToken::Alts(tokens)), + } + } else { + Err(Error::InvalidPattern(format!( + "pattern element must be a string or list of strings (got {})", + value.get_type() + ))) + } +} + +fn parse_examples<'v>(examples: UnpackList>) -> Result>> { + examples.items.into_iter().map(parse_example).collect() +} + +fn parse_literal_absolute_path(raw: &str) -> Result { + if !Path::new(raw).is_absolute() { + return Err(Error::InvalidRule(format!( + "host_executable paths must be absolute (got {raw})" + ))); + } + + AbsolutePathBuf::try_from(raw.to_string()) + .map_err(|error| Error::InvalidRule(format!("invalid absolute path `{raw}`: {error}"))) +} + +fn validate_host_executable_name(name: &str) -> Result<()> { + if name.is_empty() { + return Err(Error::InvalidRule( + "host_executable name cannot be empty".to_string(), + )); + } + + let path = Path::new(name); + if path.components().count() != 1 + || path.file_name().and_then(|value| value.to_str()) != Some(name) + { + return Err(Error::InvalidRule(format!( + "host_executable name must be a bare executable name (got {name})" + ))); + } + + Ok(()) +} + +fn parse_network_rule_decision(raw: &str) -> Result { + match raw { + "deny" => Ok(Decision::Forbidden), + other => Decision::parse(other), + } +} + +fn error_location_from_file_span(span: FileSpan) -> ErrorLocation { + let resolved = span.resolve_span(); + ErrorLocation { + path: span.filename().to_string(), + range: TextRange { + start: TextPosition { + line: resolved.begin.line + 1, + column: resolved.begin.column + 1, + }, + end: TextPosition { + line: resolved.end.line + 1, + column: resolved.end.column + 1, + }, + }, + } +} + +fn attach_validation_location(error: Error, location: Option) -> Error { + match location { + Some(location) => error.with_location(location), + None => error, + } +} + +fn parse_example<'v>(value: Value<'v>) -> Result> { + if let Some(raw) = value.unpack_str() { + parse_string_example(raw) + } else if let Some(list) = ListRef::from_value(value) { + parse_list_example(list) + } else { + Err(Error::InvalidExample(format!( + "example must be a string or list of strings (got {})", + value.get_type() + ))) + } +} + +fn parse_string_example(raw: &str) -> Result> { + let tokens = shlex::split(raw).ok_or_else(|| { + Error::InvalidExample("example string has invalid shell syntax".to_string()) + })?; + + if tokens.is_empty() { + Err(Error::InvalidExample( + "example cannot be an empty string".to_string(), + )) + } else { + Ok(tokens) + } +} + +fn parse_list_example(list: &ListRef) -> Result> { + let tokens: Vec = list + .content() + .iter() + .map(|value| { + value + .unpack_str() + .ok_or_else(|| { + Error::InvalidExample(format!( + "example tokens must be strings (got {})", + value.get_type() + )) + }) + .map(str::to_string) + }) + .collect::>()?; + + if tokens.is_empty() { + Err(Error::InvalidExample( + "example cannot be an empty list".to_string(), + )) + } else { + Ok(tokens) + } +} + +fn policy_builder<'v, 'a>(eval: &Evaluator<'v, 'a, '_>) -> RefMut<'a, PolicyBuilder> { + #[expect(clippy::expect_used)] + eval.extra + .as_ref() + .expect("policy_builder requires Evaluator.extra to be populated") + .downcast_ref::>() + .expect("Evaluator.extra must contain a PolicyBuilder") + .borrow_mut() +} + +#[starlark_module] +fn policy_builtins(builder: &mut GlobalsBuilder) { + fn prefix_rule<'v>( + pattern: UnpackList>, + decision: Option<&'v str>, + r#match: Option>>, + not_match: Option>>, + justification: Option<&'v str>, + eval: &mut Evaluator<'v, '_, '_>, + ) -> anyhow::Result { + let decision = match decision { + Some(raw) => Decision::parse(raw)?, + None => Decision::Allow, + }; + + let justification = match justification { + Some(raw) if raw.trim().is_empty() => { + return Err(Error::InvalidRule("justification cannot be empty".to_string()).into()); + } + Some(raw) => Some(raw.to_string()), + None => None, + }; + + let pattern_tokens = parse_pattern(pattern)?; + + let matches: Vec> = + r#match.map(parse_examples).transpose()?.unwrap_or_default(); + let not_matches: Vec> = not_match + .map(parse_examples) + .transpose()? + .unwrap_or_default(); + let location = eval + .call_stack_top_location() + .map(error_location_from_file_span); + + let mut builder = policy_builder(eval); + + let (first_token, remaining_tokens) = pattern_tokens + .split_first() + .ok_or_else(|| Error::InvalidPattern("pattern cannot be empty".to_string()))?; + + let rest: Arc<[PatternToken]> = remaining_tokens.to_vec().into(); + + let rules: Vec = first_token + .alternatives() + .iter() + .map(|head| { + Arc::new(PrefixRule { + pattern: PrefixPattern { + first: Arc::from(head.as_str()), + rest: rest.clone(), + }, + decision, + justification: justification.clone(), + }) as RuleRef + }) + .collect(); + + builder.add_pending_example_validation(rules.clone(), matches, not_matches, location); + rules.into_iter().for_each(|rule| builder.add_rule(rule)); + Ok(NoneType) + } + + fn network_rule<'v>( + host: &'v str, + protocol: &'v str, + decision: &'v str, + justification: Option<&'v str>, + eval: &mut Evaluator<'v, '_, '_>, + ) -> anyhow::Result { + let protocol = NetworkRuleProtocol::parse(protocol)?; + let decision = parse_network_rule_decision(decision)?; + let justification = match justification { + Some(raw) if raw.trim().is_empty() => { + return Err(Error::InvalidRule("justification cannot be empty".to_string()).into()); + } + Some(raw) => Some(raw.to_string()), + None => None, + }; + + let mut builder = policy_builder(eval); + builder.add_network_rule(NetworkRule { + host: crate::rule::normalize_network_rule_host(host)?, + protocol, + decision, + justification, + }); + Ok(NoneType) + } + + fn host_executable<'v>( + name: &'v str, + paths: UnpackList>, + eval: &mut Evaluator<'v, '_, '_>, + ) -> anyhow::Result { + validate_host_executable_name(name)?; + + let mut parsed_paths = Vec::new(); + for value in paths.items { + let raw = value.unpack_str().ok_or_else(|| { + Error::InvalidRule(format!( + "host_executable paths must be strings (got {})", + value.get_type() + )) + })?; + let path = parse_literal_absolute_path(raw)?; + let Some(path_name) = executable_path_lookup_key(path.as_path()) else { + return Err(Error::InvalidRule(format!( + "host_executable path `{raw}` must have basename `{name}`" + )) + .into()); + }; + if path_name != executable_lookup_key(name) { + return Err(Error::InvalidRule(format!( + "host_executable path `{raw}` must have basename `{name}`" + )) + .into()); + } + if !parsed_paths.iter().any(|existing| existing == &path) { + parsed_paths.push(path); + } + } + + policy_builder(eval).add_host_executable(executable_lookup_key(name), parsed_paths); + Ok(NoneType) + } +} diff --git a/codex-rs/execpolicy/src/policy.rs b/codex-rs/execpolicy/src/policy.rs new file mode 100644 index 0000000000000000000000000000000000000000..78790517fd2b4bd462b02a304f24e20e6696e796 --- /dev/null +++ b/codex-rs/execpolicy/src/policy.rs @@ -0,0 +1,412 @@ +use crate::decision::Decision; +use crate::error::Error; +use crate::error::Result; +use crate::executable_name::executable_path_lookup_key; +use crate::rule::NetworkRule; +use crate::rule::NetworkRuleProtocol; +use crate::rule::PatternToken; +use crate::rule::PrefixPattern; +use crate::rule::PrefixRule; +use crate::rule::RuleMatch; +use crate::rule::RuleRef; +use crate::rule::normalize_network_rule_host; +use codex_utils_absolute_path::AbsolutePathBuf; +use multimap::MultiMap; +use serde::Deserialize; +use serde::Serialize; +use std::collections::HashMap; +use std::sync::Arc; + +type HeuristicsFallback<'a> = Option<&'a dyn Fn(&[String]) -> Decision>; + +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct MatchOptions { + pub resolve_host_executables: bool, +} + +#[derive(Clone, Debug)] +pub struct Policy { + rules_by_program: MultiMap, + network_rules: Vec, + host_executables_by_name: HashMap>, +} + +/// A requirements-owned execution policy with order-independent equality. +#[derive(Debug, Clone)] +pub struct RequirementsExecPolicy { + policy: Policy, +} + +impl RequirementsExecPolicy { + pub fn new(policy: Policy) -> Self { + Self { policy } + } + + pub fn fingerprint(&self) -> Vec { + let mut entries = Vec::new(); + for (program, rules) in self.policy.rules().iter_all() { + for rule in rules { + entries.push(format!("{program}:{rule:?}")); + } + } + entries.sort(); + entries + } +} + +impl PartialEq for RequirementsExecPolicy { + fn eq(&self, other: &Self) -> bool { + self.fingerprint() == other.fingerprint() + } +} + +impl Eq for RequirementsExecPolicy {} + +impl AsRef for RequirementsExecPolicy { + fn as_ref(&self) -> &Policy { + &self.policy + } +} + +impl Policy { + pub fn new(rules_by_program: MultiMap) -> Self { + Self::from_parts(rules_by_program, Vec::new(), HashMap::new()) + } + + pub fn from_parts( + rules_by_program: MultiMap, + network_rules: Vec, + host_executables_by_name: HashMap>, + ) -> Self { + Self { + rules_by_program, + network_rules, + host_executables_by_name, + } + } + + pub fn empty() -> Self { + Self::new(MultiMap::new()) + } + + pub fn rules(&self) -> &MultiMap { + &self.rules_by_program + } + + pub fn network_rules(&self) -> &[NetworkRule] { + &self.network_rules + } + + pub fn host_executables(&self) -> &HashMap> { + &self.host_executables_by_name + } + + pub fn get_allowed_prefixes(&self) -> Vec> { + let mut prefixes = Vec::new(); + + for (_program, rules) in self.rules_by_program.iter_all() { + for rule in rules { + let Some(prefix_rule) = rule.as_any().downcast_ref::() else { + continue; + }; + if prefix_rule.decision != Decision::Allow { + continue; + } + + let mut prefix = Vec::with_capacity(prefix_rule.pattern.rest.len() + 1); + prefix.push(prefix_rule.pattern.first.as_ref().to_string()); + prefix.extend(prefix_rule.pattern.rest.iter().map(render_pattern_token)); + prefixes.push(prefix); + } + } + + prefixes.sort(); + prefixes.dedup(); + prefixes + } + + pub fn add_prefix_rule(&mut self, prefix: &[String], decision: Decision) -> Result<()> { + let (first_token, rest) = prefix + .split_first() + .ok_or_else(|| Error::InvalidPattern("prefix cannot be empty".to_string()))?; + + let rule: RuleRef = Arc::new(PrefixRule { + pattern: PrefixPattern { + first: Arc::from(first_token.as_str()), + rest: rest + .iter() + .map(|token| PatternToken::Single(token.clone())) + .collect::>() + .into(), + }, + decision, + justification: None, + }); + + self.rules_by_program.insert(first_token.clone(), rule); + Ok(()) + } + + pub fn add_network_rule( + &mut self, + host: &str, + protocol: NetworkRuleProtocol, + decision: Decision, + justification: Option, + ) -> Result<()> { + let host = normalize_network_rule_host(host)?; + if let Some(raw) = justification.as_deref() + && raw.trim().is_empty() + { + return Err(Error::InvalidRule( + "justification cannot be empty".to_string(), + )); + } + self.network_rules.push(NetworkRule { + host, + protocol, + decision, + justification, + }); + Ok(()) + } + + pub fn set_host_executable_paths(&mut self, name: String, paths: Vec) { + self.host_executables_by_name.insert(name, paths.into()); + } + + pub fn merge_overlay(&self, overlay: &Policy) -> Policy { + let mut combined_rules = self.rules_by_program.clone(); + for (program, rules) in overlay.rules_by_program.iter_all() { + for rule in rules { + combined_rules.insert(program.clone(), rule.clone()); + } + } + + let mut combined_network_rules = self.network_rules.clone(); + combined_network_rules.extend(overlay.network_rules.iter().cloned()); + + let mut host_executables_by_name = self.host_executables_by_name.clone(); + host_executables_by_name.extend( + overlay + .host_executables_by_name + .iter() + .map(|(name, paths)| (name.clone(), paths.clone())), + ); + + Policy::from_parts( + combined_rules, + combined_network_rules, + host_executables_by_name, + ) + } + + pub fn compiled_network_domains(&self) -> (Vec, Vec) { + let mut allowed = Vec::new(); + let mut denied = Vec::new(); + + for rule in &self.network_rules { + match rule.decision { + Decision::Allow => { + denied.retain(|entry| entry != &rule.host); + upsert_domain(&mut allowed, &rule.host); + } + Decision::Forbidden => { + allowed.retain(|entry| entry != &rule.host); + upsert_domain(&mut denied, &rule.host); + } + Decision::Prompt => {} + } + } + + (allowed, denied) + } + + pub fn check(&self, cmd: &[String], heuristics_fallback: &F) -> Evaluation + where + F: Fn(&[String]) -> Decision, + { + let matched_rules = self.matches_for_command_with_options( + cmd, + Some(heuristics_fallback), + &MatchOptions::default(), + ); + Evaluation::from_matches(matched_rules) + } + + pub fn check_with_options( + &self, + cmd: &[String], + heuristics_fallback: &F, + options: &MatchOptions, + ) -> Evaluation + where + F: Fn(&[String]) -> Decision, + { + let matched_rules = + self.matches_for_command_with_options(cmd, Some(heuristics_fallback), options); + Evaluation::from_matches(matched_rules) + } + + /// Checks multiple commands and aggregates the results. + pub fn check_multiple( + &self, + commands: Commands, + heuristics_fallback: &F, + ) -> Evaluation + where + Commands: IntoIterator, + Commands::Item: AsRef<[String]>, + F: Fn(&[String]) -> Decision, + { + self.check_multiple_with_options(commands, heuristics_fallback, &MatchOptions::default()) + } + + pub fn check_multiple_with_options( + &self, + commands: Commands, + heuristics_fallback: &F, + options: &MatchOptions, + ) -> Evaluation + where + Commands: IntoIterator, + Commands::Item: AsRef<[String]>, + F: Fn(&[String]) -> Decision, + { + let matched_rules: Vec = commands + .into_iter() + .flat_map(|command| { + self.matches_for_command_with_options( + command.as_ref(), + Some(heuristics_fallback), + options, + ) + }) + .collect(); + + Evaluation::from_matches(matched_rules) + } + + /// Returns matching rules for the given command. If no rules match and + /// `heuristics_fallback` is provided, returns a single + /// `HeuristicsRuleMatch` with the decision rendered by + /// `heuristics_fallback`. + /// + /// If `heuristics_fallback.is_some()`, then the returned vector is + /// guaranteed to be non-empty. + pub fn matches_for_command( + &self, + cmd: &[String], + heuristics_fallback: HeuristicsFallback<'_>, + ) -> Vec { + self.matches_for_command_with_options(cmd, heuristics_fallback, &MatchOptions::default()) + } + + pub fn matches_for_command_with_options( + &self, + cmd: &[String], + heuristics_fallback: HeuristicsFallback<'_>, + options: &MatchOptions, + ) -> Vec { + let matched_rules = self + .match_exact_rules(cmd) + .filter(|matched_rules| !matched_rules.is_empty()) + .or_else(|| { + options + .resolve_host_executables + .then(|| self.match_host_executable_rules(cmd)) + .filter(|matched_rules| !matched_rules.is_empty()) + }) + .unwrap_or_default(); + + if matched_rules.is_empty() + && let Some(heuristics_fallback) = heuristics_fallback + { + vec![RuleMatch::HeuristicsRuleMatch { + command: cmd.to_vec(), + decision: heuristics_fallback(cmd), + }] + } else { + matched_rules + } + } + + fn match_exact_rules(&self, cmd: &[String]) -> Option> { + let first = cmd.first()?; + Some( + self.rules_by_program + .get_vec(first) + .map(|rules| rules.iter().filter_map(|rule| rule.matches(cmd)).collect()) + .unwrap_or_default(), + ) + } + + fn match_host_executable_rules(&self, cmd: &[String]) -> Vec { + let Some(first) = cmd.first() else { + return Vec::new(); + }; + let Ok(program) = AbsolutePathBuf::try_from(first.clone()) else { + return Vec::new(); + }; + let Some(basename) = executable_path_lookup_key(program.as_path()) else { + return Vec::new(); + }; + let Some(rules) = self.rules_by_program.get_vec(&basename) else { + return Vec::new(); + }; + if let Some(paths) = self.host_executables_by_name.get(&basename) + && !paths.iter().any(|path| path == &program) + { + return Vec::new(); + } + + let basename_command = std::iter::once(basename) + .chain(cmd.iter().skip(1).cloned()) + .collect::>(); + rules + .iter() + .filter_map(|rule| rule.matches(&basename_command)) + .map(|rule_match| rule_match.with_resolved_program(&program)) + .collect() + } +} + +fn upsert_domain(entries: &mut Vec, host: &str) { + entries.retain(|entry| entry != host); + entries.push(host.to_string()); +} + +fn render_pattern_token(token: &PatternToken) -> String { + match token { + PatternToken::Single(value) => value.clone(), + PatternToken::Alts(alternatives) => format!("[{}]", alternatives.join("|")), + } +} + +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct Evaluation { + pub decision: Decision, + #[serde(rename = "matchedRules")] + pub matched_rules: Vec, +} + +impl Evaluation { + pub fn is_match(&self) -> bool { + self.matched_rules + .iter() + .any(|rule_match| !matches!(rule_match, RuleMatch::HeuristicsRuleMatch { .. })) + } + + /// Caller is responsible for ensuring that `matched_rules` is non-empty. + fn from_matches(matched_rules: Vec) -> Self { + let decision = matched_rules.iter().map(RuleMatch::decision).max(); + #[expect(clippy::expect_used)] + let decision = decision.expect("invariant failed: matched_rules must be non-empty"); + + Self { + decision, + matched_rules, + } + } +} diff --git a/codex-rs/execpolicy/src/rule.rs b/codex-rs/execpolicy/src/rule.rs new file mode 100644 index 0000000000000000000000000000000000000000..6642cb1d9c1b67c0fd00a991660a019d50f636fd --- /dev/null +++ b/codex-rs/execpolicy/src/rule.rs @@ -0,0 +1,306 @@ +use crate::decision::Decision; +use crate::error::Error; +use crate::error::Result; +use crate::policy::MatchOptions; +use crate::policy::Policy; +use codex_utils_absolute_path::AbsolutePathBuf; +use serde::Deserialize; +use serde::Serialize; +use shlex::try_join; +use std::any::Any; +use std::fmt::Debug; +use std::sync::Arc; + +/// Matches a single command token, either a fixed string or one of several allowed alternatives. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum PatternToken { + Single(String), + Alts(Vec), +} + +impl PatternToken { + fn matches(&self, token: &str) -> bool { + match self { + Self::Single(expected) => expected == token, + Self::Alts(alternatives) => alternatives.iter().any(|alt| alt == token), + } + } + + pub fn alternatives(&self) -> &[String] { + match self { + Self::Single(expected) => std::slice::from_ref(expected), + Self::Alts(alternatives) => alternatives, + } + } +} + +/// Prefix matcher for commands with support for alternative match tokens. +/// First token is fixed since we key by the first token in policy. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct PrefixPattern { + pub first: Arc, + pub rest: Arc<[PatternToken]>, +} + +impl PrefixPattern { + pub fn matches_prefix(&self, cmd: &[String]) -> Option> { + let pattern_length = self.rest.len() + 1; + if cmd.len() < pattern_length || cmd[0] != self.first.as_ref() { + return None; + } + + for (pattern_token, cmd_token) in self.rest.iter().zip(&cmd[1..pattern_length]) { + if !pattern_token.matches(cmd_token) { + return None; + } + } + + Some(cmd[..pattern_length].to_vec()) + } +} + +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum RuleMatch { + PrefixRuleMatch { + #[serde(rename = "matchedPrefix")] + matched_prefix: Vec, + decision: Decision, + #[serde(rename = "resolvedProgram", skip_serializing_if = "Option::is_none")] + resolved_program: Option, + /// Optional rationale for why this rule exists. + /// + /// This can be supplied for any decision and may be surfaced in different contexts + /// (e.g., prompt reasons or rejection messages). + #[serde(skip_serializing_if = "Option::is_none")] + justification: Option, + }, + HeuristicsRuleMatch { + command: Vec, + decision: Decision, + }, +} + +impl RuleMatch { + pub fn decision(&self) -> Decision { + match self { + Self::PrefixRuleMatch { decision, .. } => *decision, + Self::HeuristicsRuleMatch { decision, .. } => *decision, + } + } + + pub fn with_resolved_program(self, resolved_program: &AbsolutePathBuf) -> Self { + match self { + Self::PrefixRuleMatch { + matched_prefix, + decision, + justification, + .. + } => Self::PrefixRuleMatch { + matched_prefix, + decision, + resolved_program: Some(resolved_program.clone()), + justification, + }, + other => other, + } + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct PrefixRule { + pub pattern: PrefixPattern, + pub decision: Decision, + pub justification: Option, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum NetworkRuleProtocol { + Http, + Https, + Socks5Tcp, + Socks5Udp, +} + +impl NetworkRuleProtocol { + pub fn parse(raw: &str) -> Result { + match raw { + "http" => Ok(Self::Http), + "https" | "https_connect" | "http-connect" => Ok(Self::Https), + "socks5_tcp" => Ok(Self::Socks5Tcp), + "socks5_udp" => Ok(Self::Socks5Udp), + other => Err(Error::InvalidRule(format!( + "network_rule protocol must be one of http, https, socks5_tcp, socks5_udp (got {other})" + ))), + } + } + + pub fn as_policy_string(self) -> &'static str { + match self { + Self::Http => "http", + Self::Https => "https", + Self::Socks5Tcp => "socks5_tcp", + Self::Socks5Udp => "socks5_udp", + } + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct NetworkRule { + pub host: String, + pub protocol: NetworkRuleProtocol, + pub decision: Decision, + pub justification: Option, +} + +pub(crate) fn normalize_network_rule_host(raw: &str) -> Result { + let mut host = raw.trim(); + if host.is_empty() { + return Err(Error::InvalidRule( + "network_rule host cannot be empty".to_string(), + )); + } + if host.contains("://") || host.contains('/') || host.contains('?') || host.contains('#') { + return Err(Error::InvalidRule( + "network_rule host must be a hostname or IP literal (without scheme or path)" + .to_string(), + )); + } + + if let Some(stripped) = host.strip_prefix('[') { + let Some((inside, rest)) = stripped.split_once(']') else { + return Err(Error::InvalidRule( + "network_rule host has an invalid bracketed IPv6 literal".to_string(), + )); + }; + let port_ok = rest + .strip_prefix(':') + .is_some_and(|port| !port.is_empty() && port.chars().all(|c| c.is_ascii_digit())); + if !rest.is_empty() && !port_ok { + return Err(Error::InvalidRule(format!( + "network_rule host contains an unsupported suffix: {raw}" + ))); + } + host = inside; + } else if host.matches(':').count() == 1 + && let Some((candidate, port)) = host.rsplit_once(':') + && !candidate.is_empty() + && !port.is_empty() + && port.chars().all(|c| c.is_ascii_digit()) + { + host = candidate; + } + + let normalized = host.trim_end_matches('.').trim().to_ascii_lowercase(); + if normalized.is_empty() { + return Err(Error::InvalidRule( + "network_rule host cannot be empty".to_string(), + )); + } + if normalized.contains('*') { + return Err(Error::InvalidRule( + "network_rule host must be a specific host; wildcards are not allowed".to_string(), + )); + } + if normalized.chars().any(char::is_whitespace) { + return Err(Error::InvalidRule( + "network_rule host cannot contain whitespace".to_string(), + )); + } + + Ok(normalized) +} + +pub trait Rule: Any + Debug + Send + Sync { + fn program(&self) -> &str; + + fn matches(&self, cmd: &[String]) -> Option; + + fn as_any(&self) -> &dyn Any; +} + +pub type RuleRef = Arc; + +impl Rule for PrefixRule { + fn program(&self) -> &str { + self.pattern.first.as_ref() + } + + fn matches(&self, cmd: &[String]) -> Option { + self.pattern + .matches_prefix(cmd) + .map(|matched_prefix| RuleMatch::PrefixRuleMatch { + matched_prefix, + decision: self.decision, + resolved_program: None, + justification: self.justification.clone(), + }) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +/// Count how many rules match each provided example and error if any example is unmatched. +pub(crate) fn validate_match_examples( + policy: &Policy, + rules: &[RuleRef], + matches: &[Vec], +) -> Result<()> { + let mut unmatched_examples = Vec::new(); + let options = MatchOptions { + resolve_host_executables: true, + }; + + for example in matches { + if !policy + .matches_for_command_with_options(example, /*heuristics_fallback*/ None, &options) + .is_empty() + { + continue; + } + + unmatched_examples.push( + try_join(example.iter().map(String::as_str)) + .unwrap_or_else(|_| "unable to render example".to_string()), + ); + } + + if unmatched_examples.is_empty() { + Ok(()) + } else { + Err(Error::ExampleDidNotMatch { + rules: rules.iter().map(|rule| format!("{rule:?}")).collect(), + examples: unmatched_examples, + location: None, + }) + } +} + +/// Ensure that no rule matches any provided negative example. +pub(crate) fn validate_not_match_examples( + policy: &Policy, + _rules: &[RuleRef], + not_matches: &[Vec], +) -> Result<()> { + let options = MatchOptions { + resolve_host_executables: true, + }; + + for example in not_matches { + if let Some(rule) = policy + .matches_for_command_with_options(example, /*heuristics_fallback*/ None, &options) + .first() + { + return Err(Error::ExampleDidMatch { + rule: format!("{rule:?}"), + example: try_join(example.iter().map(String::as_str)) + .unwrap_or_else(|_| "unable to render example".to_string()), + location: None, + }); + } + } + + Ok(()) +} diff --git a/codex-rs/execpolicy/src/sandbox_migration.rs b/codex-rs/execpolicy/src/sandbox_migration.rs new file mode 100644 index 0000000000000000000000000000000000000000..a2a57c1c4cf4703399a5eff5cf3693f4a52d25da --- /dev/null +++ b/codex-rs/execpolicy/src/sandbox_migration.rs @@ -0,0 +1,123 @@ +use std::collections::HashSet; +use std::io; +use std::io::SeekFrom; +use std::io::Write as _; +use std::path::Path; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncSeekExt; +use tokio::io::AsyncWriteExt; + +const MIGRATION_MARKER_FILENAME: &str = ".sandbox_migration"; + +/// removes legacy allow rules that newer codex versions no longer offer. +/// +/// this migration is intentionally one-shot. once complete, a marker in `codex_home` prevents +/// policies saved by newer codex versions from being removed on later startups. +pub async fn prefix_rule_migration( + codex_home: &Path, + policy_path: &Path, + banned_prefixes: &[&[&str]], +) -> io::Result<()> { + let marker_path = codex_home.join(MIGRATION_MARKER_FILENAME); + if tokio::fs::try_exists(&marker_path).await? { + return Ok(()); + } + clean_rules_file(policy_path, banned_prefixes).await?; + + write_migration_marker(codex_home, &marker_path).await?; + Ok(()) +} + +// atomically writes the marker after creating codex home when needed. +async fn write_migration_marker(codex_home: &Path, marker_path: &Path) -> io::Result<()> { + tokio::fs::create_dir_all(codex_home).await?; + let codex_home = codex_home.to_owned(); + let marker_path = marker_path.to_owned(); + tokio::task::spawn_blocking(move || { + let mut marker = tempfile::NamedTempFile::new_in(codex_home)?; + marker.write_all(b"v1\n")?; + match marker.persist_noclobber(marker_path) { + Ok(_) => Ok(()), + Err(err) if err.error.kind() == io::ErrorKind::AlreadyExists => Ok(()), + Err(err) => Err(err.error), + } + }) + .await + .map_err(io::Error::other)? +} + +// removes exact banned allow rules only when the policy needs changing. +async fn clean_rules_file(policy_path: &Path, banned_prefixes: &[&[&str]]) -> io::Result<()> { + let contents = match tokio::fs::read_to_string(policy_path).await { + Ok(contents) => contents, + Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(()), + Err(err) => return Err(err), + }; + if strip_banned_allow_rules(&contents, banned_prefixes) == contents { + return Ok(()); + } + + let mut file = match tokio::fs::OpenOptions::new() + .read(true) + .write(true) + .open(policy_path) + .await + { + Ok(file) => file, + Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(()), + Err(err) => return Err(err), + }; + + let mut contents = String::new(); + file.read_to_string(&mut contents).await?; + let retained = strip_banned_allow_rules(&contents, banned_prefixes); + if retained == contents { + return Ok(()); + } + + file.seek(SeekFrom::Start(0)).await?; + file.write_all(retained.as_bytes()).await?; + file.set_len(retained.len() as u64).await?; + Ok(()) +} + +// returns the policy text without exact banned allow rules. +fn strip_banned_allow_rules(contents: &str, banned_prefixes: &[&[&str]]) -> String { + let banned_prefixes = banned_prefixes + .iter() + .map(|prefix| { + prefix + .iter() + .map(|token| token.to_ascii_lowercase()) + .collect::>() + }) + .collect::>(); + contents + .split_inclusive('\n') + .filter(|line| !should_remove_rule(line, &banned_prefixes)) + .collect() +} + +// checks whether a line is an exact banned allow rule. +fn should_remove_rule(line: &str, banned_prefixes: &HashSet>) -> bool { + let line = line.strip_suffix('\n').unwrap_or(line); + let line = line.strip_suffix('\r').unwrap_or(line); + let Some(pattern) = line + .strip_prefix("prefix_rule(pattern=") + .and_then(|line| line.strip_suffix(r#", decision="allow")"#)) + else { + return false; + }; + let Ok(prefix) = serde_json::from_str::>(pattern) else { + return false; + }; + let prefix = prefix + .iter() + .map(|token| token.to_ascii_lowercase()) + .collect::>(); + banned_prefixes.contains(&prefix) +} + +#[cfg(test)] +#[path = "sandbox_migration_tests.rs"] +mod tests; diff --git a/codex-rs/execpolicy/src/sandbox_migration_tests.rs b/codex-rs/execpolicy/src/sandbox_migration_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..5bb6af6e256cbf1a66d0057503db2e7266a063c7 --- /dev/null +++ b/codex-rs/execpolicy/src/sandbox_migration_tests.rs @@ -0,0 +1,58 @@ +use super::*; +use pretty_assertions::assert_eq; +use tempfile::tempdir; + +#[tokio::test] +async fn removes_banned_allow_rules_once() { + const BANNED_PREFIXES: &[&[&str]] = &[ + &["cmd.exe", "/k"], + &["git"], + &["pwsh", "-ec"], + &["pwsh", "-f"], + ]; + let codex_home = tempdir().expect("create codex home"); + let policy_path = codex_home.path().join("rules/default.rules"); + std::fs::create_dir_all(policy_path.parent().expect("rules directory")) + .expect("create rules directory"); + std::fs::write( + &policy_path, + r#"prefix_rule(pattern=["git"], decision="allow") +prefix_rule(pattern=["git"], decision="prompt") +prefix_rule(pattern=["git"], decision="deny") +prefix_rule(pattern=["git", "status"], decision="allow") +prefix_rule(pattern=["CMD.EXE", "/K"], decision="allow") +prefix_rule(pattern=["PWSH", "-EC"], decision="allow") +prefix_rule(pattern=["PwSh", "-F"], decision="allow") +network_rule(host="api.github.com", protocol="https", decision="allow") +"#, + ) + .expect("write legacy policy"); + + prefix_rule_migration(codex_home.path(), &policy_path, BANNED_PREFIXES) + .await + .expect("run sandbox migration"); + assert_eq!( + std::fs::read_to_string(&policy_path).expect("read migrated policy"), + r#"prefix_rule(pattern=["git"], decision="prompt") +prefix_rule(pattern=["git"], decision="deny") +prefix_rule(pattern=["git", "status"], decision="allow") +network_rule(host="api.github.com", protocol="https", decision="allow") +"# + ); + assert_eq!( + std::fs::read_to_string(codex_home.path().join(MIGRATION_MARKER_FILENAME)) + .expect("read migration marker"), + "v1\n" + ); + + let post_migration_policy = r#"prefix_rule(pattern=["git"], decision="allow") +"#; + std::fs::write(&policy_path, post_migration_policy).expect("write post-migration policy"); + prefix_rule_migration(codex_home.path(), &policy_path, BANNED_PREFIXES) + .await + .expect("rerun sandbox migration"); + assert_eq!( + std::fs::read_to_string(&policy_path).expect("read post-migration policy"), + post_migration_policy + ); +} diff --git a/codex-rs/execpolicy/tests/basic.rs b/codex-rs/execpolicy/tests/basic.rs new file mode 100644 index 0000000000000000000000000000000000000000..f6a86adb27423fbc49823407284d57c0e6fc95b1 --- /dev/null +++ b/codex-rs/execpolicy/tests/basic.rs @@ -0,0 +1,963 @@ +#![allow(clippy::expect_used)] +use std::any::Any; +use std::fs; +use std::path::PathBuf; +use std::sync::Arc; + +use anyhow::Context; +use anyhow::Result; +use codex_execpolicy::Decision; +use codex_execpolicy::Error; +use codex_execpolicy::Evaluation; +use codex_execpolicy::MatchOptions; +use codex_execpolicy::NetworkRuleProtocol; +use codex_execpolicy::PatternToken; +use codex_execpolicy::Policy; +use codex_execpolicy::PolicyParser; +use codex_execpolicy::PrefixPattern; +use codex_execpolicy::PrefixRule; +use codex_execpolicy::RuleMatch; +use codex_execpolicy::RuleRef; +use codex_execpolicy::blocking_append_allow_prefix_rule; +use codex_utils_absolute_path::AbsolutePathBuf; +use pretty_assertions::assert_eq; +use tempfile::tempdir; + +fn tokens(cmd: &[&str]) -> Vec { + cmd.iter().map(std::string::ToString::to_string).collect() +} + +fn allow_all(_: &[String]) -> Decision { + Decision::Allow +} + +fn prompt_all(_: &[String]) -> Decision { + Decision::Prompt +} + +fn absolute_path(path: &str) -> AbsolutePathBuf { + AbsolutePathBuf::try_from(path.to_string()).expect("path should be absolute") +} + +fn host_absolute_path(segments: &[&str]) -> String { + let mut path = if cfg!(windows) { + PathBuf::from(r"C:\") + } else { + PathBuf::from("/") + }; + for segment in segments { + path.push(segment); + } + path.to_string_lossy().into_owned() +} + +fn host_executable_name(name: &str) -> String { + if cfg!(windows) { + format!("{name}.exe") + } else { + name.to_string() + } +} + +fn starlark_string(value: &str) -> String { + value.replace('\\', "\\\\").replace('"', "\\\"") +} + +#[derive(Clone, Debug, Eq, PartialEq)] +enum RuleSnapshot { + Prefix(PrefixRule), +} + +fn rule_snapshots(rules: &[RuleRef]) -> Vec { + rules + .iter() + .map(|rule| { + let rule_any = rule.as_ref() as &dyn Any; + if let Some(prefix_rule) = rule_any.downcast_ref::() { + RuleSnapshot::Prefix(prefix_rule.clone()) + } else { + panic!("unexpected rule type in RuleRef: {rule:?}"); + } + }) + .collect() +} + +#[test] +fn append_allow_prefix_rule_dedupes_existing_rule() -> Result<()> { + let tmp = tempdir().context("create temp dir")?; + let policy_path = tmp.path().join("rules").join("default.rules"); + let prefix = tokens(&["python3"]); + + blocking_append_allow_prefix_rule(&policy_path, &prefix)?; + blocking_append_allow_prefix_rule(&policy_path, &prefix)?; + + let contents = fs::read_to_string(&policy_path).context("read policy")?; + assert_eq!( + contents, + r#"prefix_rule(pattern=["python3"], decision="allow") +"# + ); + Ok(()) +} + +#[test] +fn network_rules_compile_into_domain_lists() -> Result<()> { + let policy_src = r#" +network_rule(host = "google.com", protocol = "http", decision = "allow") +network_rule(host = "api.github.com", protocol = "https", decision = "allow") +network_rule(host = "blocked.example.com", protocol = "https", decision = "deny") +network_rule(host = "prompt-only.example.com", protocol = "https", decision = "prompt") + "#; + let mut parser = PolicyParser::new(); + parser.parse("network.rules", policy_src)?; + let policy = parser.build(); + + assert_eq!(policy.network_rules().len(), 4); + assert_eq!( + policy.network_rules()[1].protocol, + NetworkRuleProtocol::Https + ); + + let (allowed, denied) = policy.compiled_network_domains(); + assert_eq!( + allowed, + vec!["google.com".to_string(), "api.github.com".to_string()] + ); + assert_eq!(denied, vec!["blocked.example.com".to_string()]); + Ok(()) +} + +#[test] +fn network_rule_rejects_wildcard_hosts() { + let mut parser = PolicyParser::new(); + let err = parser + .parse( + "network.rules", + r#"network_rule(host="*", protocol="http", decision="allow")"#, + ) + .expect_err("wildcard network_rule host should fail"); + assert!(err.to_string().contains("wildcards are not allowed")); +} + +#[test] +fn basic_match() -> Result<()> { + let policy_src = r#" +prefix_rule( + pattern = ["git", "status"], +) + "#; + let mut parser = PolicyParser::new(); + parser.parse("test.rules", policy_src)?; + let policy = parser.build(); + let cmd = tokens(&["git", "status"]); + let evaluation = policy.check(&cmd, &allow_all); + assert_eq!( + Evaluation { + decision: Decision::Allow, + matched_rules: vec![RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["git", "status"]), + decision: Decision::Allow, + resolved_program: None, + justification: None, + }], + }, + evaluation + ); + Ok(()) +} + +#[test] +fn justification_is_attached_to_forbidden_matches() -> Result<()> { + let policy_src = r#" +prefix_rule( + pattern = ["rm"], + decision = "forbidden", + justification = "destructive command", +) + "#; + let mut parser = PolicyParser::new(); + parser.parse("test.rules", policy_src)?; + let policy = parser.build(); + + let evaluation = policy.check( + &tokens(&["rm", "-rf", "/some/important/folder"]), + &allow_all, + ); + assert_eq!( + Evaluation { + decision: Decision::Forbidden, + matched_rules: vec![RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["rm"]), + decision: Decision::Forbidden, + resolved_program: None, + justification: Some("destructive command".to_string()), + }], + }, + evaluation + ); + Ok(()) +} + +#[test] +fn justification_can_be_used_with_allow_decision() -> Result<()> { + let policy_src = r#" +prefix_rule( + pattern = ["ls"], + decision = "allow", + justification = "safe and commonly used", +) + "#; + let mut parser = PolicyParser::new(); + parser.parse("test.rules", policy_src)?; + let policy = parser.build(); + + let evaluation = policy.check(&tokens(&["ls", "-l"]), &prompt_all); + assert_eq!( + Evaluation { + decision: Decision::Allow, + matched_rules: vec![RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["ls"]), + decision: Decision::Allow, + resolved_program: None, + justification: Some("safe and commonly used".to_string()), + }], + }, + evaluation + ); + Ok(()) +} + +#[test] +fn justification_cannot_be_empty() { + let policy_src = r#" +prefix_rule( + pattern = ["ls"], + decision = "prompt", + justification = " ", +) + "#; + let mut parser = PolicyParser::new(); + let err = parser + .parse("test.rules", policy_src) + .expect_err("expected parse error"); + assert!( + err.to_string() + .contains("invalid rule: justification cannot be empty") + ); +} + +#[test] +fn add_prefix_rule_extends_policy() -> Result<()> { + let mut policy = Policy::empty(); + policy.add_prefix_rule(&tokens(&["ls", "-l"]), Decision::Prompt)?; + + let rules = rule_snapshots(policy.rules().get_vec("ls").context("missing ls rules")?); + assert_eq!( + vec![RuleSnapshot::Prefix(PrefixRule { + pattern: PrefixPattern { + first: Arc::from("ls"), + rest: vec![PatternToken::Single(String::from("-l"))].into(), + }, + decision: Decision::Prompt, + justification: None, + })], + rules + ); + + let evaluation = policy.check(&tokens(&["ls", "-l", "/some/important/folder"]), &allow_all); + assert_eq!( + Evaluation { + decision: Decision::Prompt, + matched_rules: vec![RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["ls", "-l"]), + decision: Decision::Prompt, + resolved_program: None, + justification: None, + }], + }, + evaluation + ); + Ok(()) +} + +#[test] +fn add_prefix_rule_rejects_empty_prefix() -> Result<()> { + let mut policy = Policy::empty(); + let result = policy.add_prefix_rule(&[], Decision::Allow); + + match result.unwrap_err() { + Error::InvalidPattern(message) => assert_eq!(message, "prefix cannot be empty"), + other => panic!("expected InvalidPattern(..), got {other:?}"), + } + Ok(()) +} + +#[test] +fn parses_multiple_policy_files() -> Result<()> { + let first_policy = r#" +prefix_rule( + pattern = ["git"], + decision = "prompt", +) + "#; + let second_policy = r#" +prefix_rule( + pattern = ["git", "commit"], + decision = "forbidden", +) + "#; + let mut parser = PolicyParser::new(); + parser.parse("first.rules", first_policy)?; + parser.parse("second.rules", second_policy)?; + let policy = parser.build(); + + let git_rules = rule_snapshots(policy.rules().get_vec("git").context("missing git rules")?); + assert_eq!( + vec![ + RuleSnapshot::Prefix(PrefixRule { + pattern: PrefixPattern { + first: Arc::from("git"), + rest: Vec::::new().into(), + }, + decision: Decision::Prompt, + justification: None, + }), + RuleSnapshot::Prefix(PrefixRule { + pattern: PrefixPattern { + first: Arc::from("git"), + rest: vec![PatternToken::Single("commit".to_string())].into(), + }, + decision: Decision::Forbidden, + justification: None, + }), + ], + git_rules + ); + + let status_eval = policy.check(&tokens(&["git", "status"]), &allow_all); + assert_eq!( + Evaluation { + decision: Decision::Prompt, + matched_rules: vec![RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["git"]), + decision: Decision::Prompt, + resolved_program: None, + justification: None, + }], + }, + status_eval + ); + + let commit_eval = policy.check(&tokens(&["git", "commit", "-m", "hi"]), &allow_all); + assert_eq!( + Evaluation { + decision: Decision::Forbidden, + matched_rules: vec![ + RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["git"]), + decision: Decision::Prompt, + resolved_program: None, + justification: None, + }, + RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["git", "commit"]), + decision: Decision::Forbidden, + resolved_program: None, + justification: None, + }, + ], + }, + commit_eval + ); + Ok(()) +} + +#[test] +fn only_first_token_alias_expands_to_multiple_rules() -> Result<()> { + let policy_src = r#" +prefix_rule( + pattern = [["bash", "sh"], ["-c", "-l"]], +) + "#; + let mut parser = PolicyParser::new(); + parser.parse("test.rules", policy_src)?; + let policy = parser.build(); + + let bash_rules = rule_snapshots( + policy + .rules() + .get_vec("bash") + .context("missing bash rules")?, + ); + let sh_rules = rule_snapshots(policy.rules().get_vec("sh").context("missing sh rules")?); + assert_eq!( + vec![RuleSnapshot::Prefix(PrefixRule { + pattern: PrefixPattern { + first: Arc::from("bash"), + rest: vec![PatternToken::Alts(vec!["-c".to_string(), "-l".to_string()])].into(), + }, + decision: Decision::Allow, + justification: None, + })], + bash_rules + ); + assert_eq!( + vec![RuleSnapshot::Prefix(PrefixRule { + pattern: PrefixPattern { + first: Arc::from("sh"), + rest: vec![PatternToken::Alts(vec!["-c".to_string(), "-l".to_string()])].into(), + }, + decision: Decision::Allow, + justification: None, + })], + sh_rules + ); + + let bash_eval = policy.check(&tokens(&["bash", "-c", "echo", "hi"]), &allow_all); + assert_eq!( + Evaluation { + decision: Decision::Allow, + matched_rules: vec![RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["bash", "-c"]), + decision: Decision::Allow, + resolved_program: None, + justification: None, + }], + }, + bash_eval + ); + + let sh_eval = policy.check(&tokens(&["sh", "-l", "echo", "hi"]), &allow_all); + assert_eq!( + Evaluation { + decision: Decision::Allow, + matched_rules: vec![RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["sh", "-l"]), + decision: Decision::Allow, + resolved_program: None, + justification: None, + }], + }, + sh_eval + ); + Ok(()) +} + +#[test] +fn tail_aliases_are_not_cartesian_expanded() -> Result<()> { + let policy_src = r#" +prefix_rule( + pattern = ["npm", ["i", "install"], ["--legacy-peer-deps", "--no-save"]], +) + "#; + let mut parser = PolicyParser::new(); + parser.parse("test.rules", policy_src)?; + let policy = parser.build(); + + let rules = rule_snapshots(policy.rules().get_vec("npm").context("missing npm rules")?); + assert_eq!( + vec![RuleSnapshot::Prefix(PrefixRule { + pattern: PrefixPattern { + first: Arc::from("npm"), + rest: vec![ + PatternToken::Alts(vec!["i".to_string(), "install".to_string()]), + PatternToken::Alts(vec![ + "--legacy-peer-deps".to_string(), + "--no-save".to_string(), + ]), + ] + .into(), + }, + decision: Decision::Allow, + justification: None, + })], + rules + ); + + let npm_i = policy.check(&tokens(&["npm", "i", "--legacy-peer-deps"]), &allow_all); + assert_eq!( + Evaluation { + decision: Decision::Allow, + matched_rules: vec![RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["npm", "i", "--legacy-peer-deps"]), + decision: Decision::Allow, + resolved_program: None, + justification: None, + }], + }, + npm_i + ); + + let npm_install = policy.check( + &tokens(&["npm", "install", "--no-save", "leftpad"]), + &allow_all, + ); + assert_eq!( + Evaluation { + decision: Decision::Allow, + matched_rules: vec![RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["npm", "install", "--no-save"]), + decision: Decision::Allow, + resolved_program: None, + justification: None, + }], + }, + npm_install + ); + Ok(()) +} + +#[test] +fn match_and_not_match_examples_are_enforced() -> Result<()> { + let policy_src = r#" +prefix_rule( + pattern = ["git", "status"], + match = [["git", "status"], "git status"], + not_match = [ + ["git", "--config", "color.status=always", "status"], + "git --config color.status=always status", + ], +) + "#; + let mut parser = PolicyParser::new(); + parser.parse("test.rules", policy_src)?; + let policy = parser.build(); + let match_eval = policy.check(&tokens(&["git", "status"]), &allow_all); + assert_eq!( + Evaluation { + decision: Decision::Allow, + matched_rules: vec![RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["git", "status"]), + decision: Decision::Allow, + resolved_program: None, + justification: None, + }], + }, + match_eval + ); + + let no_match_eval = policy.check( + &tokens(&["git", "--config", "color.status=always", "status"]), + &allow_all, + ); + assert_eq!( + Evaluation { + decision: Decision::Allow, + matched_rules: vec![RuleMatch::HeuristicsRuleMatch { + command: tokens(&["git", "--config", "color.status=always", "status",]), + decision: Decision::Allow, + }], + }, + no_match_eval + ); + Ok(()) +} + +#[test] +fn strictest_decision_wins_across_matches() -> Result<()> { + let policy_src = r#" +prefix_rule( + pattern = ["git"], + decision = "prompt", +) +prefix_rule( + pattern = ["git", "commit"], + decision = "forbidden", +) + "#; + let mut parser = PolicyParser::new(); + parser.parse("test.rules", policy_src)?; + let policy = parser.build(); + + let commit = policy.check(&tokens(&["git", "commit", "-m", "hi"]), &allow_all); + assert_eq!( + Evaluation { + decision: Decision::Forbidden, + matched_rules: vec![ + RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["git"]), + decision: Decision::Prompt, + resolved_program: None, + justification: None, + }, + RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["git", "commit"]), + decision: Decision::Forbidden, + resolved_program: None, + justification: None, + }, + ], + }, + commit + ); + Ok(()) +} + +#[test] +fn strictest_decision_across_multiple_commands() -> Result<()> { + let policy_src = r#" +prefix_rule( + pattern = ["git"], + decision = "prompt", +) +prefix_rule( + pattern = ["git", "commit"], + decision = "forbidden", +) + "#; + let mut parser = PolicyParser::new(); + parser.parse("test.rules", policy_src)?; + let policy = parser.build(); + + let commands = vec![ + tokens(&["git", "status"]), + tokens(&["git", "commit", "-m", "hi"]), + ]; + + let evaluation = policy.check_multiple(&commands, &allow_all); + assert_eq!( + Evaluation { + decision: Decision::Forbidden, + matched_rules: vec![ + RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["git"]), + decision: Decision::Prompt, + resolved_program: None, + justification: None, + }, + RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["git"]), + decision: Decision::Prompt, + resolved_program: None, + justification: None, + }, + RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["git", "commit"]), + decision: Decision::Forbidden, + resolved_program: None, + justification: None, + }, + ], + }, + evaluation + ); + Ok(()) +} + +#[test] +fn heuristics_match_is_returned_when_no_policy_matches() { + let policy = Policy::empty(); + let command = tokens(&["python"]); + + let evaluation = policy.check(&command, &prompt_all); + assert_eq!( + Evaluation { + decision: Decision::Prompt, + matched_rules: vec![RuleMatch::HeuristicsRuleMatch { + command, + decision: Decision::Prompt, + }], + }, + evaluation + ); +} + +#[test] +fn parses_host_executable_paths() -> Result<()> { + let homebrew_git = host_absolute_path(&["opt", "homebrew", "bin", "git"]); + let usr_git = host_absolute_path(&["usr", "bin", "git"]); + let homebrew_git_literal = starlark_string(&homebrew_git); + let usr_git_literal = starlark_string(&usr_git); + let policy_src = format!( + r#" +host_executable( + name = "git", + paths = [ + "{homebrew_git_literal}", + "{usr_git_literal}", + "{usr_git_literal}", + ], +) +"# + ); + let mut parser = PolicyParser::new(); + parser.parse("test.rules", &policy_src)?; + let policy = parser.build(); + + assert_eq!( + policy + .host_executables() + .get("git") + .expect("missing git host executable") + .as_ref(), + [absolute_path(&homebrew_git), absolute_path(&usr_git)] + ); + Ok(()) +} + +#[test] +fn host_executable_rejects_non_absolute_path() { + let policy_src = r#" +host_executable(name = "git", paths = ["git"]) +"#; + let mut parser = PolicyParser::new(); + let err = parser + .parse("test.rules", policy_src) + .expect_err("expected parse error"); + assert!( + err.to_string() + .contains("host_executable paths must be absolute") + ); +} + +#[test] +fn host_executable_rejects_name_with_path_separator() { + let git_path = host_absolute_path(&["usr", "bin", "git"]); + let git_path_literal = starlark_string(&git_path); + let policy_src = + format!(r#"host_executable(name = "{git_path_literal}", paths = ["{git_path_literal}"])"#); + let mut parser = PolicyParser::new(); + let err = parser + .parse("test.rules", &policy_src) + .expect_err("expected parse error"); + assert!( + err.to_string() + .contains("host_executable name must be a bare executable name") + ); +} + +#[test] +fn host_executable_rejects_path_with_wrong_basename() { + let rg_path = host_absolute_path(&["usr", "bin", "rg"]); + let rg_path_literal = starlark_string(&rg_path); + let policy_src = format!(r#"host_executable(name = "git", paths = ["{rg_path_literal}"])"#); + let mut parser = PolicyParser::new(); + let err = parser + .parse("test.rules", &policy_src) + .expect_err("expected parse error"); + assert!(err.to_string().contains("must have basename `git`")); +} + +#[test] +fn host_executable_last_definition_wins() -> Result<()> { + let usr_git = host_absolute_path(&["usr", "bin", "git"]); + let homebrew_git = host_absolute_path(&["opt", "homebrew", "bin", "git"]); + let usr_git_literal = starlark_string(&usr_git); + let homebrew_git_literal = starlark_string(&homebrew_git); + let mut parser = PolicyParser::new(); + parser.parse( + "shared.rules", + &format!(r#"host_executable(name = "git", paths = ["{usr_git_literal}"])"#), + )?; + parser.parse( + "user.rules", + &format!(r#"host_executable(name = "git", paths = ["{homebrew_git_literal}"])"#), + )?; + let policy = parser.build(); + + assert_eq!( + policy + .host_executables() + .get("git") + .expect("missing git host executable") + .as_ref(), + [absolute_path(&homebrew_git)] + ); + Ok(()) +} + +#[test] +fn host_executable_resolution_uses_basename_rule_when_allowed() -> Result<()> { + let git_name = host_executable_name("git"); + let git_path = host_absolute_path(&["usr", "bin", &git_name]); + let git_path_literal = starlark_string(&git_path); + let policy_src = format!( + r#" +prefix_rule(pattern = ["git", "status"], decision = "prompt") +host_executable(name = "git", paths = ["{git_path_literal}"]) +"# + ); + let mut parser = PolicyParser::new(); + parser.parse("test.rules", &policy_src)?; + let policy = parser.build(); + + let evaluation = policy.check_with_options( + &[git_path.clone(), "status".to_string()], + &allow_all, + &MatchOptions { + resolve_host_executables: true, + }, + ); + assert_eq!( + evaluation, + Evaluation { + decision: Decision::Prompt, + matched_rules: vec![RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["git", "status"]), + decision: Decision::Prompt, + resolved_program: Some(absolute_path(&git_path)), + justification: None, + }], + } + ); + Ok(()) +} + +#[test] +fn prefix_rule_examples_honor_host_executable_resolution() -> Result<()> { + let allowed_git_name = host_executable_name("git"); + let allowed_git = host_absolute_path(&["usr", "bin", &allowed_git_name]); + let other_git = host_absolute_path(&["opt", "homebrew", "bin", &allowed_git_name]); + let allowed_git_literal = starlark_string(&allowed_git); + let other_git_literal = starlark_string(&other_git); + let policy_src = format!( + r#" +prefix_rule( + pattern = ["git", "status"], + match = [["{allowed_git_literal}", "status"]], + not_match = [["{other_git_literal}", "status"]], +) +host_executable(name = "git", paths = ["{allowed_git_literal}"]) +"# + ); + + let mut parser = PolicyParser::new(); + parser.parse("test.rules", &policy_src)?; + + Ok(()) +} + +#[test] +fn host_executable_resolution_respects_explicit_empty_allowlist() -> Result<()> { + let policy_src = r#" +prefix_rule(pattern = ["git"], decision = "prompt") +host_executable(name = "git", paths = []) +"#; + let mut parser = PolicyParser::new(); + parser.parse("test.rules", policy_src)?; + let policy = parser.build(); + let git_path = host_absolute_path(&["usr", "bin", "git"]); + + let evaluation = policy.check_with_options( + &[git_path.clone(), "status".to_string()], + &allow_all, + &MatchOptions { + resolve_host_executables: true, + }, + ); + assert_eq!( + evaluation, + Evaluation { + decision: Decision::Allow, + matched_rules: vec![RuleMatch::HeuristicsRuleMatch { + command: vec![git_path, "status".to_string()], + decision: Decision::Allow, + }], + } + ); + Ok(()) +} + +#[test] +fn host_executable_resolution_ignores_path_not_in_allowlist() -> Result<()> { + let allowed_git = host_absolute_path(&["usr", "bin", "git"]); + let other_git = host_absolute_path(&["opt", "homebrew", "bin", "git"]); + let allowed_git_literal = starlark_string(&allowed_git); + let policy_src = format!( + r#" +prefix_rule(pattern = ["git"], decision = "prompt") +host_executable(name = "git", paths = ["{allowed_git_literal}"]) +"# + ); + let mut parser = PolicyParser::new(); + parser.parse("test.rules", &policy_src)?; + let policy = parser.build(); + + let evaluation = policy.check_with_options( + &[other_git.clone(), "status".to_string()], + &allow_all, + &MatchOptions { + resolve_host_executables: true, + }, + ); + assert_eq!( + evaluation, + Evaluation { + decision: Decision::Allow, + matched_rules: vec![RuleMatch::HeuristicsRuleMatch { + command: vec![other_git, "status".to_string()], + decision: Decision::Allow, + }], + } + ); + Ok(()) +} + +#[test] +fn host_executable_resolution_falls_back_without_mapping() -> Result<()> { + let policy_src = r#" +prefix_rule(pattern = ["git"], decision = "prompt") +"#; + let mut parser = PolicyParser::new(); + parser.parse("test.rules", policy_src)?; + let policy = parser.build(); + let git_path = host_absolute_path(&["usr", "bin", "git"]); + + let evaluation = policy.check_with_options( + &[git_path.clone(), "status".to_string()], + &allow_all, + &MatchOptions { + resolve_host_executables: true, + }, + ); + assert_eq!( + evaluation, + Evaluation { + decision: Decision::Prompt, + matched_rules: vec![RuleMatch::PrefixRuleMatch { + matched_prefix: tokens(&["git"]), + decision: Decision::Prompt, + resolved_program: Some(absolute_path(&git_path)), + justification: None, + }], + } + ); + Ok(()) +} + +#[test] +fn host_executable_resolution_does_not_override_exact_match() -> Result<()> { + let git_path = host_absolute_path(&["usr", "bin", "git"]); + let git_path_literal = starlark_string(&git_path); + let policy_src = format!( + r#" +prefix_rule(pattern = ["{git_path_literal}"], decision = "allow") +prefix_rule(pattern = ["git"], decision = "prompt") +host_executable(name = "git", paths = ["{git_path_literal}"]) +"# + ); + let mut parser = PolicyParser::new(); + parser.parse("test.rules", &policy_src)?; + let policy = parser.build(); + + let evaluation = policy.check_with_options( + &[git_path.clone(), "status".to_string()], + &allow_all, + &MatchOptions { + resolve_host_executables: true, + }, + ); + assert_eq!( + evaluation, + Evaluation { + decision: Decision::Allow, + matched_rules: vec![RuleMatch::PrefixRuleMatch { + matched_prefix: vec![git_path], + decision: Decision::Allow, + resolved_program: None, + justification: None, + }], + } + ); + Ok(()) +} diff --git a/codex-rs/file-system/src/find_up.rs b/codex-rs/file-system/src/find_up.rs new file mode 100644 index 0000000000000000000000000000000000000000..a5a071d2843ccd19f11a73aed3af54e676b9e876 --- /dev/null +++ b/codex-rs/file-system/src/find_up.rs @@ -0,0 +1,126 @@ +use crate::ExecutorFileSystem; +use crate::FileSystemResult; +use crate::FileSystemSandboxContext; +use codex_utils_absolute_path::AbsolutePathBuf; +use codex_utils_path_uri::PathUri; +use futures::StreamExt; +use std::io; + +// Keep enough ordinary metadata calls in flight to cover typical ancestor chains in one remote +// round trip, while leaving room for independent startup discovery to run at the same time. +const MAX_CONCURRENT_PROBES: usize = 256; + +/// Controls how an upward marker search handles metadata errors other than `NotFound`. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum FindUpErrorPolicy { + /// Return the first error in lexical search order. + Propagate, + /// Treat errors as missing markers and continue searching. + Ignore, +} + +/// Finds the nearest ancestor containing one of the provided marker names. +/// +/// Marker paths are probed in lexical order from `start` toward the filesystem root. A bounded +/// number of ordinary metadata calls are kept in flight so remote filesystems can pipeline them +/// without requiring a batch protocol operation. +pub async fn find_nearest_ancestor_with_markers( + file_system: &dyn ExecutorFileSystem, + start: &PathUri, + markers: Vec, + error_policy: FindUpErrorPolicy, + sandbox: Option<&FileSystemSandboxContext>, +) -> FileSystemResult> { + find_nearest_ancestor( + file_system, + start.clone(), + markers, + PathUri::parent, + |ancestor, marker| { + ancestor + .join(marker) + .map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err)) + }, + error_policy, + sandbox, + ) + .await +} + +/// Finds the nearest native ancestor containing one of the provided marker names. +/// +/// Ancestors and marker paths remain native until each complete probe is converted to a URI. This +/// preserves paths that require an opaque [`PathUri`] fallback. +pub async fn find_nearest_native_ancestor_with_markers( + file_system: &dyn ExecutorFileSystem, + start: &AbsolutePathBuf, + markers: Vec, + error_policy: FindUpErrorPolicy, + sandbox: Option<&FileSystemSandboxContext>, +) -> FileSystemResult> { + find_nearest_ancestor( + file_system, + start.clone(), + markers, + AbsolutePathBuf::parent, + |ancestor, marker| Ok(PathUri::from_abs_path(&ancestor.join(marker))), + error_policy, + sandbox, + ) + .await +} + +async fn find_nearest_ancestor( + file_system: &dyn ExecutorFileSystem, + start: P, + markers: Vec, + parent: Parent, + mut marker_path: MarkerPath, + error_policy: FindUpErrorPolicy, + sandbox: Option<&FileSystemSandboxContext>, +) -> FileSystemResult> +where + P: Clone + Send, + Parent: FnMut(&P) -> Option

+ Send, + MarkerPath: FnMut(&P, &str) -> FileSystemResult + Send, +{ + let mut ancestors = std::iter::successors(Some(start), parent); + let mut ancestor = ancestors.next(); + let mut marker_index = 0; + let probes = std::iter::from_fn(move || { + let current_ancestor = ancestor.clone()?; + let marker = markers.get(marker_index)?; + let marker_path = marker_path(¤t_ancestor, marker); + + marker_index += 1; + if marker_index == markers.len() { + marker_index = 0; + ancestor = ancestors.next(); + } + + Some((current_ancestor, marker_path)) + }); + let mut results = futures::stream::iter(probes) + .map(|(ancestor, marker_path)| async move { + let marker_path = marker_path?; + match file_system + .get_metadata(&marker_path, crate::GetMetadataOptions::default(), sandbox) + .await + { + Ok(_) => Ok(Some(ancestor)), + Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(None), + Err(err) => match error_policy { + FindUpErrorPolicy::Propagate => Err(err), + FindUpErrorPolicy::Ignore => Ok(None), + }, + } + }) + .buffered(MAX_CONCURRENT_PROBES); + + while let Some(result) = results.next().await { + if let Some(ancestor) = result? { + return Ok(Some(ancestor)); + } + } + Ok(None) +} diff --git a/codex-rs/file-system/src/lib.rs b/codex-rs/file-system/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..ba825c5ce8119d254f6fd4319073d6b858c0301d --- /dev/null +++ b/codex-rs/file-system/src/lib.rs @@ -0,0 +1,592 @@ +mod find_up; + +use bytes::Bytes; +use codex_protocol::config_types::WindowsSandboxLevel; +use codex_protocol::config_types::WindowsSandboxProxySettingsMode; +use codex_protocol::models::ManagedFileSystemPermissions; +use codex_protocol::models::PermissionProfile; +use codex_protocol::models::SandboxEnforcement; +use codex_protocol::permissions::FileSystemAccessMode; +use codex_protocol::permissions::FileSystemPath; +use codex_protocol::permissions::FileSystemSandboxEntry; +use codex_protocol::permissions::FileSystemSandboxEntryMissingPathBehavior; +use codex_protocol::permissions::FileSystemSandboxKind; +use codex_protocol::permissions::FileSystemSandboxPolicy; +use codex_protocol::permissions::FileSystemSandboxPolicyContext; +use codex_protocol::permissions::FileSystemSpecialPath; +use codex_protocol::permissions::NetworkSandboxPolicy; +use codex_protocol::protocol::SandboxPolicy; +use codex_utils_path_uri::PathUri; +pub use find_up::FindUpErrorPolicy; +pub use find_up::find_nearest_ancestor_with_markers; +pub use find_up::find_nearest_native_ancestor_with_markers; +use futures::Stream; +use serde::Deserialize; +use serde::Serialize; +use std::future::Future; +use std::io; +use std::num::NonZeroUsize; +use std::path::Path; +use std::pin::Pin; +use std::task::Context; +use std::task::Poll; + +/// Maximum chunk size returned by [`ExecutorFileSystem::read_file_stream`]. +pub const FILE_READ_CHUNK_SIZE: usize = 1024 * 1024; +/// Maximum accepted directory depth for a filesystem walk. +pub const MAX_WALK_DEPTH: usize = 64; +/// Maximum accepted directory count, including the walk root. +pub const MAX_WALK_DIRECTORIES: usize = 10_000; +/// Maximum accepted number of directory entries to examine. +pub const MAX_WALK_ENTRIES: usize = 50_000; +/// Maximum estimated size of a walk response. +pub const MAX_WALK_RESPONSE_BYTES: usize = 4 * 1024 * 1024; +/// Per-entry or per-error overhead charged to the walk response budget. +pub const WALK_RESPONSE_ITEM_OVERHEAD_BYTES: usize = 64; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct ReadFileOptions { + pub follow_symlinks: bool, +} + +impl Default for ReadFileOptions { + fn default() -> Self { + Self { + follow_symlinks: true, + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct WriteFileOptions { + pub follow_symlinks: bool, +} + +impl Default for WriteFileOptions { + fn default() -> Self { + Self { + follow_symlinks: true, + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct GetMetadataOptions { + pub follow_symlinks: bool, +} + +impl Default for GetMetadataOptions { + fn default() -> Self { + Self { + follow_symlinks: true, + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct CreateDirectoryOptions { + pub recursive: bool, + pub follow_symlinks: bool, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct RemoveOptions { + pub recursive: bool, + pub force: bool, + pub follow_symlinks: bool, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct CopyOptions { + pub recursive: bool, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct FileMetadata { + pub is_directory: bool, + pub is_file: bool, + pub is_symlink: bool, + /// Size in bytes. + pub size: u64, + pub created_at_ms: i64, + pub modified_at_ms: i64, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ReadDirectoryEntry { + pub file_name: String, + pub is_directory: bool, + pub is_file: bool, +} + +/// Bounds for a recursive filesystem walk. +#[derive(Clone, Copy, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub struct WalkOptions { + /// Maximum directory depth below the root that may be traversed. + pub max_depth: usize, + /// Maximum number of directories that may be traversed, including the root. + pub max_directories: usize, + /// Maximum number of directory entries that may be examined. + pub max_entries: usize, + /// Whether directory symlinks should be followed. + pub follow_directory_symlinks: bool, + /// Whether directories whose names start with `.` should be returned but not traversed. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub prune_hidden_directories: bool, +} + +/// Type of a filesystem entry returned by a walk. +#[derive(Clone, Copy, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub enum WalkEntryKind { + Directory, + File, +} + +/// One entry returned by a walk. +#[derive(Clone, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub struct WalkEntry { + pub path: PathUri, + pub kind: WalkEntryKind, +} + +/// A descendant that could not be inspected during a walk. +#[derive(Clone, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub struct WalkError { + pub path: PathUri, + pub message: String, +} + +/// Entries and recoverable errors collected by a bounded walk. +#[derive(Clone, Debug, Default, Eq, PartialEq, serde::Deserialize, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub struct WalkOutcome { + pub entries: Vec, + pub errors: Vec, + pub truncated: bool, +} + +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ExecFileSystemPath { + Path { path: PathUri }, + GlobPattern { pattern: String }, + Special { value: FileSystemSpecialPath }, +} + +impl From for ExecFileSystemPath { + fn from(value: FileSystemPath) -> Self { + match value { + FileSystemPath::Path { path } => Self::Path { path }, + FileSystemPath::GlobPattern { pattern } => Self::GlobPattern { pattern }, + FileSystemPath::Special { value } => Self::Special { value }, + } + } +} + +impl TryFrom for FileSystemPath { + type Error = io::Error; + + fn try_from(value: ExecFileSystemPath) -> Result { + Ok(match value { + ExecFileSystemPath::Path { path } => { + path.to_abs_path()?; + Self::Path { path } + } + ExecFileSystemPath::GlobPattern { pattern } => Self::GlobPattern { pattern }, + ExecFileSystemPath::Special { value } => Self::Special { value }, + }) + } +} + +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +pub struct ExecFileSystemSandboxEntry { + pub path: ExecFileSystemPath, + pub access: FileSystemAccessMode, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub missing_path_behavior: Option, +} + +impl From for ExecFileSystemSandboxEntry { + fn from(value: FileSystemSandboxEntry) -> Self { + Self { + path: value.path.into(), + access: value.access, + missing_path_behavior: value.missing_path_behavior, + } + } +} + +impl TryFrom for FileSystemSandboxEntry { + type Error = io::Error; + + fn try_from(value: ExecFileSystemSandboxEntry) -> Result { + Ok(Self { + path: value.path.try_into()?, + access: value.access, + missing_path_behavior: value.missing_path_behavior, + }) + } +} + +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ExecManagedFileSystemPermissions { + Restricted { + entries: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + glob_scan_max_depth: Option, + }, + Unrestricted, +} + +impl From for ExecManagedFileSystemPermissions { + fn from(value: ManagedFileSystemPermissions) -> Self { + match value { + ManagedFileSystemPermissions::Restricted { + entries, + glob_scan_max_depth, + } => Self::Restricted { + entries: entries.into_iter().map(Into::into).collect(), + glob_scan_max_depth, + }, + ManagedFileSystemPermissions::Unrestricted => Self::Unrestricted, + } + } +} + +impl TryFrom for ManagedFileSystemPermissions { + type Error = io::Error; + + fn try_from(value: ExecManagedFileSystemPermissions) -> Result { + Ok(match value { + ExecManagedFileSystemPermissions::Restricted { + entries, + glob_scan_max_depth, + } => Self::Restricted { + entries: entries + .into_iter() + .map(TryInto::try_into) + .collect::>()?, + glob_scan_max_depth, + }, + ExecManagedFileSystemPermissions::Unrestricted => Self::Unrestricted, + }) + } +} + +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ExecPermissionProfile { + Managed { + file_system: ExecManagedFileSystemPermissions, + network: NetworkSandboxPolicy, + }, + Disabled, + External { + network: NetworkSandboxPolicy, + }, +} + +impl From for ExecPermissionProfile { + fn from(value: PermissionProfile) -> Self { + match value { + PermissionProfile::Managed { + file_system, + network, + } => Self::Managed { + file_system: file_system.into(), + network, + }, + PermissionProfile::Disabled => Self::Disabled, + PermissionProfile::External { network } => Self::External { network }, + } + } +} + +impl TryFrom for PermissionProfile { + type Error = io::Error; + + fn try_from(value: ExecPermissionProfile) -> Result { + Ok(match value { + ExecPermissionProfile::Managed { + file_system, + network, + } => Self::Managed { + file_system: file_system.try_into()?, + network, + }, + ExecPermissionProfile::Disabled => Self::Disabled, + ExecPermissionProfile::External { network } => Self::External { network }, + }) + } +} + +/// Windows sandbox choice encoded in executor RPCs. +/// +/// The serialized field retains its legacy `windowsSandboxLevel` name for compatibility, but MXC +/// is a sandbox implementation rather than a RestrictedToken level. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub enum WindowsSandboxSelection { + #[default] + Disabled, + RestrictedToken, + Elevated, + Mxc, +} + +impl From for WindowsSandboxSelection { + fn from(level: WindowsSandboxLevel) -> Self { + match level { + WindowsSandboxLevel::Disabled => Self::Disabled, + WindowsSandboxLevel::RestrictedToken => Self::RestrictedToken, + WindowsSandboxLevel::Elevated => Self::Elevated, + } + } +} + +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct FileSystemSandboxContext { + pub permissions: ExecPermissionProfile, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cwd: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub workspace_roots: Vec, + /// Executor-local user home used to resolve home-relative policy paths. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub user_home_dir: Option, + /// Executor-local default directories used to resolve `:tmpdir` policy entries. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub temporary_directories: Option>, + #[serde(rename = "windowsSandboxLevel")] + pub windows_sandbox_selection: WindowsSandboxSelection, + #[serde(default)] + pub windows_sandbox_private_desktop: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub windows_sandbox_proxy_settings_mode: Option, + #[serde(default)] + pub use_legacy_landlock: bool, +} + +impl FileSystemSandboxContext { + pub fn from_legacy_sandbox_policy( + sandbox_policy: SandboxPolicy, + cwd: PathUri, + ) -> io::Result { + // Legacy policy projection materializes native roots, so convert at the receiving-host + // boundary while retaining the URI in the resulting sandbox context. + let native_cwd = cwd.to_abs_path()?; + let file_system_sandbox_policy = + FileSystemSandboxPolicy::from_legacy_sandbox_policy_for_cwd( + &sandbox_policy, + &native_cwd, + ); + let permissions = PermissionProfile::from_runtime_permissions_with_enforcement( + SandboxEnforcement::from_legacy_sandbox_policy(&sandbox_policy), + &file_system_sandbox_policy, + NetworkSandboxPolicy::from(&sandbox_policy), + ); + Ok(Self::from_permission_profile_with_cwd(permissions, cwd)) + } + + pub fn from_permission_profile(permissions: PermissionProfile) -> Self { + Self::from_permissions_and_cwd(permissions, /*cwd*/ None) + } + + pub fn from_permission_profile_with_cwd(permissions: PermissionProfile, cwd: PathUri) -> Self { + Self::from_permissions_and_cwd(permissions, Some(cwd)) + } + + fn from_permissions_and_cwd(permissions: PermissionProfile, cwd: Option) -> Self { + let workspace_roots = cwd.iter().cloned().collect(); + Self { + permissions: permissions.into(), + cwd, + workspace_roots, + user_home_dir: None, + temporary_directories: None, + windows_sandbox_selection: WindowsSandboxSelection::Disabled, + windows_sandbox_private_desktop: false, + windows_sandbox_proxy_settings_mode: None, + use_legacy_landlock: false, + } + } + + pub fn should_run_in_sandbox(&self) -> bool { + let Ok(permissions) = PermissionProfile::try_from(self.permissions.clone()) else { + // A sandbox context for another host must not select the unsandboxed filesystem. + return true; + }; + let file_system_policy = permissions.file_system_sandbox_policy(); + matches!(file_system_policy.kind, FileSystemSandboxKind::Restricted) + && !file_system_policy.has_full_disk_write_access() + } + + /// Whether this context selects either supported Windows sandbox implementation. + pub fn windows_sandbox_is_requested(&self) -> bool { + self.windows_sandbox_selection != WindowsSandboxSelection::Disabled + } + + /// Borrows the executor-owned paths needed to interpret filesystem policy entries. + /// + /// Returns `None` after the cwd has been omitted from a transport context. + pub fn policy_context(&self) -> Option> { + Some(FileSystemSandboxPolicyContext { + cwd: self.cwd.as_ref()?, + workspace_roots: &self.workspace_roots, + user_home_dir: self.user_home_dir.as_ref(), + temporary_directories: self.temporary_directories.as_deref(), + }) + } + + pub fn has_cwd_dependent_permissions(&self) -> bool { + match &self.permissions { + ExecPermissionProfile::Managed { + file_system: ExecManagedFileSystemPermissions::Restricted { entries, .. }, + .. + } => entries.iter().any(|entry| match &entry.path { + ExecFileSystemPath::GlobPattern { pattern } => !Path::new(pattern).is_absolute(), + ExecFileSystemPath::Special { + value: FileSystemSpecialPath::ProjectRoots { .. }, + } => true, + ExecFileSystemPath::Path { .. } | ExecFileSystemPath::Special { .. } => false, + }), + ExecPermissionProfile::Managed { + file_system: ExecManagedFileSystemPermissions::Unrestricted, + .. + } + | ExecPermissionProfile::Disabled + | ExecPermissionProfile::External { .. } => false, + } + } + + pub fn drop_cwd_if_unused(mut self) -> Self { + if !self.has_cwd_dependent_permissions() { + self.cwd = None; + self.workspace_roots.clear(); + } + self + } +} + +pub type FileSystemResult = io::Result; + +/// Future returned by [`ExecutorFileSystem`] operations. +pub type ExecutorFileSystemFuture<'a, T> = + Pin> + Send + 'a>>; + +/// Stream of immutable chunks read from an [`ExecutorFileSystem`]. +pub struct FileSystemReadStream { + inner: Pin> + Send + 'static>>, +} + +impl FileSystemReadStream { + /// Wraps a filesystem byte stream. + pub fn new(stream: impl Stream> + Send + 'static) -> Self { + Self { + inner: Box::pin(stream), + } + } +} + +impl Stream for FileSystemReadStream { + type Item = FileSystemResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.inner.as_mut().poll_next(cx) + } +} + +/// Abstract filesystem access used by components that may operate locally or via +/// a remote environment. +pub trait ExecutorFileSystem: Send + Sync { + /// Resolves a path within this filesystem. + fn canonicalize<'a>( + &'a self, + path: &'a PathUri, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, PathUri>; + + fn read_file<'a>( + &'a self, + path: &'a PathUri, + options: ReadFileOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, Vec>; + + /// Reads a file as a stream of chunks no larger than [`FILE_READ_CHUNK_SIZE`]. + fn read_file_stream<'a>( + &'a self, + path: &'a PathUri, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, FileSystemReadStream>; + + /// Reads a file and decodes it as UTF-8 text. + fn read_file_text<'a>( + &'a self, + path: &'a PathUri, + options: ReadFileOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, String> { + Box::pin(async move { + let bytes = self.read_file(path, options, sandbox).await?; + String::from_utf8(bytes).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err)) + }) + } + + fn write_file<'a>( + &'a self, + path: &'a PathUri, + contents: Vec, + options: WriteFileOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, ()>; + + fn create_directory<'a>( + &'a self, + path: &'a PathUri, + create_directory_options: CreateDirectoryOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, ()>; + + fn get_metadata<'a>( + &'a self, + path: &'a PathUri, + options: GetMetadataOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, FileMetadata>; + + fn read_directory<'a>( + &'a self, + path: &'a PathUri, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, Vec>; + + /// Recursively lists descendants, optionally following directory symlinks. + fn walk<'a>( + &'a self, + path: &'a PathUri, + options: WalkOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, WalkOutcome>; + + fn remove<'a>( + &'a self, + path: &'a PathUri, + remove_options: RemoveOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, ()>; + + fn copy<'a>( + &'a self, + source_path: &'a PathUri, + destination_path: &'a PathUri, + copy_options: CopyOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, ()>; +} diff --git a/codex-rs/history/src/guardian_history.rs b/codex-rs/history/src/guardian_history.rs new file mode 100644 index 0000000000000000000000000000000000000000..1a94fe8b61ab255f739cdc59510de5a13b43f033 --- /dev/null +++ b/codex-rs/history/src/guardian_history.rs @@ -0,0 +1,21 @@ +//! Model-invisible checkpoint of the host's bounded Guardian transcript. + +use codex_protocol::models::ResponseItem; +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; + +/// Original review evidence, separate from the compacted model conversation. +/// Hosts enforce transcript retention limits both when saving and restoring it. +#[derive(Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +#[serde(transparent)] +pub struct GuardianHistoryCheckpoint(pub Vec); + +impl std::fmt::Debug for GuardianHistoryCheckpoint { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("GuardianHistoryCheckpoint") + .field("items", &self.0.len()) + .finish() + } +} diff --git a/codex-rs/history/src/lib.rs b/codex-rs/history/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..2be359013bc89c92e881c8918dfb26f723980525 --- /dev/null +++ b/codex-rs/history/src/lib.rs @@ -0,0 +1,517 @@ +//! Model-history and persisted-rollout domain types. + +use std::borrow::Borrow; +use std::ops::Deref; +use std::ops::DerefMut; +use std::path::PathBuf; +use std::sync::Arc; + +use codex_protocol::ThreadId; +use codex_protocol::capabilities::SelectedCapabilityRoot; +use codex_protocol::dynamic_tools::DynamicToolSpec; +use codex_protocol::mcp::McpResourceOriginCheckpoint; +use codex_protocol::models::BaseInstructions; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::InterAgentCommunication; +use codex_protocol::protocol::MultiAgentVersion; +use codex_protocol::protocol::SessionMeta; +use codex_protocol::protocol::SessionMetaLine; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadHistoryMode; +use codex_protocol::protocol::ThreadSource; +use codex_protocol::protocol::TokenUsageRecord; +use codex_protocol::protocol::TurnContextItem; +use codex_protocol::protocol::WorldStateItem; +use codex_protocol::realtime::RealtimeItem; +use codex_protocol::security_risk::SecurityRiskScore; +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Deserializer; +use serde::Serialize; +use serde::Serializer; +use serde::de::Error as _; + +/// A model-history item with room for history-only metadata. +/// +/// Persistence keeps the response item intact and stores its metadata separately. +#[derive(Debug, Clone, PartialEq)] +pub struct ResponseItemEnvelope { + pub item: ResponseItem, + pub metadata: Option, +} + +/// Metadata owned by the Codex harness and persisted with a response item. +/// +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] +pub struct CodexHarnessMetadata { + /// Whether a developer message was supplied by an app-server client. + #[serde(default)] + pub client_authored: bool, + + /// The originating history budget, including any tool-specific allowance. + /// Measured in tokens and reused when replaying persisted history. + #[serde( + default, + rename = "fallback_token_limit_override", + skip_serializing_if = "Option::is_none" + )] + pub history_truncation_token_limit: Option, + + /// Whether a response configuration update was created by the Codex harness itself. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub harness_authored_configuration: bool, + + /// Producer compatibility for an opaque compaction item, never the currently selected model. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub compaction_model_hash: Option, + + /// Thread acceptance order, independent of when queued user input reaches model history. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub user_input_order: Option, + + /// Copied parent context stays model-visible but must not become child-local authorization. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub inherited_user_message: bool, +} + +impl ResponseItemEnvelope { + /// Wraps a raw Responses API item for persisted history. + pub fn new(item: ResponseItem) -> Self { + Self { + item, + metadata: None, + } + } + + /// Unwraps the raw Responses API item. + pub fn into_item(self) -> ResponseItem { + self.item + } +} + +impl From for ResponseItemEnvelope { + fn from(item: ResponseItem) -> Self { + Self::new(item) + } +} + +impl Deref for ResponseItemEnvelope { + type Target = ResponseItem; + + fn deref(&self) -> &Self::Target { + &self.item + } +} + +impl DerefMut for ResponseItemEnvelope { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.item + } +} + +impl Borrow for ResponseItemEnvelope { + fn borrow(&self) -> &ResponseItem { + &self.item + } +} + +/// Persisted rollout item used by core history and rollout storage. +#[derive(Debug, Clone)] +pub enum RolloutItem { + SessionMeta(SessionMetaLine), + ResponseItem(ResponseItemEnvelope), + InterAgentCommunication(InterAgentCommunication), + InterAgentCommunicationMetadata { + trigger_turn: bool, + }, + Compacted(CompactedItem), + TurnContext(TurnContextItem), + TokenUsageRecord(TokenUsageRecord), + WorldState(WorldStateItem), + SecurityRiskScore(SecurityRiskScore), + RetainedContext(RetainedContextEvent), + EventMsg(EventMsg), + /// Sparse, model-invisible facts used to reconstruct realtime presentation. + RealtimeItem(RealtimeItem), +} + +impl Serialize for RolloutItem { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + rollout_payload::RolloutItemWire::from(self).serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for RolloutItem { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + rollout_payload::RolloutItemWire::deserialize(deserializer).map(Into::into) + } +} + +impl JsonSchema for RolloutItem { + fn schema_name() -> String { + "RolloutItem".to_string() + } + + fn schema_id() -> std::borrow::Cow<'static, str> { + std::borrow::Cow::Borrowed(concat!(module_path!(), "::RolloutItem")) + } + + fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::schema::Schema { + rollout_payload::RolloutItemWire::json_schema(generator) + } +} + +mod guardian_history; +mod reconciled_retained_context; +mod retained_context; + +pub use reconciled_retained_context::ReconciledRetainedContext; +pub use retained_context::RetainedContext; +pub use retained_context::RetainedContextEntry; +pub use retained_context::RetainedContextEvent; +pub use retained_context::RetainedContextOrder; +pub use retained_context::RetainedInputSource; +pub use retained_context::RetainedUserMessage; +pub use retained_context::VerifiedAnswer; +pub use retained_context::VerifiedQuestionAnswer; +mod rollout_payload; + +pub use guardian_history::GuardianHistoryCheckpoint; + +#[derive(Clone, Debug, PartialEq)] +pub struct CompactedItem { + pub message: String, + pub replacement_history: Option>, + pub guardian_history: Option, + pub retained_context: Option, + pub mcp_resource_origins: Option, + pub window_number: Option, + pub first_window_id: Option, + pub previous_window_id: Option, + pub window_id: Option, + /// Responses API ID for the model-backed compaction request, when one exists. + pub compaction_response_id: Option, + /// Snapshot of the latest reachable token usage record when this compaction was written. + /// + /// `thread/resume` can restore token usage totals from this field without scanning arbitrarily + /// far past the compaction. + pub latest_token_usage_record: Option, +} + +impl Serialize for CompactedItem { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + rollout_payload::CompactedItemWire::from(self).serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for CompactedItem { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + rollout_payload::CompactedItemWire::deserialize(deserializer)? + .try_into() + .map_err(D::Error::custom) + } +} + +impl JsonSchema for CompactedItem { + fn schema_name() -> String { + "CompactedItem".to_string() + } + + fn schema_id() -> std::borrow::Cow<'static, str> { + std::borrow::Cow::Borrowed(concat!(module_path!(), "::CompactedItem")) + } + + fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::schema::Schema { + rollout_payload::CompactedItemWire::json_schema(generator) + } +} + +impl From for ResponseItem { + fn from(value: CompactedItem) -> Self { + ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: value.message, + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + } + } +} + +/// One persisted rollout JSONL record. +/// +/// This intentionally does not implement Deserialize: JSONL readers must use +/// codex_rollout's canonical parser so nested decimal values survive the flattened envelope. +#[derive(Serialize, Clone, JsonSchema)] +pub struct RolloutLine { + pub timestamp: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ordinal: Option, + #[serde(flatten)] + pub item: RolloutItem, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ResumedHistory { + pub conversation_id: ThreadId, + pub history: Arc>, + pub rollout_path: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub enum InitialHistory { + New, + Cleared, + Resumed(ResumedHistory), + Forked(Vec), +} + +impl InitialHistory { + pub fn scan_rollout_items(&self, mut predicate: impl FnMut(&RolloutItem) -> bool) -> bool { + match self { + Self::New | Self::Cleared => false, + Self::Resumed(resumed) => resumed.history.iter().any(&mut predicate), + Self::Forked(items) => items.iter().any(predicate), + } + } + + pub fn forked_from_id(&self) -> Option { + match self { + Self::New | Self::Cleared => None, + Self::Resumed(resumed) => resumed.history.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => meta_line.meta.forked_from_id, + _ => None, + }), + Self::Forked(items) => items.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => Some(meta_line.meta.id), + _ => None, + }), + } + } + + pub fn session_cwd(&self) -> Option { + match self { + Self::New | Self::Cleared => None, + Self::Resumed(resumed) => session_cwd_from_items(&resumed.history), + Self::Forked(items) => session_cwd_from_items(items), + } + } + + pub fn get_rollout_items(&self) -> &[RolloutItem] { + match self { + Self::New | Self::Cleared => &[], + Self::Resumed(resumed) => &resumed.history, + Self::Forked(items) => items, + } + } + + pub fn get_event_msgs(&self) -> Option> { + match self { + Self::New | Self::Cleared => None, + Self::Resumed(resumed) => Some( + resumed + .history + .iter() + .filter_map(|item| match item { + RolloutItem::EventMsg(event) => Some(event.clone()), + _ => None, + }) + .collect(), + ), + Self::Forked(items) => Some( + items + .iter() + .filter_map(|item| match item { + RolloutItem::EventMsg(event) => Some(event.clone()), + _ => None, + }) + .collect(), + ), + } + } + + pub fn get_base_instructions(&self) -> Option { + match self { + Self::New | Self::Cleared => None, + Self::Resumed(resumed) => resumed.history.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => meta_line.meta.base_instructions.clone(), + _ => None, + }), + Self::Forked(items) => items.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => meta_line.meta.base_instructions.clone(), + _ => None, + }), + } + } + + pub fn get_dynamic_tools(&self) -> Option> { + match self { + Self::New | Self::Cleared => None, + Self::Resumed(resumed) => resumed.history.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => meta_line.meta.dynamic_tools.clone(), + _ => None, + }), + Self::Forked(items) => items.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => meta_line.meta.dynamic_tools.clone(), + _ => None, + }), + } + } + + pub fn get_selected_capability_roots(&self) -> Vec { + self.get_session_meta() + .map(|meta| meta.selected_capability_roots.clone()) + .unwrap_or_default() + } + + pub fn get_multi_agent_version(&self) -> Option { + match self { + Self::New | Self::Cleared => None, + Self::Resumed(resumed) => { + multi_agent_version_from_items(&resumed.history, Some(resumed.conversation_id)) + } + Self::Forked(items) => multi_agent_version_from_items(items, /*thread_id*/ None), + } + } + + pub fn get_history_mode(&self, default_history_mode: ThreadHistoryMode) -> ThreadHistoryMode { + match self { + Self::New | Self::Cleared => default_history_mode, + Self::Resumed(_) | Self::Forked(_) => self + .get_session_meta() + .map(|meta| meta.history_mode) + .unwrap_or(default_history_mode), + } + } + + pub fn get_resumed_session_sources(&self) -> Option<(SessionSource, Option)> { + let meta = self.get_resumed_session_meta()?; + Some((meta.source.clone(), meta.thread_source.clone())) + } + + pub fn get_resumed_thread_source(&self) -> Option { + self.get_resumed_session_meta() + .and_then(|meta| meta.thread_source.clone()) + } + + pub fn get_session_originator(&self) -> Option { + self.get_session_meta() + .map(|meta| meta.originator.clone()) + .filter(|originator| !originator.is_empty()) + } + + pub fn get_resumed_parent_thread_id(&self) -> Option { + self.get_resumed_session_meta() + .and_then(|meta| meta.parent_thread_id) + } + + fn get_session_meta(&self) -> Option<&SessionMeta> { + match self { + Self::New | Self::Cleared => None, + Self::Resumed(resumed) => resumed.history.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => Some(&meta_line.meta), + _ => None, + }), + Self::Forked(items) => items.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => Some(&meta_line.meta), + _ => None, + }), + } + } + + fn get_resumed_session_meta(&self) -> Option<&SessionMeta> { + match self { + Self::New | Self::Cleared | Self::Forked(_) => None, + Self::Resumed(resumed) => resumed.history.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => Some(&meta_line.meta), + _ => None, + }), + } + } +} + +fn session_cwd_from_items(items: &[RolloutItem]) -> Option { + items.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => Some(meta_line.meta.cwd.clone()), + _ => None, + }) +} + +/// Returns a thread's latest plugin selection, with a turn-context fallback. +/// +/// Forked history may contain ancestor snapshots, and compaction may append a +/// frozen turn context after an update. Neither can replace thread-owned settings. +/// Without an owned snapshot, only the latest turn context supplies the initial +/// selection; a missing field must not resurrect a selection from an older turn. +pub fn latest_disabled_plugin_ids(items: &[RolloutItem], thread_id: ThreadId) -> Option<&[String]> { + if let Some(ids) = items.iter().rev().find_map(|item| { + if let RolloutItem::EventMsg(EventMsg::ThreadSettingsApplied(event)) = item + && event.thread_id == Some(thread_id) + { + Some(event.thread_settings.disabled_plugin_ids.as_slice()) + } else { + None + } + }) { + return Some(ids); + } + items + .iter() + .rev() + .find_map(|item| match item { + RolloutItem::TurnContext(context) => Some(context), + _ => None, + }) + .and_then(|context| context.disabled_plugin_ids.as_deref()) +} + +fn multi_agent_version_from_items( + items: &[RolloutItem], + thread_id: Option, +) -> Option { + let session_meta_version = items.iter().rev().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) + if thread_id.is_none_or(|thread_id| meta_line.meta.id == thread_id) => + { + meta_line.meta.multi_agent_version + } + _ => None, + }); + + session_meta_version.or_else(|| { + items.iter().rev().find_map(|item| match item { + RolloutItem::TurnContext(turn_context) => turn_context.multi_agent_version, + RolloutItem::SessionMeta(_) + | RolloutItem::ResponseItem(_) + | RolloutItem::InterAgentCommunication(_) + | RolloutItem::InterAgentCommunicationMetadata { .. } + | RolloutItem::Compacted(_) + | RolloutItem::TokenUsageRecord(_) + | RolloutItem::WorldState(_) + | RolloutItem::RetainedContext(_) + | RolloutItem::SecurityRiskScore(_) + | RolloutItem::RealtimeItem(_) + | RolloutItem::EventMsg(_) => None, + }) + }) +} + +#[cfg(test)] +#[path = "tests.rs"] +mod tests; diff --git a/codex-rs/history/src/reconciled_retained_context.rs b/codex-rs/history/src/reconciled_retained_context.rs new file mode 100644 index 0000000000000000000000000000000000000000..1343be63474efa74b0a97901889689db118120aa --- /dev/null +++ b/codex-rs/history/src/reconciled_retained_context.rs @@ -0,0 +1,101 @@ +//! Read-only recovery of retained instructions from eligible live messages. +//! Missing retained instructions use persisted acceptance order; unknown ordering stays incomplete. +//! Recovering surviving sources cannot clear a checkpoint's record of missing instructions. + +use std::collections::HashSet; + +use crate::RetainedContext; +use crate::RetainedContextEntry; +use crate::RetainedContextOrder; +use crate::RetainedUserMessage; + +/// Retained evidence supplemented by surviving user messages without changing the checkpoint. +pub struct ReconciledRetainedContext<'a> { + retained_context: Option<&'a RetainedContext>, + recovered: Vec<(RetainedContextOrder, RetainedUserMessage)>, + /// Instructions were already missing, or a recovered source could not be ordered. + pub missing_user_messages: bool, +} + +impl<'a> ReconciledRetainedContext<'a> { + /// Reconciles eligible live user messages with retained sources in acceptance order. + /// The caller supplies original source identities and filters out non-user context. + /// Inherited sources must supply no local order; retained identities match before ordering. + /// Live messages are consumed only when retained user instructions are incomplete. + pub fn new( + retained_context: Option<&'a RetainedContext>, + live_messages: impl IntoIterator, RetainedUserMessage)>, + ) -> Self { + let mut missing_user_messages = + retained_context.is_none_or(RetainedContext::has_missing_user_messages); + let retained_entries = retained_context + .into_iter() + .flat_map(RetainedContext::ordered_entries) + .collect::>(); + let mut recovered = Vec::new(); + if retained_context.is_none_or(|context| !context.user_messages_complete()) { + let mut matched_entries = HashSet::new(); + let mut used_orders = retained_entries + .iter() + .map(|(order, _)| *order) + .collect::>(); + for (order, message) in live_messages { + if retained_entries + .iter() + .enumerate() + .any(|(index, (_, entry))| { + let RetainedContextEntry::UserMessage(retained) = entry else { + return false; + }; + let matches_source = if let Some(id) = &retained.message_id { + message.message_id.as_ref() == Some(id) + } else { + message.turn_id == retained.turn_id && message.text == retained.text + }; + matches_source && matched_entries.insert(index) + }) + { + continue; + } + // Queued steering can enter raw history after a later-accepted answer. + // Compare their persisted sequence numbers, including checkpoint-only answers. + // Parent counters cannot establish an adopted prefix position. + let Some(order) = order + .map(RetainedContextOrder::Local) + .filter(|order| used_orders.insert(*order)) + else { + missing_user_messages = true; + continue; + }; + recovered.push((order, message)); + } + } + Self { + retained_context, + recovered, + missing_user_messages, + } + } + + /// Returns retained and recovered evidence in persisted acceptance order. + pub fn ordered_entries( + &self, + ) -> impl DoubleEndedIterator)> { + let mut entries = self + .retained_context + .into_iter() + .flat_map(RetainedContext::ordered_entries) + .collect::>(); + entries.extend( + self.recovered + .iter() + .map(|(order, message)| (*order, RetainedContextEntry::UserMessage(message))), + ); + entries.sort_by_key(|(order, _)| *order); + entries.into_iter() + } +} + +#[cfg(test)] +#[path = "reconciled_retained_context_tests.rs"] +mod tests; diff --git a/codex-rs/history/src/reconciled_retained_context_tests.rs b/codex-rs/history/src/reconciled_retained_context_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..baeb79d0dcbe25f1b4daea3b53267035a1f8a3d9 --- /dev/null +++ b/codex-rs/history/src/reconciled_retained_context_tests.rs @@ -0,0 +1,131 @@ +use super::*; +use crate::RetainedContextEvent; +use crate::RetainedInputSource; +use crate::VerifiedAnswer; +use crate::VerifiedQuestionAnswer; +use pretty_assertions::assert_eq; + +fn instruction(text: &str) -> RetainedUserMessage { + RetainedUserMessage { + turn_id: "turn-1".to_owned(), + message_id: None, + text: text.to_owned(), + complete: false, + } +} + +#[test] +fn recovery_preserves_source_identity_acceptance_order_and_checkpoint_gaps() { + let initial = instruction("Inspect the deployment."); + let retained_excerpt = RetainedUserMessage { + message_id: Some("retained-message".to_owned()), + ..instruction("") + }; + let original = RetainedUserMessage { + text: "Deploy the reviewed change.".to_owned(), + ..retained_excerpt.clone() + }; + let mut retained = RetainedContext::default(); + retained.record_user_message(initial.clone(), RetainedInputSource::Local(Some(0))); + retained.record_user_message(retained_excerpt, RetainedInputSource::Local(Some(1))); + retained.record(&RetainedContextEvent::VerifiedAnswer { + answer: VerifiedAnswer { + turn_id: "answer-turn".to_owned(), + call_id: "publish-question".to_owned(), + questions: vec![VerifiedQuestionAnswer { + question: "Publish?".to_owned(), + answer: "Never publicly.".to_owned(), + }], + }, + acceptance_order: Some(3), + }); + retained.mark_user_messages_incomplete(); + let checkpoint = retained.clone(); + + let reconciled = ReconciledRetainedContext::new( + Some(&retained), + [ + // Existing identities match before requiring an order, including an omitted excerpt. + (None, initial.clone()), + (None, original), + (Some(4), instruction("Do not deploy after all.")), + // This steer arrived before the checkpoint-only answer but was recorded later. + (Some(2), instruction("Make the deployment public.")), + // A second identical message is a distinct source once the retained one matched. + (Some(5), initial), + ], + ); + + assert_eq!( + ( + reconciled + .ordered_entries() + .map(|(order, entry)| match entry { + RetainedContextEntry::UserMessage(message) => (order, message.text.as_str()), + RetainedContextEntry::VerifiedAnswer(answer) => { + (order, answer.questions[0].answer.as_str()) + } + }) + .collect::>(), + reconciled.missing_user_messages, + ), + ( + vec![ + (RetainedContextOrder::Local(0), "Inspect the deployment."), + (RetainedContextOrder::Local(1), ""), + ( + RetainedContextOrder::Local(2), + "Make the deployment public." + ), + (RetainedContextOrder::Local(3), "Never publicly."), + (RetainedContextOrder::Local(4), "Do not deploy after all."), + (RetainedContextOrder::Local(5), "Inspect the deployment."), + ], + true, + ), + ); + assert_eq!(retained, checkpoint); +} + +#[test] +fn recovery_marks_missing_and_conflicting_orders_incomplete() { + let mut retained = RetainedContext::default(); + retained.record_user_message( + instruction("Inspect the deployment."), + RetainedInputSource::Local(Some(1)), + ); + assert!(!retained.has_missing_user_messages()); + let checkpoint = retained.clone(); + + for invalid_order in [None, Some(1), Some(2)] { + let reconciled = ReconciledRetainedContext::new( + Some(&retained), + [ + (Some(2), instruction("Do not deploy.")), + (invalid_order, instruction("Invalid-order instruction.")), + ], + ); + assert_eq!( + ( + reconciled + .ordered_entries() + .map(|(order, entry)| match entry { + RetainedContextEntry::UserMessage(message) => + (order, message.text.as_str()), + RetainedContextEntry::VerifiedAnswer(_) => panic!("unexpected answer"), + }) + .collect::>(), + reconciled.missing_user_messages, + ), + ( + vec![ + (RetainedContextOrder::Local(1), "Inspect the deployment."), + (RetainedContextOrder::Local(2), "Do not deploy.") + ], + true, + ), + "invalid order: {invalid_order:?}", + ); + } + assert_eq!(retained, checkpoint); +} diff --git a/codex-rs/history/src/retained_context.rs b/codex-rs/history/src/retained_context.rs new file mode 100644 index 0000000000000000000000000000000000000000..ecc5c364c9ed7339e338e52c46e4c799bfe7c27b --- /dev/null +++ b/codex-rs/history/src/retained_context.rs @@ -0,0 +1,460 @@ +//! Bounded host-owned facts outside the model's compaction contract. +//! Checkpoints and live recording use the same admission rules; rendering is a consumer concern. +//! Adopted parent instructions precede local facts without sharing the local acceptance counter. +//! Their parent verified answers are unavailable, so adopted authorization stays incomplete. + +use std::collections::VecDeque; + +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; + +use crate::CodexHarnessMetadata; +use crate::ResponseItemEnvelope; + +const MAX_FAMILY_RECORDS: usize = 8; +const MAX_RECORD_BYTES: usize = 16_384; +const MAX_FAMILY_BYTES: usize = 65_536; + +/// Original assistant question and host-verified user reply, never an inferred permission. +#[derive(Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct VerifiedQuestionAnswer { + pub question: String, + pub answer: String, +} + +/// One accepted request_user_input response. Identity is local to the owning thread. +/// An omitted payload records incomplete evidence, rather than keeping a partial grant. +#[derive(Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct VerifiedAnswer { + pub turn_id: String, + pub call_id: String, + pub questions: Vec, +} + +/// Original user instruction, retained outside model summarization for delegated review. +/// Non-text input is not reconstructed as text; missing evidence remains explicit. +#[derive(Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct RetainedUserMessage { + pub turn_id: String, + pub message_id: Option, + pub text: String, + pub complete: bool, +} + +/// Local facts use their acceptance counter; copied parent instructions use prefix order. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum RetainedInputSource { + Local(Option), + Inherited, +} + +impl RetainedInputSource { + /// Parent counters never establish order within the owning thread. + pub fn acceptance_order(self) -> Option { + match self { + Self::Local(order) => order, + Self::Inherited => None, + } + } +} + +impl From> for RetainedInputSource { + fn from(metadata: Option<&CodexHarnessMetadata>) -> Self { + if metadata.is_some_and(|metadata| metadata.inherited_user_message) { + Self::Inherited + } else { + Self::Local(metadata.and_then(|metadata| metadata.user_input_order)) + } + } +} + +impl std::fmt::Debug for RetainedUserMessage { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("RetainedUserMessage") + .field("complete", &self.complete) + .finish_non_exhaustive() + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +struct Ordered { + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + inherited: bool, + #[serde(default)] + order: u64, + #[serde(flatten)] + value: T, +} + +impl Ordered { + fn key(&self) -> RetainedContextOrder { + if self.inherited { + RetainedContextOrder::Inherited(self.order) + } else { + RetainedContextOrder::Local(self.order) + } + } +} + +/// Inherited prefix order precedes the owning thread's local acceptance order. +/// The counters have separate scopes and must not be compared without their origin. +#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] +pub enum RetainedContextOrder { + Inherited(u64), + Local(u64), +} + +/// Borrowed host evidence across retained families. +pub enum RetainedContextEntry<'a> { + UserMessage(&'a RetainedUserMessage), + VerifiedAnswer(&'a VerifiedAnswer), +} + +impl std::fmt::Debug for VerifiedAnswer { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("VerifiedAnswer") + .field("questions", &self.questions.len()) + .finish_non_exhaustive() + } +} + +/// Sparse, model-invisible updates. Only the host may produce these records. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum RetainedContextEvent { + VerifiedAnswer { + #[serde(flatten)] + answer: VerifiedAnswer, + /// Absent in legacy events, which retain their recorded ordering. + #[serde(default, skip_serializing_if = "Option::is_none")] + acceptance_order: Option, + }, +} + +/// Bounded snapshot of retained families, persisted with the parent compaction checkpoint. +/// Facts live until their instruction boundary is rolled back; compaction does not expire them. +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct RetainedContext { + verified_answers: VecDeque>, + /// Lost evidence cannot be treated as a complete authorization history. + #[serde(rename = "incomplete")] + verified_answers_incomplete: bool, + #[serde(default)] + user_messages: VecDeque>, + /// Old checkpoints did not preserve user restrictions for delegated review. + #[serde(default = "legacy_user_messages_incomplete")] + user_messages_incomplete: bool, + #[serde(default)] + next_order: u64, +} + +fn legacy_user_messages_incomplete() -> bool { + true +} + +fn bound_family(items: &mut VecDeque>, incomplete: &mut bool) { + // Queued instructions can be recorded after later-accepted answers. Evict by + // acceptance order, not by the order in which persistence happened to finish. + items.make_contiguous().sort_by_key(Ordered::key); + while items.len() > MAX_FAMILY_RECORDS + || serde_json::to_vec(items).map_or(true, |bytes| bytes.len() > MAX_FAMILY_BYTES) + { + items.pop_front(); + *incomplete = true; + } +} + +impl RetainedUserMessage { + fn bound(&mut self) { + if serde_json::to_vec(self).map_or(true, |bytes| bytes.len() > MAX_RECORD_BYTES) { + self.text.clear(); + self.complete = false; + self.turn_id + .truncate(self.turn_id.floor_char_boundary(1_024)); + if let Some(id) = &mut self.message_id { + id.truncate(id.floor_char_boundary(1_024)); + } + } + } +} + +impl RetainedContextEvent { + /// Bounds a persisted event before it enters the rollout or the live snapshot. + pub fn bound(&mut self) { + match self { + Self::VerifiedAnswer { answer, .. } => { + if serde_json::to_vec(answer).map_or(true, |bytes| bytes.len() > MAX_RECORD_BYTES) { + answer.questions.clear(); + // IDs are correlation metadata, not model-visible authorization text. + answer + .turn_id + .truncate(answer.turn_id.floor_char_boundary(1_024)); + answer + .call_id + .truncate(answer.call_id.floor_char_boundary(1_024)); + } + } + } + } +} + +impl RetainedContext { + /// Reserves order without retaining pending input that hooks may reject or cancel. + pub fn reserve_order(&mut self) -> u64 { + let order = self.next_order; + self.next_order = self.next_order.saturating_add(1); + order + } + + fn record_order(&mut self, acceptance_order: Option) -> u64 { + let order = acceptance_order.unwrap_or(self.next_order); + self.next_order = self.next_order.max(order.saturating_add(1)); + order + } + + pub fn verified_answers(&self) -> impl DoubleEndedIterator { + self.verified_answers.iter().map(|entry| &entry.value) + } + + pub fn verified_answers_complete(&self) -> bool { + !self.verified_answers_incomplete + && self + .verified_answers + .iter() + .all(|answer| !answer.value.questions.is_empty()) + } + + pub fn user_messages_complete(&self) -> bool { + !self.has_missing_user_messages() + && self + .user_messages + .iter() + .all(|message| message.value.complete) + } + + /// Whether instruction records were lost or omitted by legacy capture. + /// Bounded excerpts keep their source record and do not set this marker. + pub fn has_missing_user_messages(&self) -> bool { + self.user_messages_incomplete + } + + /// A skipped instruction leaves a gap that later checkpoints must preserve. + pub fn mark_user_messages_incomplete(&mut self) { + self.user_messages_incomplete = true; + } + + /// Whether a root checkpoint already contains its adopted instruction prefix. + pub fn has_inherited_user_messages(&self) -> bool { + self.user_messages.iter().any(|entry| entry.inherited) + } + + /// Returns retained evidence with its origin and persisted order across both families. + /// Explicit acceptance order survives delayed recording; inherited prefix order stays separate. + pub fn ordered_entries( + &self, + ) -> impl DoubleEndedIterator)> { + let mut entries = self + .verified_answers + .iter() + .map(|entry| { + ( + entry.key(), + RetainedContextEntry::VerifiedAnswer(&entry.value), + ) + }) + .chain( + self.user_messages + .iter() + .map(|entry| (entry.key(), RetainedContextEntry::UserMessage(&entry.value))), + ) + .collect::>(); + entries.sort_by_key(|(order, _)| *order); + entries.into_iter() + } + + /// Records a delivered user item with its acceptance order. Legacy items without + /// this metadata use recording order; checkpoint/suffix replay uses the same path. + /// Inherited instructions use prefix order because their original counters belong to parents. + pub fn record_user_message( + &mut self, + mut message: RetainedUserMessage, + source: RetainedInputSource, + ) { + message.bound(); + let inherited = source == RetainedInputSource::Inherited; + // Worker forks omit parent answer records. Adopting their instructions + // cannot establish whether an omitted answer restricted an inherited grant. + self.verified_answers_incomplete |= inherited; + if let Some(index) = self.user_messages.iter().position(|entry| { + message.message_id.is_some() && entry.value.message_id == message.message_id + }) { + if self.user_messages[index].value == message + && self.user_messages[index].inherited == inherited + { + return; + } + self.user_messages.remove(index); + } + // Parent and worker counters have different scopes. Adopt the copied prefix + // in history order without advancing the worker's local acceptance counter. + let order = if inherited { + self.user_messages + .iter() + .filter(|entry| entry.inherited) + .map(|entry| entry.order.saturating_add(1)) + .max() + .unwrap_or_default() + } else { + self.record_order(source.acceptance_order()) + }; + self.user_messages.push_back(Ordered { + inherited, + order, + value: message, + }); + bound_family(&mut self.user_messages, &mut self.user_messages_incomplete); + } + + /// Recovers omitted text from an exact source without changing identity, order, or + /// completeness. Recovered excerpts must still satisfy the retained storage limits. + pub fn recover_user_message_excerpts( + &mut self, + mut excerpt_for_id: impl FnMut(&str) -> Option, + ) { + for entry in &mut self.user_messages { + let message = &mut entry.value; + if message.text.is_empty() + && !message.complete + && let Some(text) = message.message_id.as_deref().and_then(&mut excerpt_for_id) + { + message.text = text; + message.bound(); + } + } + bound_family(&mut self.user_messages, &mut self.user_messages_incomplete); + } + + /// Same-event delivery is idempotent; changed contents replace that source's record. + pub fn record(&mut self, event: &RetainedContextEvent) -> bool { + let mut event = event.clone(); + event.bound(); + match event { + RetainedContextEvent::VerifiedAnswer { + answer, + acceptance_order, + } => { + if let Some(index) = self.verified_answers.iter().position(|existing| { + existing.value.turn_id == answer.turn_id + && existing.value.call_id == answer.call_id + }) { + if self.verified_answers[index].value == answer { + return false; + } + self.verified_answers.remove(index); + } + let order = self.record_order(acceptance_order); + self.verified_answers.push_back(Ordered { + inherited: false, + order, + value: answer, + }); + bound_family( + &mut self.verified_answers, + &mut self.verified_answers_incomplete, + ); + } + } + true + } + + /// Restoring a saved thread must not bypass the live storage limits. + /// A missing checkpoint cannot establish complete historical user instructions. + /// Surviving local input metadata also advances the counter when its retained record is missing. + pub fn restore(&mut self, checkpoint: Option<&Self>, surviving_items: &[ResponseItemEnvelope]) { + *self = checkpoint.cloned().unwrap_or_else(|| Self { + user_messages_incomplete: true, + ..Self::default() + }); + for entry in &mut self.verified_answers { + let mut event = RetainedContextEvent::VerifiedAnswer { + answer: entry.value.clone(), + acceptance_order: Some(entry.order), + }; + event.bound(); + let RetainedContextEvent::VerifiedAnswer { answer, .. } = event; + entry.value = answer; + self.next_order = self.next_order.max(entry.order.saturating_add(1)); + } + for entry in &mut self.user_messages { + entry.value.bound(); + // Also repair checkpoints written before adoption recorded the answer gap. + self.verified_answers_incomplete |= entry.inherited; + if !entry.inherited { + self.next_order = self.next_order.max(entry.order.saturating_add(1)); + } + } + for order in surviving_items + .iter() + .filter_map(|item| item.metadata.as_ref()) + .filter(|metadata| !metadata.inherited_user_message) + .filter_map(|metadata| metadata.user_input_order) + { + self.next_order = self.next_order.max(order.saturating_add(1)); + } + bound_family( + &mut self.verified_answers, + &mut self.verified_answers_incomplete, + ); + bound_family(&mut self.user_messages, &mut self.user_messages_incomplete); + } + + /// Keeps legacy answers whose source calls survive when no retained instruction boundary exists. + pub fn retain_answers(&mut self, mut keep: impl FnMut(&VerifiedAnswer) -> bool) { + self.verified_answers.retain(|answer| keep(&answer.value)); + } + + /// Rolls back at the original user-message boundary, including later-accepted facts. + /// The explicit order also covers checkpoints made before a queued message was delivered. + /// Steering can share a turn ID. Legacy sources without message identity fall back to + /// source-turn removal and cannot establish complete retained user instructions. + pub fn rollback( + &mut self, + turn_ids: &[&str], + first_removed_message_id: Option<&str>, + source: RetainedInputSource, + ) { + let boundary = source.acceptance_order().map(RetainedContextOrder::Local); + if let Some(boundary) = boundary.or_else(|| { + first_removed_message_id.and_then(|id| { + self.user_messages + .iter() + .find(|message| message.value.message_id.as_deref() == Some(id)) + .map(Ordered::key) + }) + }) { + self.verified_answers.retain(|entry| entry.key() < boundary); + self.user_messages.retain(|entry| entry.key() < boundary); + return; + } + self.user_messages_incomplete |= first_removed_message_id.is_some() + || self + .user_messages + .iter() + .any(|message| turn_ids.contains(&message.value.turn_id.as_str())); + self.verified_answers.retain(|answer| { + source != RetainedInputSource::Inherited + && !turn_ids.contains(&answer.value.turn_id.as_str()) + }); + self.user_messages.retain(|message| { + (source != RetainedInputSource::Inherited || message.inherited) + && !turn_ids.contains(&message.value.turn_id.as_str()) + }); + } +} + +#[cfg(test)] +#[path = "retained_context_tests.rs"] +mod tests; diff --git a/codex-rs/history/src/retained_context_tests.rs b/codex-rs/history/src/retained_context_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..9b5bf65f99a95e23b2fc8e001e646a89019ccf9d --- /dev/null +++ b/codex-rs/history/src/retained_context_tests.rs @@ -0,0 +1,388 @@ +use super::*; +use crate::CodexHarnessMetadata; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; +use pretty_assertions::assert_eq; + +fn publish_answer() -> RetainedContextEvent { + RetainedContextEvent::VerifiedAnswer { + answer: VerifiedAnswer { + turn_id: "turn-1".to_owned(), + call_id: "ask-1".to_owned(), + questions: vec![VerifiedQuestionAnswer { + question: "Publish?".to_owned(), + answer: "Yes, but never publicly.".to_owned(), + }], + }, + acceptance_order: None, + } +} + +#[test] +fn retained_evidence_preserves_order_through_recovery_checkpoint_and_rollback() { + let mut context = RetainedContext::default(); + let first = publish_answer(); + assert!(context.record(&first)); + let before_restriction = context.clone(); + context.record_user_message( + RetainedUserMessage { + turn_id: "revocation-turn".to_owned(), + message_id: Some("revocation".to_owned()), + text: String::new(), + complete: false, + }, + RetainedInputSource::Local(None), + ); + let mut expected = context.clone(); + expected.user_messages[0].value.text = "Do not publish after all.".to_owned(); + context.recover_user_message_excerpts(|id| { + assert_eq!(id, "revocation"); + Some("Do not publish after all.".to_owned()) + }); + assert_eq!(context, expected); + context.recover_user_message_excerpts(|_| panic!("existing text must not be replaced")); + let snapshot = context.clone(); + assert!(!context.record(&first)); + assert_eq!(context, snapshot); + let checkpoint = + serde_json::from_str(&serde_json::to_string(&context).expect("retained answer fixture")) + .expect("retained answer fixture"); + let mut restored = RetainedContext::default(); + restored.restore(Some(&checkpoint), &[]); + assert_eq!(restored, snapshot); + assert_eq!( + restored + .ordered_entries() + .map(|(_, entry)| match entry { + RetainedContextEntry::VerifiedAnswer(answer) => answer.questions[0].answer.as_str(), + RetainedContextEntry::UserMessage(message) => message.text.as_str(), + }) + .collect::>(), + vec!["Yes, but never publicly.", "Do not publish after all."] + ); + restored.rollback( + &["revocation-turn"], + Some("revocation"), + RetainedInputSource::Local(None), + ); + assert_eq!( + restored, + RetainedContext { + next_order: snapshot.next_order, + ..before_restriction + } + ); +} + +#[test] +fn retained_families_enforce_storage_limits_without_changing_snapshots() { + let mut context = RetainedContext::default(); + let first = publish_answer(); + context.record(&first); + let snapshot = context.clone(); + + for index in 2..=10 { + context.record(&RetainedContextEvent::VerifiedAnswer { + answer: VerifiedAnswer { + turn_id: "turn-2".to_owned(), + call_id: format!("ask-{index}"), + questions: vec![VerifiedQuestionAnswer { + question: "Continue?".to_owned(), + answer: "Yes".to_owned(), + }], + }, + acceptance_order: None, + }); + } + assert!(!context.verified_answers_complete()); + assert_eq!(context.verified_answers().count(), MAX_FAMILY_RECORDS); + context.rollback( + &["turn-2"], + /*first_removed_message_id*/ None, + RetainedInputSource::Local(None), + ); + assert_eq!(context.verified_answers().count(), 0); + assert!(!context.verified_answers_complete()); + + let mut restored = snapshot.clone(); + let mut oversized = first.clone(); + let RetainedContextEvent::VerifiedAnswer { answer, .. } = &mut oversized; + answer.questions[0].answer = "a".repeat(MAX_RECORD_BYTES); + restored.record(&oversized); + assert!(!restored.verified_answers_complete()); + assert!( + restored + .verified_answers() + .next() + .expect("retained answer fixture") + .questions + .is_empty() + ); + assert_eq!( + snapshot + .verified_answers() + .next() + .expect("retained answer fixture") + .questions[0] + .answer, + "Yes, but never publicly." + ); + for index in 0..=MAX_FAMILY_RECORDS { + restored.record_user_message( + RetainedUserMessage { + turn_id: "later-turn".to_owned(), + message_id: Some(format!("message-{index}")), + text: "Keep the repository private.".to_owned(), + complete: true, + }, + RetainedInputSource::Local(None), + ); + } + assert!(!restored.user_messages_complete()); + assert_eq!( + restored + .ordered_entries() + .filter(|(_, entry)| matches!(entry, RetainedContextEntry::UserMessage(_))) + .count(), + MAX_FAMILY_RECORDS + ); + let recent = restored.clone(); + restored.record_user_message( + RetainedUserMessage { + turn_id: "earlier-turn".to_owned(), + message_id: Some("delayed-message".to_owned()), + text: "An older queued instruction.".to_owned(), + complete: true, + }, + RetainedInputSource::Local(Some(0)), + ); + assert_eq!( + restored, recent, + "delayed old input must not evict newer evidence" + ); + restored.record_user_message( + RetainedUserMessage { + turn_id: "oversized-turn".to_owned(), + message_id: Some("oversized-message".to_owned()), + text: "restriction ".repeat(MAX_RECORD_BYTES), + complete: true, + }, + RetainedInputSource::Local(None), + ); + let Some((_, RetainedContextEntry::UserMessage(message))) = + restored.ordered_entries().next_back() + else { + panic!("latest user evidence"); + }; + assert_eq!((&message.text, message.complete), (&String::new(), false)); +} + +#[test] +fn recovered_excerpts_obey_record_and_family_limits() { + let mut context = RetainedContext::default(); + for index in 0..MAX_FAMILY_RECORDS { + context.record_user_message( + RetainedUserMessage { + turn_id: "turn-1".to_owned(), + message_id: Some(format!("message-{index}")), + text: String::new(), + complete: false, + }, + RetainedInputSource::Local(None), + ); + } + let unchanged = context.clone(); + context.recover_user_message_excerpts(|_| Some("x".repeat(MAX_RECORD_BYTES))); + assert_eq!(context, unchanged); + context.recover_user_message_excerpts(|_| Some("x".repeat(MAX_RECORD_BYTES - 1_024))); + assert!(context.user_messages_incomplete); + assert!(context.user_messages.len() < MAX_FAMILY_RECORDS); + assert!(serde_json::to_vec(&context.user_messages).unwrap().len() <= MAX_FAMILY_BYTES); +} + +#[test] +fn legacy_checkpoints_mark_user_messages_incomplete() { + let RetainedContextEvent::VerifiedAnswer { answer, .. } = publish_answer(); + let mut wire = serde_json::json!({ + "verified_answers": [answer], "incomplete": false + }); + let legacy: RetainedContext = + serde_json::from_value(wire.clone()).expect("legacy retained-answer checkpoint"); + assert!(legacy.verified_answers_complete()); + assert!(!legacy.user_messages_complete()); + + // Retain the original wire key even though the internal field names its family. + wire["verified_answers"][0]["order"] = serde_json::json!(0); + wire["user_messages"] = serde_json::json!([]); + wire["user_messages_incomplete"] = serde_json::json!(true); + wire["next_order"] = serde_json::json!(0); + assert_eq!(serde_json::to_value(&legacy).unwrap(), wire); + + let mut restored = RetainedContext::default(); + restored.restore(/*checkpoint*/ None, &[]); + assert!(!restored.user_messages_complete()); +} + +#[test] +fn restored_input_order_accounts_for_surviving_local_sources() { + let surviving_items = [(Some(7), false), (Some(100), true), (None, false)].map( + |(user_input_order, inherited_user_message)| ResponseItemEnvelope { + item: ResponseItem::Message { + id: None, + role: "user".to_owned(), + content: vec![ContentItem::InputText { + text: "Inspect the deployment.".to_owned(), + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + }, + metadata: Some(CodexHarnessMetadata { + user_input_order, + inherited_user_message, + ..Default::default() + }), + }, + ); + let checkpoint = RetainedContext { + next_order: 20, + ..Default::default() + }; + for (checkpoint, expected_order) in [(None, 8), (Some(&checkpoint), 20)] { + let mut restored = RetainedContext::default(); + restored.restore(checkpoint, &surviving_items); + assert_eq!( + ( + restored.reserve_order(), + restored.ordered_entries().count(), + restored.user_messages_complete(), + ), + (expected_order, 0, checkpoint.is_some()), + ); + } +} + +#[test] +fn accepted_order_survives_delayed_recording_and_checkpoint_replay() { + let mut context = RetainedContext::default(); + // Rejected/canceled input consumes a sequence but never becomes evidence. + context.reserve_order(); + let steer_order = context.reserve_order(); + let answer_order = context.reserve_order(); + let RetainedContextEvent::VerifiedAnswer { answer, .. } = publish_answer(); + let event = RetainedContextEvent::VerifiedAnswer { + answer, + acceptance_order: Some(answer_order), + }; + context.record(&event); + let checkpoint = context.clone(); + let instruction = RetainedUserMessage { + turn_id: "turn-1".to_owned(), + message_id: Some("steer".to_owned()), + text: "Keep the repository private.".to_owned(), + complete: true, + }; + context.record_user_message( + instruction.clone(), + RetainedInputSource::Local(Some(steer_order)), + ); + let mut resumed = RetainedContext::default(); + resumed.restore(Some(&checkpoint), &[]); + resumed.record_user_message(instruction, RetainedInputSource::Local(Some(steer_order))); + assert_eq!(resumed, context); + assert_eq!( + resumed + .ordered_entries() + .map(|(order, entry)| match entry { + RetainedContextEntry::UserMessage(message) => (order, message.text.as_str()), + RetainedContextEntry::VerifiedAnswer(answer) => { + (order, answer.questions[0].answer.as_str()) + } + }) + .collect::>(), + vec![ + ( + RetainedContextOrder::Local(steer_order), + "Keep the repository private." + ), + ( + RetainedContextOrder::Local(answer_order), + "Yes, but never publicly." + ), + ], + ); + // This checkpoint predates model delivery of the queued instruction. Rollback + // still removes its later-accepted answer using the persisted boundary order. + let mut before_delivery = checkpoint; + before_delivery.rollback( + &["turn-1"], + Some("steer"), + RetainedInputSource::Local(Some(steer_order)), + ); + resumed.rollback( + &["turn-1"], + Some("steer"), + RetainedInputSource::Local(Some(steer_order)), + ); + assert_eq!(before_delivery, resumed); + assert_eq!( + resumed, + RetainedContext { + next_order: 3, + ..Default::default() + } + ); +} + +#[test] +fn adopted_instructions_preserve_local_order_and_rollback_scope() { + let mut context = RetainedContext::default(); + context.record(&publish_answer()); + for index in 0..2 { + let message = RetainedUserMessage { + turn_id: "parent-turn".to_owned(), + message_id: Some(format!("parent-{index}")), + text: format!("Parent instruction {index}"), + complete: true, + }; + context.record_user_message(message.clone(), RetainedInputSource::Inherited); + context.record_user_message(message, RetainedInputSource::Inherited); + } + assert_eq!(context.reserve_order(), 1); + assert!(!context.verified_answers_complete()); + assert_eq!( + context + .ordered_entries() + .map(|(order, _)| order) + .collect::>(), + vec![ + RetainedContextOrder::Inherited(0), + RetainedContextOrder::Inherited(1), + RetainedContextOrder::Local(0) + ], + ); + let checkpoint = serde_json::from_slice(&serde_json::to_vec(&context).unwrap()).unwrap(); + let mut restored = RetainedContext::default(); + restored.restore(Some(&checkpoint), &[]); + assert_eq!(restored, context); + // Older roots adopted complete instructions without recording the missing parent answers. + let mut legacy_checkpoint = serde_json::to_value(&checkpoint).unwrap(); + legacy_checkpoint["incomplete"] = serde_json::json!(false); + let legacy_checkpoint = serde_json::from_value(legacy_checkpoint).unwrap(); + restored.restore(Some(&legacy_checkpoint), &[]); + assert_eq!(restored, context); + restored.rollback( + &["turn-1"], + /*first_removed_message_id*/ None, + RetainedInputSource::Local(Some(0)), + ); + let mut expected = context; + expected.verified_answers.clear(); + assert_eq!(restored, expected); + restored.rollback( + &["parent-turn"], + Some("parent-1"), + RetainedInputSource::Inherited, + ); + expected.user_messages.pop_back(); + assert_eq!(restored, expected); +} diff --git a/codex-rs/history/src/rollout_payload.rs b/codex-rs/history/src/rollout_payload.rs new file mode 100644 index 0000000000000000000000000000000000000000..0a491fed6dd8736de776637fef7897e1fa2edbdb --- /dev/null +++ b/codex-rs/history/src/rollout_payload.rs @@ -0,0 +1,294 @@ +use std::borrow::Cow; + +use super::CodexHarnessMetadata; +use super::CompactedItem; +use super::EventMsg; +use super::InterAgentCommunication; +use super::McpResourceOriginCheckpoint; +use super::RealtimeItem; +use super::ResponseItem; +use super::ResponseItemEnvelope; +use super::RolloutItem; +use super::SecurityRiskScore; +use super::SessionMetaLine; +use super::TokenUsageRecord; +use super::TurnContextItem; +use super::WorldStateItem; +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; + +/// Persisted rollout item used by core history and rollout storage. +#[derive(Serialize, Deserialize, JsonSchema)] +#[serde(tag = "type", rename_all = "snake_case")] +pub(super) enum RolloutItemWire<'a> { + SessionMeta { + payload: Cow<'a, SessionMetaLine>, + }, + ResponseItem { + payload: Cow<'a, ResponseItem>, + #[serde(default, skip_serializing_if = "Option::is_none")] + metadata: Option>, + }, + InterAgentCommunication { + payload: Cow<'a, InterAgentCommunication>, + }, + InterAgentCommunicationMetadata { + payload: InterAgentCommunicationMetadataPayload, + }, + Compacted { + payload: Cow<'a, CompactedItem>, + }, + TurnContext { + payload: Cow<'a, TurnContextItem>, + }, + TokenUsageRecord { + payload: Cow<'a, TokenUsageRecord>, + }, + WorldState { + payload: Cow<'a, WorldStateItem>, + }, + RetainedContext { + payload: Cow<'a, crate::RetainedContextEvent>, + }, + SecurityRiskScore { + payload: Cow<'a, SecurityRiskScore>, + }, + EventMsg { + payload: Cow<'a, EventMsg>, + }, + RealtimeItem { + payload: Cow<'a, RealtimeItem>, + }, +} + +impl<'a> From<&'a RolloutItem> for RolloutItemWire<'a> { + fn from(item: &'a RolloutItem) -> Self { + match item { + RolloutItem::SessionMeta(payload) => Self::SessionMeta { + payload: Cow::Borrowed(payload), + }, + RolloutItem::ResponseItem(envelope) => Self::ResponseItem { + payload: Cow::Borrowed(&envelope.item), + metadata: envelope.metadata.as_ref().map(Cow::Borrowed), + }, + RolloutItem::InterAgentCommunication(payload) => Self::InterAgentCommunication { + payload: Cow::Borrowed(payload), + }, + RolloutItem::InterAgentCommunicationMetadata { trigger_turn } => { + Self::InterAgentCommunicationMetadata { + payload: InterAgentCommunicationMetadataPayload { + trigger_turn: *trigger_turn, + }, + } + } + RolloutItem::Compacted(payload) => Self::Compacted { + payload: Cow::Borrowed(payload), + }, + RolloutItem::TurnContext(payload) => Self::TurnContext { + payload: Cow::Borrowed(payload), + }, + RolloutItem::TokenUsageRecord(payload) => Self::TokenUsageRecord { + payload: Cow::Borrowed(payload), + }, + RolloutItem::WorldState(payload) => Self::WorldState { + payload: Cow::Borrowed(payload), + }, + RolloutItem::RetainedContext(payload) => Self::RetainedContext { + payload: Cow::Borrowed(payload), + }, + RolloutItem::SecurityRiskScore(payload) => Self::SecurityRiskScore { + payload: Cow::Borrowed(payload), + }, + RolloutItem::EventMsg(payload) => Self::EventMsg { + payload: Cow::Borrowed(payload), + }, + RolloutItem::RealtimeItem(payload) => Self::RealtimeItem { + payload: Cow::Borrowed(payload), + }, + } + } +} + +impl From> for RolloutItem { + fn from(item: RolloutItemWire<'_>) -> Self { + match item { + RolloutItemWire::SessionMeta { payload } => Self::SessionMeta(payload.into_owned()), + RolloutItemWire::ResponseItem { payload, metadata } => { + Self::ResponseItem(ResponseItemEnvelope { + item: payload.into_owned(), + metadata: metadata.map(Cow::into_owned), + }) + } + RolloutItemWire::InterAgentCommunication { payload } => { + Self::InterAgentCommunication(payload.into_owned()) + } + RolloutItemWire::InterAgentCommunicationMetadata { payload } => { + Self::InterAgentCommunicationMetadata { + trigger_turn: payload.trigger_turn, + } + } + RolloutItemWire::Compacted { payload } => Self::Compacted(payload.into_owned()), + RolloutItemWire::TurnContext { payload } => Self::TurnContext(payload.into_owned()), + RolloutItemWire::TokenUsageRecord { payload } => { + Self::TokenUsageRecord(payload.into_owned()) + } + RolloutItemWire::WorldState { payload } => Self::WorldState(payload.into_owned()), + RolloutItemWire::RetainedContext { payload } => { + Self::RetainedContext(payload.into_owned()) + } + RolloutItemWire::SecurityRiskScore { payload } => { + Self::SecurityRiskScore(payload.into_owned()) + } + RolloutItemWire::EventMsg { payload } => Self::EventMsg(payload.into_owned()), + RolloutItemWire::RealtimeItem { payload } => Self::RealtimeItem(payload.into_owned()), + } + } +} + +#[derive(Serialize, Deserialize, JsonSchema)] +pub(super) struct InterAgentCommunicationMetadataPayload { + trigger_turn: bool, +} + +#[derive(Serialize, Deserialize, JsonSchema)] +pub(super) struct CompactedItemWire<'a> { + message: Cow<'a, str>, + #[serde(default, skip_serializing_if = "Option::is_none")] + replacement_history: Option>>, + #[serde(default, skip_serializing_if = "Option::is_none")] + replacement_history_metadata: Option>>, + #[serde(default, skip_serializing_if = "Option::is_none")] + guardian_history: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + retained_context: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + mcp_resource_origins: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + window_number: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + first_window_id: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + previous_window_id: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[schemars(with = "Option")] + window_id: Option>, + #[serde(default)] + compaction_response_id: Option>, + #[serde(default)] + latest_token_usage_record: Option>, +} + +impl<'a> From<&'a CompactedItem> for CompactedItemWire<'a> { + fn from(item: &'a CompactedItem) -> Self { + let replacement_history_metadata = item + .replacement_history + .as_ref() + .filter(|items| items.iter().any(|envelope| envelope.metadata.is_some())) + .map(|items| { + items + .iter() + .map(|envelope| match envelope.metadata.as_ref() { + Some(metadata) => Cow::Borrowed(metadata), + None => Cow::Owned(CodexHarnessMetadata::default()), + }) + .collect() + }); + + Self { + message: Cow::Borrowed(&item.message), + replacement_history: item.replacement_history.as_ref().map(|items| { + items + .iter() + .map(|envelope| Cow::Borrowed(&envelope.item)) + .collect() + }), + replacement_history_metadata, + guardian_history: item.guardian_history.as_ref().map(Cow::Borrowed), + retained_context: item.retained_context.as_ref().map(Cow::Borrowed), + mcp_resource_origins: item.mcp_resource_origins.as_ref().map(Cow::Borrowed), + window_number: item.window_number, + first_window_id: item.first_window_id.as_deref().map(Cow::Borrowed), + previous_window_id: item.previous_window_id.as_deref().map(Cow::Borrowed), + window_id: item + .window_id + .as_deref() + .map(|window_id| WindowIdWire::Id(Cow::Borrowed(window_id))), + compaction_response_id: item.compaction_response_id.as_deref().map(Cow::Borrowed), + latest_token_usage_record: item.latest_token_usage_record.as_ref().map(Cow::Borrowed), + } + } +} + +impl TryFrom> for CompactedItem { + type Error = String; + + fn try_from(item: CompactedItemWire<'_>) -> Result { + let replacement_history = match ( + item.replacement_history, + item.replacement_history_metadata, + ) { + (Some(items), Some(metadata)) => { + let history_len = items.len(); + let metadata_len = metadata.len(); + if history_len != metadata_len { + return Err(format!( + "replacement_history_metadata has {metadata_len} entries for {history_len} history items" + )); + } + Some( + items + .into_iter() + .zip(metadata) + .map(|(item, metadata)| ResponseItemEnvelope { + item: item.into_owned(), + metadata: Some(metadata.into_owned()), + }) + .collect(), + ) + } + (Some(items), None) => Some( + items + .into_iter() + .map(|item| ResponseItemEnvelope::new(item.into_owned())) + .collect(), + ), + (None, Some(_)) => { + return Err("replacement_history_metadata requires replacement_history".to_string()); + } + (None, None) => None, + }; + + let mut window_number = item.window_number; + let window_id = match item.window_id { + Some(WindowIdWire::Id(window_id)) => Some(window_id.into_owned()), + Some(WindowIdWire::LegacyWindowNumber(legacy_window_number)) => { + window_number.get_or_insert(legacy_window_number); + None + } + None => None, + }; + + Ok(Self { + message: item.message.into_owned(), + replacement_history, + guardian_history: item.guardian_history.map(Cow::into_owned), + retained_context: item.retained_context.map(Cow::into_owned), + mcp_resource_origins: item.mcp_resource_origins.map(Cow::into_owned), + window_number, + first_window_id: item.first_window_id.map(Cow::into_owned), + previous_window_id: item.previous_window_id.map(Cow::into_owned), + window_id, + compaction_response_id: item.compaction_response_id.map(Cow::into_owned), + latest_token_usage_record: item.latest_token_usage_record.map(Cow::into_owned), + }) + } +} + +// Older rollouts stored the numeric window number in `window_id`. +#[derive(Serialize, Deserialize)] +#[serde(untagged)] +enum WindowIdWire<'a> { + Id(Cow<'a, str>), + LegacyWindowNumber(u64), +} diff --git a/codex-rs/history/src/tests.rs b/codex-rs/history/src/tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..b04d2448db40b1b67827f2bf5196e5838d5903e3 --- /dev/null +++ b/codex-rs/history/src/tests.rs @@ -0,0 +1,829 @@ +use anyhow::Result; +use codex_protocol::models::ConfigurationReasoning; +use codex_protocol::openai_models::ReasoningEffort; +use codex_protocol::protocol::ThreadSettingsSnapshot; +use pretty_assertions::assert_eq; +use serde_json::json; + +use super::*; + +fn thread_settings_snapshot(disabled_plugin_ids: Vec) -> Result { + Ok(serde_json::from_value(json!({ + "model": "gpt-5", + "model_provider_id": "openai", + "approval_policy": "never", + "approvals_reviewer": "user", + "permission_profile": codex_protocol::models::PermissionProfile::read_only(), + "cwd": std::env::current_dir()?, + "collaboration_mode": { + "mode": "default", + "settings": { + "model": "gpt-5", + "reasoning_effort": null, + "developer_instructions": null + } + }, + "disabled_plugin_ids": disabled_plugin_ids + }))?) +} + +#[test] +fn latest_disabled_plugin_ids_preserves_owned_updates_and_clears() -> Result<()> { + let thread_id = ThreadId::new(); + let ancestor_id = ThreadId::new(); + let selected = thread_settings_snapshot(vec!["example@marketplace".to_string()])?; + let cleared = thread_settings_snapshot(Vec::new())?; + let item = |thread_id, thread_settings| { + RolloutItem::EventMsg(EventMsg::ThreadSettingsApplied( + codex_protocol::protocol::ThreadSettingsAppliedEvent { + thread_id, + thread_settings, + }, + )) + }; + let foreign_items = vec![ + item(Some(ancestor_id), selected.clone()), + item(/*thread_id*/ None, selected.clone()), + ]; + assert_eq!(latest_disabled_plugin_ids(&foreign_items, thread_id), None); + + let mut history = vec![item(Some(thread_id), selected.clone())]; + history.extend(foreign_items.clone()); + assert_eq!( + latest_disabled_plugin_ids(&history, thread_id), + Some(selected.disabled_plugin_ids.as_slice()) + ); + + history.push(item(Some(thread_id), cleared.clone())); + history.extend(foreign_items); + history.push(turn_context_with_disabled_plugins(Some( + selected.disabled_plugin_ids, + ))?); + assert_eq!( + latest_disabled_plugin_ids(&history, thread_id), + Some(cleared.disabled_plugin_ids.as_slice()) + ); + Ok(()) +} + +fn turn_context_with_disabled_plugins( + disabled_plugin_ids: Option>, +) -> Result { + let mut context: TurnContextItem = serde_json::from_value(json!({ + "cwd": std::env::current_dir()?, + "approval_policy": "never", + "sandbox_policy": { "type": "danger-full-access" }, + "model": "gpt-5", + "summary": "auto" + }))?; + context.disabled_plugin_ids = disabled_plugin_ids; + Ok(RolloutItem::TurnContext(context)) +} + +#[test] +fn latest_disabled_plugin_ids_falls_back_only_to_latest_turn_context() -> Result<()> { + let thread_id = ThreadId::new(); + let selected = vec!["example@marketplace".to_string()]; + let mut history = vec![turn_context_with_disabled_plugins(Some(selected.clone()))?]; + assert_eq!( + latest_disabled_plugin_ids(&history, thread_id), + Some(selected.as_slice()) + ); + + history.push(RolloutItem::EventMsg(EventMsg::ThreadSettingsApplied( + codex_protocol::protocol::ThreadSettingsAppliedEvent { + thread_id: Some(ThreadId::new()), + thread_settings: thread_settings_snapshot(Vec::new())?, + }, + ))); + assert_eq!( + latest_disabled_plugin_ids(&history, thread_id), + Some(selected.as_slice()) + ); + + let cleared = Vec::new(); + history.push(turn_context_with_disabled_plugins(Some(cleared.clone()))?); + assert_eq!( + latest_disabled_plugin_ids(&history, thread_id), + Some(cleared.as_slice()) + ); + + history.push(turn_context_with_disabled_plugins( + /*disabled_plugin_ids*/ None, + )?); + assert_eq!(latest_disabled_plugin_ids(&history, thread_id), None); + Ok(()) +} + +#[test] +fn older_thread_settings_snapshot_defaults_disabled_plugins_to_empty() -> Result<()> { + let expected = thread_settings_snapshot(Vec::new())?; + let mut legacy = serde_json::to_value(&expected)?; + legacy + .as_object_mut() + .unwrap() + .remove("disabled_plugin_ids"); + + assert_eq!( + serde_json::from_value::(legacy)?, + expected + ); + Ok(()) +} + +#[test] +fn response_item_envelope_accessors_preserve_item() { + let expected_item = ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: "hello".to_string(), + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + }; + let mut envelope = ResponseItemEnvelope::new(expected_item.clone()); + + assert_eq!(&*envelope, &expected_item); + let borrowed: &ResponseItem = envelope.borrow(); + assert_eq!(borrowed, &expected_item); + let replacement_item = ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: "goodbye".to_string(), + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + }; + *envelope = replacement_item.clone(); + + assert_eq!(envelope.into_item(), replacement_item); +} + +#[test] +/// Keeps legacy response-item rollout lines readable and byte-shape compatible. +fn response_item_rollout_line_preserves_shape() -> Result<()> { + let legacy_line = json!({ + "timestamp": "2025-01-03T12:00:00.000Z", + "ordinal": 7, + "type": "response_item", + "payload": { + "type": "message", + "role": "user", + "content": [{ + "type": "input_text", + "text": "hello", + }], + }, + }); + + let line = RolloutLine { + timestamp: "2025-01-03T12:00:00.000Z".to_string(), + ordinal: Some(7), + item: serde_json::from_value(legacy_line.clone())?, + }; + let RolloutItem::ResponseItem(envelope) = &line.item else { + panic!("expected response item"); + }; + assert!(matches!(&envelope.item, ResponseItem::Message { .. })); + assert_eq!(envelope.metadata, None); + assert_eq!(serde_json::to_value(line)?, legacy_line); + Ok(()) +} + +#[test] +/// Keeps harness metadata beside, rather than inside, response-item payloads. +fn response_item_envelope_stores_metadata_beside_rollout_payload() -> Result<()> { + let response_item = response_message("developer"); + let line = RolloutLine { + timestamp: "2025-01-03T12:00:00.000Z".to_string(), + ordinal: Some(7), + item: RolloutItem::ResponseItem(ResponseItemEnvelope { + item: response_item.clone(), + metadata: Some(CodexHarnessMetadata { + client_authored: true, + history_truncation_token_limit: Some(20_000), + inherited_user_message: true, + ..Default::default() + }), + }), + }; + let serialized = serde_json::to_value(&line)?; + + assert_eq!( + serialized, + json!({ + "timestamp": "2025-01-03T12:00:00.000Z", + "ordinal": 7, + "type": "response_item", + "payload": response_item, + "metadata": { + "client_authored": true, + "fallback_token_limit_override": 20_000, + "inherited_user_message": true, + }, + }) + ); + assert_eq!(serialized["payload"].get("metadata"), None); + + let restored = serde_json::from_value(serialized)?; + let RolloutItem::ResponseItem(envelope) = restored else { + panic!("expected response item"); + }; + assert_eq!( + envelope.metadata, + Some(CodexHarnessMetadata { + client_authored: true, + history_truncation_token_limit: Some(20_000), + inherited_user_message: true, + ..Default::default() + }) + ); + Ok(()) +} + +#[test] +fn response_item_envelope_preserves_harness_authored_configuration_provenance() -> Result<()> { + let response_item = ResponseItem::ConfigurationUpdate { + reasoning: ConfigurationReasoning { + effort: ReasoningEffort::High, + }, + }; + let metadata = CodexHarnessMetadata { + harness_authored_configuration: true, + ..Default::default() + }; + let rollout_item = RolloutItem::ResponseItem(ResponseItemEnvelope { + item: response_item.clone(), + metadata: Some(metadata.clone()), + }); + + let serialized = serde_json::to_value(&rollout_item)?; + assert_eq!( + serialized, + json!({ + "type": "response_item", + "payload": { + "type": "configuration_update", + "reasoning": { "effort": "high" }, + }, + "metadata": { + "client_authored": false, + "harness_authored_configuration": true, + }, + }) + ); + + let restored = serde_json::from_value::(serialized)?; + let RolloutItem::ResponseItem(envelope) = restored else { + panic!("expected response item"); + }; + assert_eq!( + envelope, + ResponseItemEnvelope { + item: response_item, + metadata: Some(metadata), + } + ); + Ok(()) +} + +#[test] +/// Keeps future metadata fields from making older binaries reject persisted items. +fn response_item_envelope_ignores_unknown_harness_metadata_fields() -> Result<()> { + let line = serde_json::from_value(json!({ + "timestamp": "2025-01-03T12:00:00.000Z", + "ordinal": 7, + "type": "response_item", + "payload": { + "type": "message", + "role": "user", + "content": [{ + "type": "input_text", + "text": "hello", + }], + }, + "metadata": { + "future_field": "value", + }, + }))?; + + let RolloutItem::ResponseItem(envelope) = line else { + panic!("expected response item"); + }; + assert_eq!(envelope.metadata, Some(CodexHarnessMetadata::default())); + + let compacted = serde_json::from_value::(json!({ + "message": "summary", + "replacement_history": [response_message("user")], + "replacement_history_metadata": [{ "future_field": "value" }], + }))?; + assert_eq!( + compacted.replacement_history.expect("replacement history")[0].metadata, + Some(CodexHarnessMetadata::default()) + ); + Ok(()) +} + +#[test] +/// Keeps legacy compacted replacement histories readable while adding nullable checkpoints. +fn response_item_replacement_history_accepts_legacy_shape() -> Result<()> { + let legacy_item = json!({ + "message": "summary", + "replacement_history": [{ + "type": "message", + "role": "user", + "content": [{ + "type": "input_text", + "text": "hello", + }], + }], + }); + + let item = serde_json::from_value::(legacy_item.clone())?; + let replacement_history = item + .replacement_history + .as_ref() + .expect("replacement history"); + assert!(matches!( + &replacement_history[0].item, + ResponseItem::Message { .. } + )); + assert_eq!(replacement_history[0].metadata, None); + let mut expected_item = legacy_item; + expected_item["compaction_response_id"] = json!(null); + expected_item["latest_token_usage_record"] = json!(null); + assert_eq!(serde_json::to_value(item)?, expected_item); + Ok(()) +} + +#[test] +/// Stores complete aligned checkpoint metadata without modifying response items. +fn compacted_replacement_history_stores_metadata_in_an_aligned_sidecar() -> Result<()> { + let developer_message = response_message("developer"); + let compaction_item = ResponseItem::Compaction { + id: None, + encrypted_content: "opaque".to_string(), + internal_chat_message_metadata_passthrough: None, + }; + let item = CompactedItem { + message: "summary".to_string(), + replacement_history: Some(vec![ + ResponseItemEnvelope { + item: developer_message.clone(), + metadata: Some(CodexHarnessMetadata { + client_authored: true, + ..Default::default() + }), + }, + ResponseItemEnvelope::new(compaction_item.clone()), + ]), + retained_context: None, + guardian_history: None, + mcp_resource_origins: None, + window_number: None, + first_window_id: None, + previous_window_id: None, + window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, + }; + + let serialized = serde_json::to_value(item)?; + assert_eq!( + serialized, + json!({ + "message": "summary", + "replacement_history": [developer_message, compaction_item], + "replacement_history_metadata": [ + { "client_authored": true }, + { "client_authored": false }, + ], + "compaction_response_id": null, + "latest_token_usage_record": null, + }) + ); + + let restored = serde_json::from_value::(serialized)?; + assert_eq!( + restored.replacement_history, + Some(vec![ + ResponseItemEnvelope { + item: developer_message, + metadata: Some(CodexHarnessMetadata { + client_authored: true, + ..Default::default() + }), + }, + ResponseItemEnvelope { + item: compaction_item, + metadata: Some(CodexHarnessMetadata::default()), + }, + ]) + ); + Ok(()) +} + +#[test] +/// Rejects checkpoint sidecars that cannot be paired unambiguously with history. +fn compacted_replacement_history_rejects_misaligned_metadata() { + let malformed_items = [ + json!({ + "message": "summary", + "replacement_history": [response_message("user")], + "replacement_history_metadata": [], + }), + json!({ + "message": "summary", + "replacement_history": [response_message("user")], + "replacement_history_metadata": [{}, {}], + }), + json!({ + "message": "summary", + "replacement_history_metadata": [{}], + }), + ]; + + for malformed in malformed_items { + let error = serde_json::from_value::(malformed) + .expect_err("misaligned checkpoint metadata must be rejected"); + assert!( + error.to_string().contains("replacement_history_metadata"), + "error: {error}" + ); + } +} + +#[test] +/// Keeps annotated checkpoints readable by binaries expecting raw response items. +fn compacted_metadata_remains_compatible_with_legacy_response_item_readers() -> Result<()> { + #[derive(Deserialize)] + #[serde(tag = "type", content = "payload", rename_all = "snake_case")] + enum LegacyRolloutItem { + ResponseItem(Box), + Compacted(LegacyCompactedItem), + } + + #[derive(Deserialize)] + struct LegacyCompactedItem { + replacement_history: Vec, + } + + let response_item = response_message("developer"); + let envelope = ResponseItemEnvelope { + item: response_item.clone(), + metadata: Some(CodexHarnessMetadata { + client_authored: true, + ..Default::default() + }), + }; + let response_line = serde_json::to_value(RolloutItem::ResponseItem(envelope.clone()))?; + let LegacyRolloutItem::ResponseItem(legacy_response) = + serde_json::from_value::(response_line)? + else { + panic!("expected legacy response item"); + }; + assert_eq!(*legacy_response, response_item); + + let checkpoint = crate::GuardianHistoryCheckpoint(vec![response_item.clone()]); + let compacted_line = serde_json::to_value(RolloutItem::Compacted(CompactedItem { + message: "summary".to_string(), + replacement_history: Some(vec![envelope]), + retained_context: None, + guardian_history: Some(checkpoint.clone()), + mcp_resource_origins: Some(McpResourceOriginCheckpoint::default()), + window_number: None, + first_window_id: None, + previous_window_id: None, + window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, + }))?; + + let restored: RolloutItem = serde_json::from_value(compacted_line.clone())?; + let RolloutItem::Compacted(restored) = restored else { + panic!("expected compacted item"); + }; + assert_eq!(restored.guardian_history, Some(checkpoint)); + let LegacyRolloutItem::Compacted(legacy) = + serde_json::from_value::(compacted_line)? + else { + panic!("expected legacy compacted item"); + }; + assert_eq!(legacy.replacement_history, vec![response_item]); + Ok(()) +} + +#[test] +/// Preserves the established tagged payload representation for every rollout variant. +fn rollout_item_variants_preserve_existing_payload_shapes() -> Result<()> { + let cwd = std::env::current_dir()?; + let fixtures = [ + json!({ + "type": "session_meta", + "payload": SessionMetaLine { + meta: SessionMeta::default(), + git: None, + }, + }), + json!({ + "type": "response_item", + "payload": response_message("user"), + }), + json!({ + "type": "inter_agent_communication", + "payload": { + "author": "/root", + "recipient": "/root/child", + "other_recipients": [], + "content": "hello", + "trigger_turn": false, + }, + }), + json!({ + "type": "inter_agent_communication_metadata", + "payload": { "trigger_turn": true }, + }), + json!({ + "type": "compacted", + "payload": { + "message": "summary", + "compaction_response_id": null, + "latest_token_usage_record": null, + }, + }), + json!({ + "type": "turn_context", + "payload": { + "cwd": cwd, + "approval_policy": "never", + "sandbox_policy": { "type": "danger-full-access" }, + "model": "gpt-5", + "summary": "auto", + }, + }), + json!({ + "type": "token_usage_record", + "payload": { + "thread_id": "0195cda5-433d-7f9a-9d7b-a9f15b60c2e2", + "turn_id": "turn-1", + "session_id": "0195cda5-433d-7f9a-9d7b-a9f15b60c2e2", + "root_turn_id": "turn-1", + "response_id": "response-1", + "usage": { + "input_tokens": 10, + "cached_input_tokens": 2, + "cache_write_input_tokens": 0, + "output_tokens": 3, + "reasoning_output_tokens": 1, + "total_tokens": 13, + }, + "turn_token_usage": { + "input_tokens": 10, + "cached_input_tokens": 2, + "cache_write_input_tokens": 0, + "output_tokens": 3, + "reasoning_output_tokens": 1, + "total_tokens": 13, + }, + "thread_token_usage": { + "input_tokens": 10, + "cached_input_tokens": 2, + "cache_write_input_tokens": 0, + "output_tokens": 3, + "reasoning_output_tokens": 1, + "total_tokens": 13, + }, + }, + }), + json!({ + "type": "world_state", + "payload": { "full": true, "state": { "cwd": "/tmp" } }, + }), + json!({ + "type": "security_risk_score", + "payload": { + "scores": { + "action_risk": 0.92, + "data_exfiltration": 0.31, + }, + }, + }), + json!({ + "type": "event_msg", + "payload": { "type": "warning", "message": "heads up" }, + }), + json!({ + "type": "retained_context", + "payload": { + "type": "verified_answer", + "turn_id": "turn-1", + "call_id": "ask-1", + "questions": [{"question": "Publish?", "answer": "Only privately."}], + }, + }), + json!({ + "type": "realtime_item", + "payload": { + "id": "segment-1", + "realtime_session_id": "session-1", + "type": "transcript_segment", + "role": "assistant", + "text": "hello", + }, + }), + ]; + + for expected in fixtures { + let item = serde_json::from_value::(expected.clone())?; + assert_eq!(serde_json::to_value(item)?, expected); + } + Ok(()) +} + +#[test] +/// Keeps the generated schema aligned with each variant's actual persisted shape. +fn rollout_item_schema_matches_tagged_payload_and_sibling_metadata() -> Result<()> { + let schema = serde_json::to_value(schemars::schema_for!(RolloutItem))?; + let variants = schema["oneOf"].as_array().expect("rollout variants"); + assert_eq!(variants.len(), 12); + + for variant in variants { + let required = variant["required"].as_array().expect("required fields"); + assert!(required.contains(&json!("type")), "schema: {variant}"); + assert!(required.contains(&json!("payload")), "schema: {variant}"); + } + + let response_item = variants + .iter() + .find(|variant| variant["properties"]["type"]["enum"] == json!(["response_item"])) + .expect("response item schema"); + assert!(response_item["properties"].get("metadata").is_some()); + assert_eq!( + response_item["properties"]["payload"]["$ref"], + json!("#/definitions/ResponseItem") + ); + + let compacted = &schema["definitions"]["CompactedItem"]; + assert_eq!( + compacted["properties"]["replacement_history"]["items"]["$ref"], + json!("#/definitions/ResponseItem") + ); + assert_eq!( + compacted["properties"]["replacement_history_metadata"]["items"]["$ref"], + json!("#/definitions/CodexHarnessMetadata") + ); + let required = compacted["required"].as_array().expect("required fields"); + assert!(!required.contains(&json!("replacement_history"))); + assert!(!required.contains(&json!("replacement_history_metadata"))); + Ok(()) +} + +fn response_message(role: &str) -> ResponseItem { + ResponseItem::Message { + id: None, + role: role.to_string(), + content: vec![ContentItem::InputText { + text: "hello".to_string(), + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + } +} + +#[test] +/// Preserves the stored compacted-item window metadata shape. +fn compacted_item_serializes_window_number_and_id() -> Result<()> { + let item = CompactedItem { + message: "summary".to_string(), + replacement_history: None, + retained_context: None, + guardian_history: None, + mcp_resource_origins: None, + window_number: Some(3), + first_window_id: Some("019b3f6e-0000-7000-8000-000000000001".to_string()), + previous_window_id: Some("019b3f6e-0000-7000-8000-000000000002".to_string()), + window_id: Some("019b3f6e-7a10-7cc3-8b6e-1d09e2f7a001".to_string()), + compaction_response_id: None, + latest_token_usage_record: None, + }; + + assert_eq!( + serde_json::to_value(item)?, + json!({ + "message": "summary", + "window_number": 3, + "first_window_id": "019b3f6e-0000-7000-8000-000000000001", + "previous_window_id": "019b3f6e-0000-7000-8000-000000000002", + "window_id": "019b3f6e-7a10-7cc3-8b6e-1d09e2f7a001", + "compaction_response_id": null, + "latest_token_usage_record": null, + }) + ); + Ok(()) +} + +#[test] +/// Keeps legacy numeric window IDs readable in stored compacted items. +fn compacted_item_migrates_legacy_numeric_window_id() -> Result<()> { + let item = serde_json::from_value::(json!({ + "message": "summary", + "window_id": 3, + }))?; + + assert_eq!( + item, + CompactedItem { + message: "summary".to_string(), + replacement_history: None, + retained_context: None, + guardian_history: None, + mcp_resource_origins: None, + window_number: Some(3), + first_window_id: None, + previous_window_id: None, + window_id: None, + compaction_response_id: None, + latest_token_usage_record: None, + } + ); + Ok(()) +} + +#[test] +fn copied_history_uses_persisted_history_mode() -> Result<()> { + let thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000001")?; + let session_meta = RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + session_id: thread_id.into(), + id: thread_id, + history_mode: ThreadHistoryMode::Legacy, + ..SessionMeta::default() + }, + git: None, + }); + let history = InitialHistory::Resumed(ResumedHistory { + conversation_id: thread_id, + history: Arc::new(vec![session_meta.clone()]), + rollout_path: None, + }); + + assert_eq!( + history.get_history_mode(ThreadHistoryMode::Paginated), + ThreadHistoryMode::Legacy + ); + assert_eq!( + InitialHistory::Forked(vec![session_meta]).get_history_mode(ThreadHistoryMode::Paginated), + ThreadHistoryMode::Legacy + ); + assert_eq!( + InitialHistory::New.get_history_mode(ThreadHistoryMode::Paginated), + ThreadHistoryMode::Paginated + ); + assert_eq!( + InitialHistory::Resumed(ResumedHistory { + conversation_id: thread_id, + history: Arc::new(Vec::new()), + rollout_path: None, + }) + .get_history_mode(ThreadHistoryMode::Paginated), + ThreadHistoryMode::Paginated + ); + Ok(()) +} + +#[test] +fn multi_agent_version_uses_newest_present_session_meta_value() -> Result<()> { + let thread_id = ThreadId::from_string("67e55044-10b1-426f-9247-bb680e5fe0c8")?; + let older_meta = SessionMetaLine { + meta: SessionMeta { + session_id: thread_id.into(), + id: thread_id, + multi_agent_version: Some(MultiAgentVersion::V2), + ..Default::default() + }, + git: None, + }; + let newer_meta_without_version = SessionMetaLine { + meta: SessionMeta { + session_id: thread_id.into(), + id: thread_id, + multi_agent_version: None, + ..Default::default() + }, + git: None, + }; + + assert_eq!( + multi_agent_version_from_items( + &[ + RolloutItem::SessionMeta(older_meta), + RolloutItem::SessionMeta(newer_meta_without_version), + ], + Some(thread_id), + ), + Some(MultiAgentVersion::V2) + ); + Ok(()) +} diff --git a/codex-rs/hooks/schema/generated/interrupt.command.input.schema.json b/codex-rs/hooks/schema/generated/interrupt.command.input.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..e0ac2a96bb719511d157dd787796e2466524a1fd --- /dev/null +++ b/codex-rs/hooks/schema/generated/interrupt.command.input.schema.json @@ -0,0 +1,55 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "NullableString": { + "type": [ + "string", + "null" + ] + } + }, + "properties": { + "cwd": { + "type": "string" + }, + "hook_event_name": { + "const": "Interrupt", + "type": "string" + }, + "model": { + "type": "string" + }, + "permission_mode": { + "enum": [ + "default", + "acceptEdits", + "plan", + "dontAsk", + "bypassPermissions" + ], + "type": "string" + }, + "session_id": { + "type": "string" + }, + "transcript_path": { + "$ref": "#/definitions/NullableString" + }, + "turn_id": { + "description": "Codex extension: expose the active turn id to internal turn-scoped hooks.", + "type": "string" + } + }, + "required": [ + "cwd", + "hook_event_name", + "model", + "permission_mode", + "session_id", + "transcript_path", + "turn_id" + ], + "title": "interrupt.command.input", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/interrupt.command.output.schema.json b/codex-rs/hooks/schema/generated/interrupt.command.output.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..4d039577456e36d226675a58e1f4ac2182b7ced0 --- /dev/null +++ b/codex-rs/hooks/schema/generated/interrupt.command.output.schema.json @@ -0,0 +1,12 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "properties": { + "systemMessage": { + "default": null, + "type": "string" + } + }, + "title": "interrupt.command.output", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/permission-request.command.input.schema.json b/codex-rs/hooks/schema/generated/permission-request.command.input.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..9ee8996db169d55b0da44d7efbb57af76676c2fe --- /dev/null +++ b/codex-rs/hooks/schema/generated/permission-request.command.input.schema.json @@ -0,0 +1,67 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "NullableString": { + "type": [ + "string", + "null" + ] + } + }, + "properties": { + "agent_id": { + "type": "string" + }, + "agent_type": { + "type": "string" + }, + "cwd": { + "type": "string" + }, + "hook_event_name": { + "const": "PermissionRequest", + "type": "string" + }, + "model": { + "type": "string" + }, + "permission_mode": { + "enum": [ + "default", + "acceptEdits", + "plan", + "dontAsk", + "bypassPermissions" + ], + "type": "string" + }, + "session_id": { + "type": "string" + }, + "tool_input": true, + "tool_name": { + "type": "string" + }, + "transcript_path": { + "$ref": "#/definitions/NullableString" + }, + "turn_id": { + "description": "Codex extension: expose the active turn id to internal turn-scoped hooks.", + "type": "string" + } + }, + "required": [ + "cwd", + "hook_event_name", + "model", + "permission_mode", + "session_id", + "tool_input", + "tool_name", + "transcript_path", + "turn_id" + ], + "title": "permission-request.command.input", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/permission-request.command.output.schema.json b/codex-rs/hooks/schema/generated/permission-request.command.output.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..347820f7ed6059b654dc8210c3b3c016d61d2d4a --- /dev/null +++ b/codex-rs/hooks/schema/generated/permission-request.command.output.schema.json @@ -0,0 +1,91 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "PermissionRequestBehaviorWire": { + "enum": [ + "allow", + "deny" + ], + "type": "string" + }, + "PermissionRequestDecisionWire": { + "additionalProperties": false, + "properties": { + "behavior": { + "$ref": "#/definitions/PermissionRequestBehaviorWire" + }, + "interrupt": { + "default": false, + "description": "Reserved for future short-circuiting semantics.\n\nPermissionRequest hooks currently fail closed if this field is `true`.", + "type": "boolean" + }, + "message": { + "default": null, + "type": "string" + }, + "updatedInput": { + "default": null, + "description": "Reserved for a future input-rewrite capability.\n\nPermissionRequest hooks currently fail closed if this field is present." + }, + "updatedPermissions": { + "default": null, + "description": "Reserved for a future permission-rewrite capability.\n\nPermissionRequest hooks currently fail closed if this field is present." + } + }, + "required": [ + "behavior" + ], + "type": "object" + }, + "PermissionRequestHookSpecificOutputWire": { + "additionalProperties": false, + "properties": { + "decision": { + "allOf": [ + { + "$ref": "#/definitions/PermissionRequestDecisionWire" + } + ], + "default": null + }, + "hookEventName": { + "const": "PermissionRequest", + "type": "string" + } + }, + "required": [ + "hookEventName" + ], + "type": "object" + } + }, + "properties": { + "continue": { + "default": true, + "type": "boolean" + }, + "hookSpecificOutput": { + "allOf": [ + { + "$ref": "#/definitions/PermissionRequestHookSpecificOutputWire" + } + ], + "default": null + }, + "stopReason": { + "default": null, + "type": "string" + }, + "suppressOutput": { + "default": false, + "type": "boolean" + }, + "systemMessage": { + "default": null, + "type": "string" + } + }, + "title": "permission-request.command.output", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/post-compact.command.input.schema.json b/codex-rs/hooks/schema/generated/post-compact.command.input.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..3131f3a776bea3b307eefb499f7809dc8f092efe --- /dev/null +++ b/codex-rs/hooks/schema/generated/post-compact.command.input.schema.json @@ -0,0 +1,58 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "NullableString": { + "type": [ + "string", + "null" + ] + } + }, + "properties": { + "agent_id": { + "type": "string" + }, + "agent_type": { + "type": "string" + }, + "cwd": { + "type": "string" + }, + "hook_event_name": { + "const": "PostCompact", + "type": "string" + }, + "model": { + "type": "string" + }, + "session_id": { + "type": "string" + }, + "transcript_path": { + "$ref": "#/definitions/NullableString" + }, + "trigger": { + "enum": [ + "manual", + "auto" + ], + "type": "string" + }, + "turn_id": { + "description": "Codex extension: expose the active turn id to internal turn-scoped hooks.", + "type": "string" + } + }, + "required": [ + "cwd", + "hook_event_name", + "model", + "session_id", + "transcript_path", + "trigger", + "turn_id" + ], + "title": "post-compact.command.input", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/post-compact.command.output.schema.json b/codex-rs/hooks/schema/generated/post-compact.command.output.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..0221ac29b445d8f5e493ca8420ad168fff03559c --- /dev/null +++ b/codex-rs/hooks/schema/generated/post-compact.command.output.schema.json @@ -0,0 +1,24 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "properties": { + "continue": { + "default": true, + "type": "boolean" + }, + "stopReason": { + "default": null, + "type": "string" + }, + "suppressOutput": { + "default": false, + "type": "boolean" + }, + "systemMessage": { + "default": null, + "type": "string" + } + }, + "title": "post-compact.command.output", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/post-tool-use.command.input.schema.json b/codex-rs/hooks/schema/generated/post-tool-use.command.input.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..f92af1dd3249ddf3511ef707c4310329e5b85c80 --- /dev/null +++ b/codex-rs/hooks/schema/generated/post-tool-use.command.input.schema.json @@ -0,0 +1,73 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "NullableString": { + "type": [ + "string", + "null" + ] + } + }, + "properties": { + "agent_id": { + "type": "string" + }, + "agent_type": { + "type": "string" + }, + "cwd": { + "type": "string" + }, + "hook_event_name": { + "const": "PostToolUse", + "type": "string" + }, + "model": { + "type": "string" + }, + "permission_mode": { + "enum": [ + "default", + "acceptEdits", + "plan", + "dontAsk", + "bypassPermissions" + ], + "type": "string" + }, + "session_id": { + "type": "string" + }, + "tool_input": true, + "tool_name": { + "type": "string" + }, + "tool_response": true, + "tool_use_id": { + "type": "string" + }, + "transcript_path": { + "$ref": "#/definitions/NullableString" + }, + "turn_id": { + "description": "Codex extension: expose the active turn id to internal turn-scoped hooks.", + "type": "string" + } + }, + "required": [ + "cwd", + "hook_event_name", + "model", + "permission_mode", + "session_id", + "tool_input", + "tool_name", + "tool_response", + "tool_use_id", + "transcript_path", + "turn_id" + ], + "title": "post-tool-use.command.input", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/post-tool-use.command.output.schema.json b/codex-rs/hooks/schema/generated/post-tool-use.command.output.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..011f887265f21b549d029e96932b75b23dde8107 --- /dev/null +++ b/codex-rs/hooks/schema/generated/post-tool-use.command.output.schema.json @@ -0,0 +1,72 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "BlockDecisionWire": { + "enum": [ + "block" + ], + "type": "string" + }, + "PostToolUseHookSpecificOutputWire": { + "additionalProperties": false, + "properties": { + "additionalContext": { + "default": null, + "type": "string" + }, + "hookEventName": { + "const": "PostToolUse", + "type": "string" + }, + "updatedMCPToolOutput": { + "default": null + } + }, + "required": [ + "hookEventName" + ], + "type": "object" + } + }, + "properties": { + "continue": { + "default": true, + "type": "boolean" + }, + "decision": { + "allOf": [ + { + "$ref": "#/definitions/BlockDecisionWire" + } + ], + "default": null + }, + "hookSpecificOutput": { + "allOf": [ + { + "$ref": "#/definitions/PostToolUseHookSpecificOutputWire" + } + ], + "default": null + }, + "reason": { + "default": null, + "type": "string" + }, + "stopReason": { + "default": null, + "type": "string" + }, + "suppressOutput": { + "default": false, + "type": "boolean" + }, + "systemMessage": { + "default": null, + "type": "string" + } + }, + "title": "post-tool-use.command.output", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/pre-compact.command.input.schema.json b/codex-rs/hooks/schema/generated/pre-compact.command.input.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..54e9a8b7f4f7b3ac642444264e6451c215431a67 --- /dev/null +++ b/codex-rs/hooks/schema/generated/pre-compact.command.input.schema.json @@ -0,0 +1,58 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "NullableString": { + "type": [ + "string", + "null" + ] + } + }, + "properties": { + "agent_id": { + "type": "string" + }, + "agent_type": { + "type": "string" + }, + "cwd": { + "type": "string" + }, + "hook_event_name": { + "const": "PreCompact", + "type": "string" + }, + "model": { + "type": "string" + }, + "session_id": { + "type": "string" + }, + "transcript_path": { + "$ref": "#/definitions/NullableString" + }, + "trigger": { + "enum": [ + "manual", + "auto" + ], + "type": "string" + }, + "turn_id": { + "description": "Codex extension: expose the active turn id to internal turn-scoped hooks.", + "type": "string" + } + }, + "required": [ + "cwd", + "hook_event_name", + "model", + "session_id", + "transcript_path", + "trigger", + "turn_id" + ], + "title": "pre-compact.command.input", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/pre-compact.command.output.schema.json b/codex-rs/hooks/schema/generated/pre-compact.command.output.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..644cd1a8b4a9443b25b995ad5ea4204a63959c52 --- /dev/null +++ b/codex-rs/hooks/schema/generated/pre-compact.command.output.schema.json @@ -0,0 +1,24 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "properties": { + "continue": { + "default": true, + "type": "boolean" + }, + "stopReason": { + "default": null, + "type": "string" + }, + "suppressOutput": { + "default": false, + "type": "boolean" + }, + "systemMessage": { + "default": null, + "type": "string" + } + }, + "title": "pre-compact.command.output", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/pre-tool-use.command.input.schema.json b/codex-rs/hooks/schema/generated/pre-tool-use.command.input.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..48dd4c571010a664e40cb174123c3bc746f12e34 --- /dev/null +++ b/codex-rs/hooks/schema/generated/pre-tool-use.command.input.schema.json @@ -0,0 +1,71 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "NullableString": { + "type": [ + "string", + "null" + ] + } + }, + "properties": { + "agent_id": { + "type": "string" + }, + "agent_type": { + "type": "string" + }, + "cwd": { + "type": "string" + }, + "hook_event_name": { + "const": "PreToolUse", + "type": "string" + }, + "model": { + "type": "string" + }, + "permission_mode": { + "enum": [ + "default", + "acceptEdits", + "plan", + "dontAsk", + "bypassPermissions" + ], + "type": "string" + }, + "session_id": { + "type": "string" + }, + "tool_input": true, + "tool_name": { + "type": "string" + }, + "tool_use_id": { + "type": "string" + }, + "transcript_path": { + "$ref": "#/definitions/NullableString" + }, + "turn_id": { + "description": "Codex extension: expose the active turn id to internal turn-scoped hooks.", + "type": "string" + } + }, + "required": [ + "cwd", + "hook_event_name", + "model", + "permission_mode", + "session_id", + "tool_input", + "tool_name", + "tool_use_id", + "transcript_path", + "turn_id" + ], + "title": "pre-tool-use.command.input", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/pre-tool-use.command.output.schema.json b/codex-rs/hooks/schema/generated/pre-tool-use.command.output.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..6730b27fd4fc80f8075d64346a554a1cfc94470a --- /dev/null +++ b/codex-rs/hooks/schema/generated/pre-tool-use.command.output.schema.json @@ -0,0 +1,93 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "PreToolUseDecisionWire": { + "enum": [ + "approve", + "block" + ], + "type": "string" + }, + "PreToolUseHookSpecificOutputWire": { + "additionalProperties": false, + "properties": { + "additionalContext": { + "default": null, + "type": "string" + }, + "hookEventName": { + "const": "PreToolUse", + "type": "string" + }, + "permissionDecision": { + "allOf": [ + { + "$ref": "#/definitions/PreToolUsePermissionDecisionWire" + } + ], + "default": null + }, + "permissionDecisionReason": { + "default": null, + "type": "string" + }, + "updatedInput": { + "default": null + } + }, + "required": [ + "hookEventName" + ], + "type": "object" + }, + "PreToolUsePermissionDecisionWire": { + "enum": [ + "allow", + "deny", + "ask" + ], + "type": "string" + } + }, + "properties": { + "continue": { + "default": true, + "type": "boolean" + }, + "decision": { + "allOf": [ + { + "$ref": "#/definitions/PreToolUseDecisionWire" + } + ], + "default": null + }, + "hookSpecificOutput": { + "allOf": [ + { + "$ref": "#/definitions/PreToolUseHookSpecificOutputWire" + } + ], + "default": null + }, + "reason": { + "default": null, + "type": "string" + }, + "stopReason": { + "default": null, + "type": "string" + }, + "suppressOutput": { + "default": false, + "type": "boolean" + }, + "systemMessage": { + "default": null, + "type": "string" + } + }, + "title": "pre-tool-use.command.output", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/session-end.command.input.schema.json b/codex-rs/hooks/schema/generated/session-end.command.input.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..de87f368726375e8a0732466f6ea6a431ff5b324 --- /dev/null +++ b/codex-rs/hooks/schema/generated/session-end.command.input.schema.json @@ -0,0 +1,40 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "NullableString": { + "type": [ + "string", + "null" + ] + } + }, + "properties": { + "cwd": { + "type": "string" + }, + "hook_event_name": { + "const": "SessionEnd", + "type": "string" + }, + "reason": { + "const": "other", + "type": "string" + }, + "session_id": { + "type": "string" + }, + "transcript_path": { + "$ref": "#/definitions/NullableString" + } + }, + "required": [ + "cwd", + "hook_event_name", + "reason", + "session_id", + "transcript_path" + ], + "title": "session-end.command.input", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/session-start.command.input.schema.json b/codex-rs/hooks/schema/generated/session-start.command.input.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..70eb76ec923c321b95d225812e16d4e7f52a9cc2 --- /dev/null +++ b/codex-rs/hooks/schema/generated/session-start.command.input.schema.json @@ -0,0 +1,61 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "NullableString": { + "type": [ + "string", + "null" + ] + } + }, + "properties": { + "cwd": { + "type": "string" + }, + "hook_event_name": { + "const": "SessionStart", + "type": "string" + }, + "model": { + "type": "string" + }, + "permission_mode": { + "enum": [ + "default", + "acceptEdits", + "plan", + "dontAsk", + "bypassPermissions" + ], + "type": "string" + }, + "session_id": { + "type": "string" + }, + "source": { + "enum": [ + "startup", + "resume", + "clear", + "compact", + "fork" + ], + "type": "string" + }, + "transcript_path": { + "$ref": "#/definitions/NullableString" + } + }, + "required": [ + "cwd", + "hook_event_name", + "model", + "permission_mode", + "session_id", + "source", + "transcript_path" + ], + "title": "session-start.command.input", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/session-start.command.output.schema.json b/codex-rs/hooks/schema/generated/session-start.command.output.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..1dd51e001030cb6b260114da3461cbe9d77ba944 --- /dev/null +++ b/codex-rs/hooks/schema/generated/session-start.command.output.schema.json @@ -0,0 +1,51 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "SessionStartHookSpecificOutputWire": { + "additionalProperties": false, + "properties": { + "additionalContext": { + "default": null, + "type": "string" + }, + "hookEventName": { + "const": "SessionStart", + "type": "string" + } + }, + "required": [ + "hookEventName" + ], + "type": "object" + } + }, + "properties": { + "continue": { + "default": true, + "type": "boolean" + }, + "hookSpecificOutput": { + "allOf": [ + { + "$ref": "#/definitions/SessionStartHookSpecificOutputWire" + } + ], + "default": null + }, + "stopReason": { + "default": null, + "type": "string" + }, + "suppressOutput": { + "default": false, + "type": "boolean" + }, + "systemMessage": { + "default": null, + "type": "string" + } + }, + "title": "session-start.command.output", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/stop.command.input.schema.json b/codex-rs/hooks/schema/generated/stop.command.input.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..dbd4a3f648071bee1af24758e9596e11ab78c782 --- /dev/null +++ b/codex-rs/hooks/schema/generated/stop.command.input.schema.json @@ -0,0 +1,63 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "NullableString": { + "type": [ + "string", + "null" + ] + } + }, + "properties": { + "cwd": { + "type": "string" + }, + "hook_event_name": { + "const": "Stop", + "type": "string" + }, + "last_assistant_message": { + "$ref": "#/definitions/NullableString" + }, + "model": { + "type": "string" + }, + "permission_mode": { + "enum": [ + "default", + "acceptEdits", + "plan", + "dontAsk", + "bypassPermissions" + ], + "type": "string" + }, + "session_id": { + "type": "string" + }, + "stop_hook_active": { + "type": "boolean" + }, + "transcript_path": { + "$ref": "#/definitions/NullableString" + }, + "turn_id": { + "description": "Codex extension: expose the active turn id to internal turn-scoped hooks.", + "type": "string" + } + }, + "required": [ + "cwd", + "hook_event_name", + "last_assistant_message", + "model", + "permission_mode", + "session_id", + "stop_hook_active", + "transcript_path", + "turn_id" + ], + "title": "stop.command.input", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/stop.command.output.schema.json b/codex-rs/hooks/schema/generated/stop.command.output.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..a2bac59cd12a7a02de5ce5b259e68263cba74b2f --- /dev/null +++ b/codex-rs/hooks/schema/generated/stop.command.output.schema.json @@ -0,0 +1,45 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "BlockDecisionWire": { + "enum": [ + "block" + ], + "type": "string" + } + }, + "properties": { + "continue": { + "default": true, + "type": "boolean" + }, + "decision": { + "allOf": [ + { + "$ref": "#/definitions/BlockDecisionWire" + } + ], + "default": null + }, + "reason": { + "default": null, + "description": "Claude requires `reason` when `decision` is `block`; we enforce that semantic rule during output parsing rather than in the JSON schema.", + "type": "string" + }, + "stopReason": { + "default": null, + "type": "string" + }, + "suppressOutput": { + "default": false, + "type": "boolean" + }, + "systemMessage": { + "default": null, + "type": "string" + } + }, + "title": "stop.command.output", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/subagent-start.command.input.schema.json b/codex-rs/hooks/schema/generated/subagent-start.command.input.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..2f6edd60387aee5af5573e3683702d94773e2dd1 --- /dev/null +++ b/codex-rs/hooks/schema/generated/subagent-start.command.input.schema.json @@ -0,0 +1,63 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "NullableString": { + "type": [ + "string", + "null" + ] + } + }, + "properties": { + "agent_id": { + "type": "string" + }, + "agent_type": { + "type": "string" + }, + "cwd": { + "type": "string" + }, + "hook_event_name": { + "const": "SubagentStart", + "type": "string" + }, + "model": { + "type": "string" + }, + "permission_mode": { + "enum": [ + "default", + "acceptEdits", + "plan", + "dontAsk", + "bypassPermissions" + ], + "type": "string" + }, + "session_id": { + "type": "string" + }, + "transcript_path": { + "$ref": "#/definitions/NullableString" + }, + "turn_id": { + "description": "Codex extension: expose the active turn id to internal turn-scoped hooks.", + "type": "string" + } + }, + "required": [ + "agent_id", + "agent_type", + "cwd", + "hook_event_name", + "model", + "permission_mode", + "session_id", + "transcript_path", + "turn_id" + ], + "title": "subagent-start.command.input", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/subagent-start.command.output.schema.json b/codex-rs/hooks/schema/generated/subagent-start.command.output.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..02493da98baeb2e789aa055b3b43d1cf03f3e300 --- /dev/null +++ b/codex-rs/hooks/schema/generated/subagent-start.command.output.schema.json @@ -0,0 +1,51 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "SubagentStartHookSpecificOutputWire": { + "additionalProperties": false, + "properties": { + "additionalContext": { + "default": null, + "type": "string" + }, + "hookEventName": { + "const": "SubagentStart", + "type": "string" + } + }, + "required": [ + "hookEventName" + ], + "type": "object" + } + }, + "properties": { + "continue": { + "default": true, + "type": "boolean" + }, + "hookSpecificOutput": { + "allOf": [ + { + "$ref": "#/definitions/SubagentStartHookSpecificOutputWire" + } + ], + "default": null + }, + "stopReason": { + "default": null, + "type": "string" + }, + "suppressOutput": { + "default": false, + "type": "boolean" + }, + "systemMessage": { + "default": null, + "type": "string" + } + }, + "title": "subagent-start.command.output", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/subagent-stop.command.input.schema.json b/codex-rs/hooks/schema/generated/subagent-stop.command.input.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..19646db1d2b18d2017eaa39056326bbfa66a814c --- /dev/null +++ b/codex-rs/hooks/schema/generated/subagent-stop.command.input.schema.json @@ -0,0 +1,75 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "NullableString": { + "type": [ + "string", + "null" + ] + } + }, + "properties": { + "agent_id": { + "type": "string" + }, + "agent_transcript_path": { + "$ref": "#/definitions/NullableString" + }, + "agent_type": { + "type": "string" + }, + "cwd": { + "type": "string" + }, + "hook_event_name": { + "const": "SubagentStop", + "type": "string" + }, + "last_assistant_message": { + "$ref": "#/definitions/NullableString" + }, + "model": { + "type": "string" + }, + "permission_mode": { + "enum": [ + "default", + "acceptEdits", + "plan", + "dontAsk", + "bypassPermissions" + ], + "type": "string" + }, + "session_id": { + "type": "string" + }, + "stop_hook_active": { + "type": "boolean" + }, + "transcript_path": { + "$ref": "#/definitions/NullableString" + }, + "turn_id": { + "description": "Codex extension: expose the active turn id to internal turn-scoped hooks.", + "type": "string" + } + }, + "required": [ + "agent_id", + "agent_transcript_path", + "agent_type", + "cwd", + "hook_event_name", + "last_assistant_message", + "model", + "permission_mode", + "session_id", + "stop_hook_active", + "transcript_path", + "turn_id" + ], + "title": "subagent-stop.command.input", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/subagent-stop.command.output.schema.json b/codex-rs/hooks/schema/generated/subagent-stop.command.output.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..0cb1dd2a5701ba07d9ab73464a9968279d2065d0 --- /dev/null +++ b/codex-rs/hooks/schema/generated/subagent-stop.command.output.schema.json @@ -0,0 +1,45 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "BlockDecisionWire": { + "enum": [ + "block" + ], + "type": "string" + } + }, + "properties": { + "continue": { + "default": true, + "type": "boolean" + }, + "decision": { + "allOf": [ + { + "$ref": "#/definitions/BlockDecisionWire" + } + ], + "default": null + }, + "reason": { + "default": null, + "description": "Claude requires `reason` when `decision` is `block`; we enforce that semantic rule during output parsing rather than in the JSON schema.", + "type": "string" + }, + "stopReason": { + "default": null, + "type": "string" + }, + "suppressOutput": { + "default": false, + "type": "boolean" + }, + "systemMessage": { + "default": null, + "type": "string" + } + }, + "title": "subagent-stop.command.output", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/user-prompt-submit.command.input.schema.json b/codex-rs/hooks/schema/generated/user-prompt-submit.command.input.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..6a10a9f75c19720b0863d11d7d7d80f190e8eacd --- /dev/null +++ b/codex-rs/hooks/schema/generated/user-prompt-submit.command.input.schema.json @@ -0,0 +1,65 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "NullableString": { + "type": [ + "string", + "null" + ] + } + }, + "properties": { + "agent_id": { + "type": "string" + }, + "agent_type": { + "type": "string" + }, + "cwd": { + "type": "string" + }, + "hook_event_name": { + "const": "UserPromptSubmit", + "type": "string" + }, + "model": { + "type": "string" + }, + "permission_mode": { + "enum": [ + "default", + "acceptEdits", + "plan", + "dontAsk", + "bypassPermissions" + ], + "type": "string" + }, + "prompt": { + "type": "string" + }, + "session_id": { + "type": "string" + }, + "transcript_path": { + "$ref": "#/definitions/NullableString" + }, + "turn_id": { + "description": "Codex extension: expose the active turn id to internal turn-scoped hooks.", + "type": "string" + } + }, + "required": [ + "cwd", + "hook_event_name", + "model", + "permission_mode", + "prompt", + "session_id", + "transcript_path", + "turn_id" + ], + "title": "user-prompt-submit.command.input", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/schema/generated/user-prompt-submit.command.output.schema.json b/codex-rs/hooks/schema/generated/user-prompt-submit.command.output.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..97875c108b9345bbe5518bb1968ce12c67db7484 --- /dev/null +++ b/codex-rs/hooks/schema/generated/user-prompt-submit.command.output.schema.json @@ -0,0 +1,69 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "additionalProperties": false, + "definitions": { + "BlockDecisionWire": { + "enum": [ + "block" + ], + "type": "string" + }, + "UserPromptSubmitHookSpecificOutputWire": { + "additionalProperties": false, + "properties": { + "additionalContext": { + "default": null, + "type": "string" + }, + "hookEventName": { + "const": "UserPromptSubmit", + "type": "string" + } + }, + "required": [ + "hookEventName" + ], + "type": "object" + } + }, + "properties": { + "continue": { + "default": true, + "type": "boolean" + }, + "decision": { + "allOf": [ + { + "$ref": "#/definitions/BlockDecisionWire" + } + ], + "default": null + }, + "hookSpecificOutput": { + "allOf": [ + { + "$ref": "#/definitions/UserPromptSubmitHookSpecificOutputWire" + } + ], + "default": null + }, + "reason": { + "default": null, + "type": "string" + }, + "stopReason": { + "default": null, + "type": "string" + }, + "suppressOutput": { + "default": false, + "type": "boolean" + }, + "systemMessage": { + "default": null, + "type": "string" + } + }, + "title": "user-prompt-submit.command.output", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/hooks/src/bin/write_hooks_schema_fixtures.rs b/codex-rs/hooks/src/bin/write_hooks_schema_fixtures.rs new file mode 100644 index 0000000000000000000000000000000000000000..303b561bf342e45182343f0e48cfc77540fd5b4f --- /dev/null +++ b/codex-rs/hooks/src/bin/write_hooks_schema_fixtures.rs @@ -0,0 +1,9 @@ +use std::path::PathBuf; + +fn main() -> anyhow::Result<()> { + let schema_root = std::env::args_os() + .nth(1) + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("schema")); + codex_hooks::write_schema_fixtures(&schema_root) +} diff --git a/codex-rs/hooks/src/config_rules.rs b/codex-rs/hooks/src/config_rules.rs new file mode 100644 index 0000000000000000000000000000000000000000..050932ecbbff46ef87a66753496e8036d5ed47e5 --- /dev/null +++ b/codex-rs/hooks/src/config_rules.rs @@ -0,0 +1,259 @@ +use std::collections::HashMap; + +use codex_config::ConfigLayerSource; +use codex_config::ConfigLayerStack; +use codex_config::HookStateToml; +use codex_config::TomlValue; + +/// Build effective hook state from config layers that are allowed to override +/// user preferences. +/// +/// This intentionally reads only user and session flag layers, including +/// disabled layers, to match the skills config behavior. Project, managed, and +/// plugin layers can discover hooks, but they do not get to write user hook +/// state. +pub fn hook_states_from_stack( + config_layer_stack: Option<&ConfigLayerStack>, +) -> HashMap { + let Some(config_layer_stack) = config_layer_stack else { + return HashMap::new(); + }; + + let mut states: HashMap = HashMap::new(); + for layer in config_layer_stack.all_layers_low_to_high() { + if !matches!( + layer.name, + ConfigLayerSource::User { .. } | ConfigLayerSource::SessionFlags + ) { + continue; + } + + let Some(state_value) = layer + .config + .get("hooks") + .and_then(|hooks| hooks.get("state")) + else { + continue; + }; + let TomlValue::Table(state_by_key) = state_value else { + continue; + }; + + for (key, state_value) in state_by_key { + let state: HookStateToml = match state_value.clone().try_into() { + Ok(state) => state, + Err(_) => { + continue; + } + }; + let key = key.trim(); + if key.is_empty() { + continue; + } + // Later layers win field-by-field so a future per-hook state write + // does not accidentally erase an existing enablement override. + let effective_state = states.entry(key.to_string()).or_default(); + if let Some(enabled) = state.enabled { + effective_state.enabled = Some(enabled); + } + if let Some(trusted_hash) = state.trusted_hash { + effective_state.trusted_hash = Some(trusted_hash); + } + } + } + + states +} + +#[cfg(test)] +mod tests { + use codex_config::ConfigLayerEntry; + use codex_config::TomlValue; + use codex_utils_absolute_path::test_support::PathBufExt; + use codex_utils_absolute_path::test_support::test_path_buf; + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn hook_states_from_stack_respects_layer_precedence() { + let key = "file:/tmp/hooks.json:pre_tool_use:0:0"; + let stack = ConfigLayerStack::new( + vec![ + ConfigLayerEntry::new( + ConfigLayerSource::User { + file: test_path_buf("/tmp/config.toml").abs(), + profile: None, + }, + config_with_hook_override(key, Some(/*enabled*/ false)), + ), + ConfigLayerEntry::new( + ConfigLayerSource::SessionFlags, + config_with_hook_override(key, Some(/*enabled*/ true)), + ), + ], + Default::default(), + Default::default(), + ) + .expect("config layer stack"); + + assert_eq!( + hook_states_from_stack(Some(&stack)), + HashMap::from([( + key.to_string(), + HookStateToml { + enabled: Some(true), + trusted_hash: None, + }, + )]) + ); + } + + #[test] + fn hook_states_from_stack_merges_fields_across_layers() { + let key = "file:/tmp/hooks.json:pre_tool_use:0:0"; + let stack = ConfigLayerStack::new( + vec![ + ConfigLayerEntry::new( + ConfigLayerSource::User { + file: test_path_buf("/tmp/config.toml").abs(), + profile: None, + }, + config_with_hook_state( + key, + HookStateToml { + enabled: Some(/*enabled*/ false), + trusted_hash: None, + }, + ), + ), + ConfigLayerEntry::new( + ConfigLayerSource::SessionFlags, + config_with_hook_state( + key, + HookStateToml { + enabled: None, + trusted_hash: Some("sha256:trusted".to_string()), + }, + ), + ), + ], + Default::default(), + Default::default(), + ) + .expect("config layer stack"); + + assert_eq!( + hook_states_from_stack(Some(&stack)), + HashMap::from([( + key.to_string(), + HookStateToml { + enabled: Some(false), + trusted_hash: Some("sha256:trusted".to_string()), + }, + )]) + ); + } + + #[test] + fn hook_states_from_stack_ignores_malformed_hook_events() { + let key = "file:/tmp/hooks.json:pre_tool_use:0:0"; + let config: TomlValue = serde_json::from_value(serde_json::json!({ + "hooks": { + "state": { + (key): { + "enabled": false, + }, + }, + "SessionStart": "not a matcher list", + }, + })) + .expect("config TOML should deserialize"); + let stack = ConfigLayerStack::new( + vec![ConfigLayerEntry::new( + ConfigLayerSource::User { + file: test_path_buf("/tmp/config.toml").abs(), + profile: None, + }, + config, + )], + Default::default(), + Default::default(), + ) + .expect("config layer stack"); + + assert_eq!( + hook_states_from_stack(Some(&stack)), + HashMap::from([( + key.to_string(), + HookStateToml { + enabled: Some(false), + trusted_hash: None, + }, + )]) + ); + } + + #[test] + fn hook_states_from_stack_ignores_malformed_state_entries() { + let key = "file:/tmp/hooks.json:pre_tool_use:0:0"; + let config: TomlValue = serde_json::from_value(serde_json::json!({ + "hooks": { + "state": { + (key): { + "enabled": false, + }, + "malformed": { + "enabled": "not a bool", + }, + }, + }, + })) + .expect("config TOML should deserialize"); + let stack = ConfigLayerStack::new( + vec![ConfigLayerEntry::new( + ConfigLayerSource::User { + file: test_path_buf("/tmp/config.toml").abs(), + profile: None, + }, + config, + )], + Default::default(), + Default::default(), + ) + .expect("config layer stack"); + + assert_eq!( + hook_states_from_stack(Some(&stack)), + HashMap::from([( + key.to_string(), + HookStateToml { + enabled: Some(false), + trusted_hash: None, + }, + )]) + ); + } + + fn config_with_hook_override(key: &str, enabled: Option) -> TomlValue { + config_with_hook_state( + key, + HookStateToml { + enabled, + trusted_hash: None, + }, + ) + } + + fn config_with_hook_state(key: &str, state: HookStateToml) -> TomlValue { + let hook_state = serde_json::to_value(state).expect("hook state should serialize"); + serde_json::from_value(serde_json::json!({ + "hooks": { + "state": { + (key): hook_state, + }, + }, + })) + .expect("config TOML should deserialize") + } +} diff --git a/codex-rs/hooks/src/declarations.rs b/codex-rs/hooks/src/declarations.rs new file mode 100644 index 0000000000000000000000000000000000000000..42ca76a2ed029bf0c71eba93e74a562da1c457f6 --- /dev/null +++ b/codex-rs/hooks/src/declarations.rs @@ -0,0 +1,102 @@ +use codex_plugin::PluginHookSource; +use codex_protocol::protocol::HookEventName; + +/// Minimal declaration metadata for one bundled plugin hook handler. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PluginHookDeclaration { + pub key: String, + pub event_name: HookEventName, +} + +/// Return the hook handlers declared by plugin bundles without projecting live runtime state. +pub fn plugin_hook_declarations(hook_sources: &[PluginHookSource]) -> Vec { + let mut declarations = Vec::new(); + + for source in hook_sources { + let key_source = plugin_hook_key_source( + source.plugin_id.as_key().as_str(), + source.source_relative_path.as_str(), + ); + for (event_name, groups) in source.hooks.clone().into_matcher_groups() { + for (group_index, group) in groups.iter().enumerate() { + for (handler_index, _) in group.hooks.iter().enumerate() { + declarations.push(PluginHookDeclaration { + key: crate::hook_key(&key_source, event_name, group_index, handler_index), + event_name, + }); + } + } + } + } + + declarations +} + +pub(crate) fn plugin_hook_key_source(plugin_id: &str, source_relative_path: &str) -> String { + format!("{plugin_id}:{source_relative_path}") +} + +#[cfg(test)] +mod tests { + use codex_config::HookEventsToml; + use codex_config::HookHandlerConfig; + use codex_config::MatcherGroup; + use codex_plugin::PluginId; + use codex_utils_absolute_path::test_support::PathBufExt; + use codex_utils_absolute_path::test_support::test_path_buf; + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn lists_declared_plugin_handlers_with_persisted_hook_keys() { + let plugin_root = test_path_buf("/tmp/plugin").abs(); + let source_path = plugin_root.join("hooks/hooks.json"); + let declarations = plugin_hook_declarations(&[PluginHookSource { + plugin_id: PluginId::parse("demo@test").expect("plugin id"), + plugin_root: plugin_root.clone(), + plugin_data_root: plugin_root.join("data"), + source_path, + source_relative_path: "hooks/hooks.json".to_string(), + hooks: HookEventsToml { + pre_tool_use: vec![MatcherGroup { + matcher: None, + hooks: vec![ + HookHandlerConfig::Prompt {}, + HookHandlerConfig::Command { + command: "echo hi".to_string(), + command_windows: None, + timeout_sec: None, + r#async: false, + status_message: None, + additional_context_limit: None, + }, + ], + }], + session_start: vec![MatcherGroup { + matcher: None, + hooks: vec![HookHandlerConfig::Agent {}], + }], + ..Default::default() + }, + }]); + + assert_eq!( + declarations, + vec![ + PluginHookDeclaration { + key: "demo@test:hooks/hooks.json:pre_tool_use:0:0".to_string(), + event_name: HookEventName::PreToolUse, + }, + PluginHookDeclaration { + key: "demo@test:hooks/hooks.json:pre_tool_use:0:1".to_string(), + event_name: HookEventName::PreToolUse, + }, + PluginHookDeclaration { + key: "demo@test:hooks/hooks.json:session_start:0:0".to_string(), + event_name: HookEventName::SessionStart, + }, + ] + ); + } +} diff --git a/codex-rs/hooks/src/engine/command_runner.rs b/codex-rs/hooks/src/engine/command_runner.rs new file mode 100644 index 0000000000000000000000000000000000000000..36914c62873f582674317c66de5e83e8225d8576 --- /dev/null +++ b/codex-rs/hooks/src/engine/command_runner.rs @@ -0,0 +1,466 @@ +use std::collections::HashMap; +#[cfg(not(windows))] +use std::ffi::OsStr; +use std::ffi::OsString; +use std::future::Future; +use std::io::ErrorKind; +use std::path::Path; +use std::process::Stdio; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::MutexGuard; +use std::time::Duration; +use std::time::Instant; + +use async_channel::Sender; +use codex_protocol::shell_environment::scrub_non_inheritable_env_vars; +#[cfg(windows)] +use codex_utils_pty::JobObject; +use futures::future::try_join; +use tokio::io::AsyncWriteExt; +use tokio::process::Command; +use tokio::sync::Semaphore; +use tokio::task::JoinSet; +use tokio::time::timeout; +use tracing::Span; + +use super::CommandShell; +use super::ConfiguredHandler; +use super::ConfiguredHandlerKind; +use super::HandlerRunResult; +use super::dispatcher::ParsedHandler; +use super::dispatcher::hook_event_name_label; +use super::dispatcher::hook_execution_mode_label; +use super::dispatcher::hook_handler_type_label; +use super::dispatcher::hook_scope_label; +use super::dispatcher::hook_source_label; +use super::dispatcher::scope_for_event; +use crate::output_spill::AdditionalContext; +use crate::output_spill::HookOutputSpiller; +use codex_protocol::ThreadId; +use codex_protocol::protocol::HookCompletedEvent; +use codex_protocol::protocol::HookHandlerType; +use codex_protocol::protocol::HookOutputEntry; +use codex_protocol::protocol::HookOutputEntryKind; + +const MAX_CONCURRENT_ASYNC_HOOKS: usize = 8; + +/// Owns command execution and bounded asynchronous work for one session. +#[derive(Clone)] +pub(crate) struct CommandHookRuntime { + shell: CommandShell, + environment: Arc>, + result_sender: Sender, + state: Arc>, + output_spiller: HookOutputSpiller, +} + +struct CommandHookRuntimeState { + concurrency_limit: Arc, + tasks: JoinSet<()>, +} + +impl Default for CommandHookRuntimeState { + fn default() -> Self { + Self { + concurrency_limit: Arc::new(Semaphore::new(MAX_CONCURRENT_ASYNC_HOOKS)), + tasks: JoinSet::new(), + } + } +} + +impl CommandHookRuntime { + pub(crate) fn new( + shell: CommandShell, + environment: Arc>, + thread_id: ThreadId, + result_sender: Sender, + ) -> Self { + Self { + shell, + environment, + result_sender, + state: Arc::new(Mutex::new(CommandHookRuntimeState::default())), + output_spiller: HookOutputSpiller::new(thread_id), + } + } + + fn lock_state(&self) -> MutexGuard<'_, CommandHookRuntimeState> { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + pub(crate) fn reconfigured(&self, shell: CommandShell) -> Self { + Self { + shell, + environment: Arc::clone(&self.environment), + result_sender: self.result_sender.clone(), + state: Arc::clone(&self.state), + output_spiller: self.output_spiller.clone(), + } + } + + pub(crate) fn output_spiller(&self) -> &HookOutputSpiller { + &self.output_spiller + } + + pub(crate) fn schedule_async_hook( + &self, + handler: ConfiguredHandler, + input_json: String, + cwd: std::path::PathBuf, + turn_id: Option, + parse: fn(&ConfiguredHandler, HandlerRunResult, Option) -> ParsedHandler, + ) { + if self.result_sender.is_closed() { + return; + } + + let result_sender = self.result_sender.clone(); + let runtime = self.clone(); + self.schedule_async_task(async move { + let result = match &handler.kind { + ConfiguredHandlerKind::Command { command, env, .. } => { + run_command(&runtime, &handler, command, env, &input_json, &cwd).await + } + ConfiguredHandlerKind::McpTool { .. } => return, + }; + let mut hook_result = parse(&handler, result, turn_id).completed; + let mut entries = Vec::new(); + let mut warnings = Vec::new(); + + for entry in std::mem::take(&mut hook_result.run.entries) { + match entry.kind { + HookOutputEntryKind::Context => { + if let Some(text) = runtime + .output_spiller + .maybe_spill_additional_contexts(vec![AdditionalContext { + text: entry.text, + limit: handler.additional_context_limit, + }]) + .await + .into_iter() + .next() + { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Context, + text, + }); + } + } + HookOutputEntryKind::Warning => warnings.push(entry), + HookOutputEntryKind::Error => entries.push(entry), + HookOutputEntryKind::Stop | HookOutputEntryKind::Feedback => {} + } + } + + entries.extend(warnings); + hook_result.run.entries = entries; + let _ = result_sender.try_send(hook_result); + }); + } + + pub(crate) fn schedule_async_task(&self, task: impl Future + Send + 'static) { + let mut state = self.lock_state(); + if state.concurrency_limit.is_closed() { + return; + } + + while state.tasks.try_join_next().is_some() {} + let concurrency_limit = Arc::clone(&state.concurrency_limit); + state.tasks.spawn(async move { + let Ok(_permit) = concurrency_limit.acquire_owned().await else { + return; + }; + task.await; + }); + } + + pub(crate) async fn shutdown(&self) { + let mut tasks = { + let mut state = self.lock_state(); + state.concurrency_limit.close(); + std::mem::take(&mut state.tasks) + }; + tasks.abort_all(); + while tasks.join_next().await.is_some() {} + } +} + +#[tracing::instrument( + name = "codex.hooks.command", + level = "trace", + skip_all, + fields( + hook.event_name = hook_event_name_label(handler.event_name), + hook.handler_type = hook_handler_type_label(HookHandlerType::Command), + hook.execution_mode = hook_execution_mode_label(handler.execution_mode()), + hook.scope = hook_scope_label(scope_for_event(handler.event_name)), + hook.source = hook_source_label(handler.source), + hook.display_order = handler.display_order, + hook.timeout_sec = handler.timeout_sec, + hook.command_outcome = tracing::field::Empty, + ) +)] +pub(crate) async fn run_command( + runtime: &CommandHookRuntime, + handler: &ConfiguredHandler, + command: &str, + env: &HashMap, + input_json: &str, + cwd: &Path, +) -> HandlerRunResult { + let started_at = chrono::Utc::now().timestamp(); + let started = Instant::now(); + + let mut command = build_command(&runtime.shell, command, &runtime.environment, env); + command + .current_dir(cwd) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true); + + #[cfg(unix)] + // Keep process-group cleanup without inheriting the controlling terminal, where + // shell startup can otherwise stop the hook on background terminal I/O. + // SAFETY: detach_from_tty only performs async-signal-safe process setup. + unsafe { + command.pre_exec(codex_utils_pty::process_group::detach_from_tty); + } + + #[cfg(windows)] + let mut process_tree_job = JobObject::create().ok(); + #[cfg(windows)] + let child = match process_tree_job.as_ref() { + Some(job) => match job.spawn_contained(&mut command) { + Ok(child) => Ok(child), + Err(_) => { + process_tree_job = None; + command.creation_flags(0); + command.spawn() + } + }, + None => command.spawn(), + }; + #[cfg(not(windows))] + let child = command.spawn(); + + let mut child = match child { + Ok(child) => child, + Err(err) => { + return finish_command_run( + started_at, + started, + CommandRunCompletion { + exit_code: None, + stdout: String::new(), + stderr: String::new(), + error: Some(err.to_string()), + outcome: "spawn_error", + }, + ); + } + }; + + let mut process_tree_guard = ProcessTreeGuard { + process_id: child.id(), + #[cfg(windows)] + job: process_tree_job, + }; + + let stdin = child.stdin.take(); + let write_stdin = async { + if let Some(mut stdin) = stdin + && let Err(err) = stdin.write_all(input_json.as_bytes()).await + && err.kind() != ErrorKind::BrokenPipe + { + return Err(("stdin_error", format!("failed to write hook stdin: {err}"))); + } + Ok(()) + }; + let wait_with_output = async { + child + .wait_with_output() + .await + .map_err(|err| ("wait_error", err.to_string())) + }; + + let timeout_duration = Duration::from_secs(handler.timeout_sec); + // Drain output while sending input so neither pipe can block the other, and + // include stdin writes in the deadline even when the hook never reads them. + match timeout(timeout_duration, try_join(write_stdin, wait_with_output)).await { + Ok(Ok(((), output))) => { + // Successfully completed hooks may intentionally leave detached helpers running. + #[cfg(windows)] + if let Some(job) = process_tree_guard.job.as_ref() { + let _ = job.preserve_descendants(); + } + process_tree_guard.process_id = None; + finish_command_run( + started_at, + started, + CommandRunCompletion { + exit_code: output.status.code(), + stdout: String::from_utf8_lossy(&output.stdout).to_string(), + stderr: String::from_utf8_lossy(&output.stderr).to_string(), + error: None, + outcome: "completed", + }, + ) + } + Ok(Err((outcome, error))) => finish_command_run( + started_at, + started, + CommandRunCompletion { + exit_code: None, + stdout: String::new(), + stderr: String::new(), + error: Some(error), + outcome, + }, + ), + Err(_) => finish_command_run( + started_at, + started, + CommandRunCompletion { + exit_code: None, + stdout: String::new(), + stderr: String::new(), + error: Some(format!("hook timed out after {}s", handler.timeout_sec)), + outcome: "timeout", + }, + ), + } +} + +// Needed only until command hooks move to the exec server, which owns process-tree cleanup. +struct ProcessTreeGuard { + process_id: Option, + #[cfg(windows)] + job: Option, +} + +impl Drop for ProcessTreeGuard { + fn drop(&mut self) { + let Some(process_id) = self.process_id else { + return; + }; + + #[cfg(unix)] + { + let _ = codex_utils_pty::process_group::kill_process_group(process_id); + } + + #[cfg(windows)] + { + if let Some(job) = self.job.as_ref() { + let _ = job.terminate(); + } else { + let _ = std::process::Command::new("taskkill") + .args(["/PID", &process_id.to_string(), "/T", "/F"]) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn(); + } + } + } +} + +struct CommandRunCompletion { + exit_code: Option, + stdout: String, + stderr: String, + error: Option, + outcome: &'static str, +} + +fn finish_command_run( + started_at: i64, + started: Instant, + completion: CommandRunCompletion, +) -> HandlerRunResult { + Span::current().record("hook.command_outcome", completion.outcome); + HandlerRunResult { + started_at, + completed_at: chrono::Utc::now().timestamp(), + duration_ms: started.elapsed().as_millis().try_into().unwrap_or(i64::MAX), + exit_code: completion.exit_code, + stdout: completion.stdout, + stderr: completion.stderr, + error: completion.error, + } +} + +fn build_command( + shell: &CommandShell, + command_line: &str, + environment: &[(OsString, OsString)], + env: &HashMap, +) -> Command { + let mut command = if shell.program.is_empty() { + default_shell_command(environment) + } else { + Command::new(&shell.program) + }; + if shell.program.is_empty() { + #[cfg(windows)] + command.raw_arg(format!(r#""{command_line}""#)); + + #[cfg(not(windows))] + command.arg(command_line); + } else { + command.args(&shell.args); + + #[cfg(windows)] + if shell.args.iter().any(|arg| arg.eq_ignore_ascii_case("/c")) { + command.raw_arg(format!(r#""{command_line}""#)); + } else { + command.arg(command_line); + } + + #[cfg(not(windows))] + command.arg(command_line); + } + // Replay the session snapshot instead of inheriting the live process environment. + command.env_clear(); + command.envs(environment.iter().cloned()); + command.envs(env); + scrub_non_inheritable_env_vars(command.as_std_mut()); + command +} + +fn default_shell_command(environment: &[(OsString, OsString)]) -> Command { + #[cfg(windows)] + let (environment_variable, fallback_program, argument) = ("COMSPEC", "cmd.exe", "/C"); + + #[cfg(not(windows))] + let (environment_variable, fallback_program, argument) = ("SHELL", "/bin/sh", "-lc"); + + let program = environment + .iter() + .find(|(key, _)| { + #[cfg(windows)] + { + key.to_str() + .is_some_and(|key| key.eq_ignore_ascii_case(environment_variable)) + } + + #[cfg(not(windows))] + { + key == OsStr::new(environment_variable) + } + }) + .map(|(_, value)| value.clone()) + .unwrap_or_else(|| OsString::from(fallback_program)); + + let mut command = Command::new(program); + command.arg(argument); + command +} + +#[cfg(test)] +#[path = "command_runner_tests.rs"] +mod tests; diff --git a/codex-rs/hooks/src/engine/command_runner_tests.rs b/codex-rs/hooks/src/engine/command_runner_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..d11da7765682ed849d75b57545bfe682899b8c4e --- /dev/null +++ b/codex-rs/hooks/src/engine/command_runner_tests.rs @@ -0,0 +1,695 @@ +use std::collections::HashMap; +use std::ffi::OsStr; +use std::ffi::OsString; +#[cfg(windows)] +use std::fs; +#[cfg(unix)] +use std::os::unix::ffi::OsStringExt; +use std::path::Path; +use std::sync::Arc; +use std::time::Duration; + +use async_channel::Receiver; +use codex_protocol::ThreadId; +use codex_protocol::protocol::HookCompletedEvent; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookOutputEntry; +use codex_protocol::protocol::HookOutputEntryKind; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookSource; +use codex_protocol::shell_environment::CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR; +use codex_utils_absolute_path::AbsolutePathBuf; +use pretty_assertions::assert_eq; +use tempfile::TempDir; +use tempfile::tempdir; +use tokio::time::sleep; +use tokio::time::timeout; + +use super::super::ClaudeHooksEngine; +use super::super::ConfiguredHandlerKind; +use super::super::tests::mcp_executor; +use super::CommandHookRuntime; +use super::CommandShell; +use super::ConfiguredHandler; +use super::MAX_CONCURRENT_ASYNC_HOOKS; +use super::build_command; +use super::run_command; +use crate::events::user_prompt_submit::UserPromptSubmitRequest; + +#[cfg(unix)] +#[tokio::test] +async fn hook_shell_startup_does_not_stop_on_controlling_terminal() { + const CHILD_ENV: &str = "CODEX_HOOK_TERMINAL_TEST_CHILD"; + const TEST_NAME: &str = + "engine::command_runner::tests::hook_shell_startup_does_not_stop_on_controlling_terminal"; + + if std::env::var_os(CHILD_ENV).is_none() { + // Re-exec under a controlling terminal even when the test runner has none. + let executable = std::env::current_exe().expect("current test executable"); + let mut env = std::env::vars().collect::>(); + env.insert(CHILD_ENV.to_string(), "1".to_string()); + let codex_utils_pty::SpawnedProcess { + session: _session, + mut stdout_rx, + exit_rx, + .. + } = codex_utils_pty::spawn_pty_process( + executable.to_str().expect("UTF-8 test executable path"), + &[ + TEST_NAME.to_string(), + "--exact".to_string(), + "--nocapture".to_string(), + ], + &std::env::current_dir().expect("current test directory"), + &env, + /*arg0*/ &None, + codex_utils_pty::TerminalSize::default(), + &[], + ) + .await + .expect("spawn test with a controlling terminal"); + + let (exit_code, output) = timeout(Duration::from_secs(10), async { + let output = async { + let mut output = Vec::new(); + while let Some(chunk) = stdout_rx.recv().await { + output.extend_from_slice(&chunk); + } + output + }; + tokio::join!(exit_rx, output) + }) + .await + .expect("terminal hook test should finish"); + let output = String::from_utf8_lossy(&output); + assert_eq!( + exit_code.expect("terminal hook test exit status"), + 0, + "{output}" + ); + assert!( + output.contains(TEST_NAME), + "child did not run test: {output}" + ); + return; + } + + std::fs::File::open("/dev/tty").expect("test process must have a controlling terminal"); + let temp = tempdir().expect("create temp dir"); + let startup_path = temp.path().join("bashenv"); + std::fs::write( + &startup_path, + "command -v stty >/dev/null || exit 1\nstty sane < /dev/tty\n", + ) + .expect("write shell startup fixture"); + let command = "printf hook-ran"; + let env = HashMap::from([( + "BASH_ENV".to_string(), + startup_path + .to_str() + .expect("UTF-8 startup path") + .to_string(), + )]); + let handler = ConfiguredHandler { + builtin: false, + event_name: HookEventName::SessionStart, + matcher: None, + timeout_sec: 2, + status_message: None, + additional_context_limit: Default::default(), + source_path: AbsolutePathBuf::try_from(temp.path().join("hooks.json")) + .expect("absolute hook configuration path") + .into(), + source: HookSource::User, + display_order: 0, + kind: ConfiguredHandlerKind::Command { + command: command.to_string(), + r#async: false, + env: env.clone(), + }, + }; + let (result_sender, _result_receiver) = async_channel::unbounded(); + let runtime = CommandHookRuntime::new( + CommandShell { + program: "/bin/bash".to_string(), + args: vec!["-c".to_string()], + }, + Arc::new(std::env::vars_os().collect()), + ThreadId::new(), + result_sender, + ); + + let result = run_command(&runtime, &handler, command, &env, "{}", temp.path()).await; + + assert_eq!(result.exit_code, Some(0), "stderr: {}", result.stderr); + assert_eq!(result.stdout, "hook-ran"); + assert_eq!(result.error, None); +} + +#[cfg(windows)] +#[tokio::test] +async fn cmd_shell_runs_quoted_hook_command_path() { + let temp = tempdir().expect("create temp dir"); + let hook_dir = temp.path().join("hook with spaces"); + fs::create_dir(&hook_dir).expect("create hook dir"); + let hook_path = hook_dir.join("hook.cmd"); + fs::write( + &hook_path, + "@echo off\r\nif not \"%~1\"==\"notify\" exit /B 7\r\necho hook-ran\r\n", + ) + .expect("write hook command"); + let source_path = + AbsolutePathBuf::try_from(hook_path.clone()).expect("absolute hook command path"); + let command = format!(r#""{}" notify"#, hook_path.display()); + let env = HashMap::new(); + let handler = ConfiguredHandler { + builtin: false, + event_name: HookEventName::SessionStart, + matcher: None, + timeout_sec: 10, + status_message: None, + additional_context_limit: Default::default(), + source_path: source_path.into(), + source: HookSource::User, + display_order: 0, + kind: ConfiguredHandlerKind::Command { + command: command.clone(), + r#async: false, + env: env.clone(), + }, + }; + let shells = [ + CommandShell { + program: String::new(), + args: Vec::new(), + }, + CommandShell { + program: std::env::var("COMSPEC").unwrap_or_else(|_| "cmd.exe".to_string()), + args: vec!["/c".to_string()], + }, + ]; + + for shell in shells { + let (result_sender, _result_receiver) = async_channel::unbounded(); + let runtime = CommandHookRuntime::new( + shell, + Arc::new(std::env::vars_os().collect()), + ThreadId::new(), + result_sender, + ); + let result = run_command(&runtime, &handler, &command, &env, "{}", temp.path()).await; + + assert_eq!(result.exit_code, Some(0), "stderr: {}", result.stderr); + assert_eq!(result.stdout.trim(), "hook-ran"); + assert!(result.error.is_none()); + } +} + +#[tokio::test] +async fn fast_exiting_hook_preserves_stdout_when_stdin_is_not_consumed() { + let temp = tempdir().expect("create temp dir"); + let source_path = AbsolutePathBuf::try_from(temp.path().join("hooks.json")) + .expect("absolute hook configuration path"); + let command = "echo hook-ran"; + let env = HashMap::new(); + let handler = ConfiguredHandler { + builtin: false, + event_name: HookEventName::SessionStart, + matcher: None, + timeout_sec: 10, + status_message: None, + additional_context_limit: Default::default(), + source_path: source_path.into(), + source: HookSource::User, + display_order: 0, + kind: ConfiguredHandlerKind::Command { + command: command.to_string(), + r#async: false, + env: env.clone(), + }, + }; + let input_json = format!(r#"{{"padding":"{}"}}"#, "x".repeat(1024 * 1024)); + let (runtime, _result_receiver) = runtime(); + + let result = run_command(&runtime, &handler, command, &env, &input_json, temp.path()).await; + + assert_eq!(result.exit_code, Some(0), "stderr: {}", result.stderr); + assert_eq!(result.stdout.trim(), "hook-ran"); + assert_eq!(result.error, None); +} + +#[tokio::test] +async fn hook_drains_output_and_times_out_while_stdin_is_blocked() { + let temp = tempdir().expect("create temp dir"); + let mut handler = write_handler( + &temp, + r#"from pathlib import Path +import sys +import time + +sys.stdout.write("x" * 1024 * 1024) +sys.stdout.flush() +sys.stderr.write("x" * 1024 * 1024) +sys.stderr.flush() +Path("output-drained").touch() +time.sleep(30) +"#, + ); + handler.timeout_sec = 2; + let ConfiguredHandlerKind::Command { command, env, .. } = &handler.kind else { + panic!("expected command hook"); + }; + let input_json = format!(r#"{{"padding":"{}"}}"#, "x".repeat(1024 * 1024)); + let (runtime, _result_receiver) = runtime(); + + // Keep user shell startup files out of the pipe I/O timeout test. + let runtime = runtime.reconfigured(if cfg!(windows) { + CommandShell { + program: "cmd.exe".to_string(), + args: vec!["/D".to_string(), "/C".to_string()], + } + } else { + CommandShell { + program: "/bin/sh".to_string(), + args: vec!["-c".to_string()], + } + }); + + let result = timeout( + Duration::from_secs(10), + run_command(&runtime, &handler, command, env, &input_json, temp.path()), + ) + .await + .expect("the hook timeout must also cover blocked stdin writes"); + + assert!(temp.path().join("output-drained").exists()); + assert_eq!(result.exit_code, None); + assert_eq!(result.error, Some("hook timed out after 2s".to_string())); +} + +#[tokio::test] +async fn command_hook_does_not_expose_configured_noise_auth_token() { + let temp = tempdir().expect("create temp dir"); + let source_path = AbsolutePathBuf::try_from(temp.path().join("hooks.json")) + .expect("absolute hook configuration path"); + let command = if cfg!(windows) { "set" } else { "env" }; + let env = HashMap::from([ + ( + CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR.to_ascii_lowercase(), + "configured-noise-token".to_string(), + ), + ("CODEX_HOOK_SAFE_ENV".to_string(), "visible".to_string()), + ]); + let handler = ConfiguredHandler { + builtin: false, + event_name: HookEventName::SessionStart, + matcher: None, + timeout_sec: 10, + status_message: None, + additional_context_limit: Default::default(), + source_path: source_path.into(), + source: HookSource::User, + display_order: 0, + kind: ConfiguredHandlerKind::Command { + command: command.to_string(), + r#async: false, + env: env.clone(), + }, + }; + let (runtime, _result_receiver) = runtime(); + + let result = run_command(&runtime, &handler, command, &env, "{}", temp.path()).await; + + assert_eq!(result.exit_code, Some(0), "stderr: {}", result.stderr); + assert!(result.stdout.contains("CODEX_HOOK_SAFE_ENV=visible")); + assert!(!result.stdout.lines().any(|line| { + line.split_once('=').is_some_and(|(name, _)| { + name.eq_ignore_ascii_case(CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR) + }) + })); + assert_eq!(result.error, None); +} + +#[test] +fn build_command_replays_snapshot_before_hook_overrides_and_scrubbing() { + #[cfg(unix)] + let non_unicode_value = OsString::from_vec(vec![b'v', 0xff]); + let environment = vec![ + ( + OsString::from("CODEX_HOOK_SNAPSHOT"), + OsString::from("captured"), + ), + ( + OsString::from("CODEX_HOOK_OVERRIDE"), + OsString::from("captured"), + ), + ( + OsString::from(CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR), + OsString::from("captured-noise-token"), + ), + #[cfg(unix)] + ( + OsString::from("CODEX_HOOK_NON_UNICODE"), + non_unicode_value.clone(), + ), + ]; + let env = HashMap::from([ + ("CODEX_HOOK_OVERRIDE".to_string(), "configured".to_string()), + ("CODEX_HOOK_SAFE_ENV".to_string(), "visible".to_string()), + ( + CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR.to_string(), + "configured-noise-token".to_string(), + ), + ]); + let command = build_command( + &CommandShell { + program: "configured-shell".to_string(), + args: vec!["-c".to_string()], + }, + "echo hook-ran", + &environment, + &env, + ); + + assert_eq!( + configured_environment_value(&command, "CODEX_HOOK_SNAPSHOT"), + Some(Some(OsString::from("captured"))) + ); + assert_eq!( + configured_environment_value(&command, "CODEX_HOOK_OVERRIDE"), + Some(Some(OsString::from("configured"))) + ); + assert_eq!( + configured_environment_value(&command, "CODEX_HOOK_SAFE_ENV"), + Some(Some(OsString::from("visible"))) + ); + assert_eq!( + configured_environment_value(&command, CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR), + None + ); + #[cfg(unix)] + assert_eq!( + configured_environment_value(&command, "CODEX_HOOK_NON_UNICODE"), + Some(Some(non_unicode_value)) + ); +} + +#[test] +fn fallback_shell_uses_snapshot() { + #[cfg(windows)] + let (name, program) = ("comspec", r"C:\captured\cmd.exe"); + #[cfg(not(windows))] + let (name, program) = ("SHELL", "/captured/shell"); + let command = build_command( + &CommandShell { + program: String::new(), + args: Vec::new(), + }, + "echo hook-ran", + &[(OsString::from(name), OsString::from(program))], + &HashMap::new(), + ); + + assert_eq!(command.as_std().get_program(), OsStr::new(program)); + #[cfg(not(windows))] + assert_eq!( + command + .as_std() + .get_args() + .map(OsStr::to_os_string) + .collect::>(), + vec![OsString::from("-lc"), OsString::from("echo hook-ran")] + ); +} + +const ASYNC_HOOK_TEST_TIMEOUT: Duration = Duration::from_secs(30); + +fn runtime() -> (CommandHookRuntime, Receiver) { + runtime_with_environment(Arc::new(std::env::vars_os().collect())) +} + +fn runtime_with_environment( + environment: Arc>, +) -> (CommandHookRuntime, Receiver) { + let thread_id = ThreadId::new(); + let (result_sender, result_receiver) = async_channel::unbounded(); + let runtime = CommandHookRuntime::new( + CommandShell { + program: String::new(), + args: Vec::new(), + }, + environment, + thread_id, + result_sender, + ); + (runtime, result_receiver) +} + +fn configured_environment_value( + command: &tokio::process::Command, + name: &str, +) -> Option> { + command + .as_std() + .get_envs() + .find(|(key, _)| *key == OsStr::new(name)) + .map(|(_, value)| value.map(OsStr::to_os_string)) +} + +fn write_handler(temp: &TempDir, source: &str) -> ConfiguredHandler { + let script_path = temp.path().join("async_hook.py"); + std::fs::write(&script_path, source).expect("write async test hook"); + ConfiguredHandler { + builtin: false, + event_name: HookEventName::UserPromptSubmit, + matcher: None, + timeout_sec: 10, + status_message: None, + additional_context_limit: Default::default(), + source_path: AbsolutePathBuf::try_from(temp.path().join("hooks.json")) + .expect("absolute test hook path") + .into(), + source: HookSource::User, + display_order: 0, + kind: ConfiguredHandlerKind::Command { + command: format!("python3 {}", script_path.display()), + r#async: true, + env: HashMap::new(), + }, + } +} + +async fn schedule(runtime: &CommandHookRuntime, handler: ConfiguredHandler, cwd: &Path) { + let engine = ClaudeHooksEngine { + handlers: vec![handler], + warnings: Vec::new(), + required_load_errors: Vec::new(), + command_runtime: runtime.clone(), + mcp_executor: mcp_executor(), + }; + engine + .run_user_prompt_submit(UserPromptSubmitRequest { + session_id: ThreadId::new(), + turn_id: "async-test-turn".to_string(), + subagent: None, + cwd: AbsolutePathBuf::try_from(cwd.to_path_buf()).expect("absolute test hook cwd"), + transcript_path: None, + model: "test-model".to_string(), + permission_mode: "default".to_string(), + prompt: "test prompt".to_string(), + }) + .await; +} + +#[tokio::test] +async fn async_hook_marks_invalid_structured_output_as_failed() { + let temp = TempDir::new().expect("async test directory"); + let (runtime, results) = runtime(); + let handler = write_handler( + &temp, + r#"import sys + +sys.stdin.read() +print('{"systemMessage": 123}') +"#, + ); + + schedule(&runtime, handler, temp.path()).await; + let hook_result = timeout(ASYNC_HOOK_TEST_TIMEOUT, results.recv()) + .await + .expect("invalid async output should still finish its background task") + .expect("result receiver should remain open"); + + assert_eq!(hook_result.run.status, HookRunStatus::Failed); + assert_eq!( + hook_result.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook returned invalid user prompt submit JSON output".to_string(), + }] + ); + + runtime.shutdown().await; +} + +#[tokio::test] +async fn async_hook_result_survives_runtime_reconfiguration() { + let temp = TempDir::new().expect("async test directory"); + let mut environment = std::env::vars_os().collect::>(); + environment.push(( + OsString::from("CODEX_HOOK_CAPTURED_ENV"), + OsString::from("captured"), + )); + let (previous, results) = runtime_with_environment(Arc::new(environment)); + let release_path = temp.path().join("release"); + let handler = write_handler( + &temp, + &format!( + r#"import json +import os +from pathlib import Path +import sys +import time + +json.load(sys.stdin) +while not Path(r"{}").exists(): + time.sleep(0.01) +print(os.environ["CODEX_HOOK_CAPTURED_ENV"]) +"#, + release_path.display() + ), + ); + schedule(&previous, handler, temp.path()).await; + + let reconfigured = previous.reconfigured(CommandShell { + program: String::new(), + args: Vec::new(), + }); + assert!(Arc::ptr_eq( + &previous.environment, + &reconfigured.environment + )); + std::fs::write(release_path, "ready").expect("release async hook"); + + let hook_result = timeout(ASYNC_HOOK_TEST_TIMEOUT, results.recv()) + .await + .expect("in-flight hook should survive runtime reconfiguration") + .expect("result receiver should remain open"); + assert_eq!( + hook_result.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Context, + text: "captured".to_string(), + }] + ); + + reconfigured.shutdown().await; +} + +#[tokio::test] +async fn async_hooks_limit_concurrent_processes_without_dropping_waiting_jobs() { + let temp = TempDir::new().expect("async test directory"); + let (runtime, results) = runtime(); + let started_dir = temp.path().join("started"); + let release_path = temp.path().join("release"); + std::fs::create_dir(&started_dir).expect("create hook marker directory"); + let mut handler = write_handler( + &temp, + &format!( + r#"import os +from pathlib import Path +import sys +import time + +sys.stdin.read() +Path(r"{started}", str(os.getpid())).touch() +while not Path(r"{release}").exists(): + time.sleep(0.01) +print("{{}}") +"#, + started = started_dir.display(), + release = release_path.display(), + ), + ); + handler.timeout_sec = ASYNC_HOOK_TEST_TIMEOUT.as_secs(); + + for _ in 0..=MAX_CONCURRENT_ASYNC_HOOKS { + schedule(&runtime, handler.clone(), temp.path()).await; + } + + let started_count = || { + std::fs::read_dir(&started_dir) + .expect("read hook marker directory") + .count() + }; + timeout(ASYNC_HOOK_TEST_TIMEOUT, async { + while started_count() < MAX_CONCURRENT_ASYNC_HOOKS { + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("all available async hook slots should start"); + assert_eq!(started_count(), MAX_CONCURRENT_ASYNC_HOOKS); + + std::fs::write(release_path, "ready").expect("release running async hooks"); + for _ in 0..=MAX_CONCURRENT_ASYNC_HOOKS { + timeout(ASYNC_HOOK_TEST_TIMEOUT, results.recv()) + .await + .expect("waiting async hook should eventually finish") + .expect("result receiver should remain open"); + } + assert_eq!(started_count(), MAX_CONCURRENT_ASYNC_HOOKS + 1); + + runtime.shutdown().await; +} + +#[tokio::test] +async fn shutdown_aborts_in_flight_async_hooks_without_delivering_context() { + let temp = TempDir::new().expect("async test directory"); + let (runtime, results) = runtime(); + let started_path = temp.path().join("started"); + let release_path = temp.path().join("release"); + let handler = write_handler( + &temp, + &format!( + r#"import json +from pathlib import Path +import sys +import time + +json.load(sys.stdin) +Path(r"{started}").write_text("started", encoding="utf-8") +while not Path(r"{release}").exists(): + time.sleep(0.01) +print(json.dumps({{ + "hookSpecificOutput": {{ + "hookEventName": "UserPromptSubmit", + "additionalContext": "must not be delivered after shutdown" + }} +}})) +"#, + started = started_path.display(), + release = release_path.display(), + ), + ); + schedule(&runtime, handler, temp.path()).await; + + timeout(ASYNC_HOOK_TEST_TIMEOUT, async { + while !started_path.exists() { + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("async hook should start before runtime shutdown"); + + runtime.shutdown().await; + drop(runtime); + std::fs::write(release_path, "ready").expect("release shutdown hook"); + assert!( + timeout(Duration::from_millis(150), results.recv()) + .await + .expect("shutdown should close the result channel") + .is_err(), + "shutdown must not deliver a late async result" + ); +} diff --git a/codex-rs/hooks/src/engine/discovery.rs b/codex-rs/hooks/src/engine/discovery.rs new file mode 100644 index 0000000000000000000000000000000000000000..5003dbd24f081d1bc56a2f897fa51ef557a8caab --- /dev/null +++ b/codex-rs/hooks/src/engine/discovery.rs @@ -0,0 +1,1741 @@ +use std::collections::HashMap; +use std::collections::HashSet; +use std::fs; +use std::path::Path; + +use codex_config::CONFIG_TOML_FILE; +use codex_config::ConfigLayerEntry; +use codex_config::ConfigLayerSource; +use codex_config::ConfigLayerStack; +use codex_config::HookEventsToml; +use codex_config::HookHandlerConfig; +use codex_config::HookStateToml; +use codex_config::HooksFile; +use codex_config::ManagedHooksRequirementsToml; +use codex_config::MatcherGroup; +use codex_config::RequirementSource; +use codex_config::TomlValue; +use codex_config::version_for_toml; +use codex_plugin::PluginHookSource; +use codex_plugin::is_allowlisted_bundled_cleanup_hook; +use codex_protocol::protocol::HookEventName; +use codex_utils_absolute_path::AbsolutePathBuf; +use serde::Deserialize; +use serde::Serialize; + +use super::ConfiguredHandler; +use super::ConfiguredHandlerKind; +use super::HookListEntry; +use super::HookListEntryHandler; +use super::dispatcher::hook_event_name_label; +use crate::config_rules::hook_states_from_stack; +use crate::events::common::matcher_pattern_for_event; +use crate::events::common::validate_matcher_pattern; +use crate::events::session_end::SESSION_END_DEFAULT_TIMEOUT_SEC; +use crate::events::session_end::SESSION_END_MAX_TIMEOUT_SEC; +use crate::output_spill::AdditionalContextLimit; +use crate::output_spill::DEFAULT_HOOK_OUTPUT_TOKEN_LIMIT; +use codex_protocol::protocol::HookSource; +use codex_protocol::protocol::HookTrustStatus; + +pub(crate) struct DiscoveryResult { + pub handlers: Vec, + pub hook_entries: Vec, + pub warnings: Vec, + pub required_load_errors: Vec, +} + +struct HookHandlerSource<'a> { + path: &'a AbsolutePathBuf, + key_source: String, + source: HookSource, + is_managed: bool, + requirement: HookRequirement<'a>, + bypass_hook_trust: bool, + hook_states: &'a HashMap, + env: HashMap, + plugin_id: Option, +} + +enum HookRequirement<'a> { + Required(&'a mut Vec), + Optional, +} + +impl HookHandlerSource<'_> { + fn record_load_failure(&mut self, warning: String, warnings: &mut Vec) { + if let HookRequirement::Required(required_load_errors) = &mut self.requirement { + required_load_errors.push(warning.clone()); + } + warnings.push(warning); + } +} + +struct NormalizedHandler { + config: HookHandlerConfig, + kind: ConfiguredHandlerKind, + timeout_sec: u64, + status_message: Option, + additional_context_limit: Option, +} + +#[derive(Clone, Copy)] +struct HookDiscoveryPolicy { + allow_managed_hooks_only: bool, + bypass_hook_trust: bool, +} + +impl HookDiscoveryPolicy { + fn allows(self, source: &HookHandlerSource<'_>) -> bool { + !self.allow_managed_hooks_only || source.is_managed + } +} + +pub(crate) fn discover_handlers( + config_layer_stack: Option<&ConfigLayerStack>, + plugin_hook_sources: Vec, + plugin_hook_load_warnings: Vec, + bypass_hook_trust: bool, +) -> DiscoveryResult { + let mut handlers = Vec::new(); + let mut hook_entries = Vec::new(); + let mut warnings = plugin_hook_load_warnings; + let mut required_load_errors = Vec::new(); + let mut display_order = 0_i64; + let mut visited_json_hook_folders = HashSet::new(); + let hook_states = hook_states_from_stack(config_layer_stack); + let policy = HookDiscoveryPolicy { + allow_managed_hooks_only: config_layer_stack.is_some_and(|config_layer_stack| { + config_layer_stack + .requirements() + .allow_managed_hooks_only + .as_ref() + .is_some_and(|requirement| requirement.value) + }), + bypass_hook_trust, + }; + + if let Some(config_layer_stack) = config_layer_stack { + required_load_errors = append_managed_requirement_handlers( + &mut handlers, + &mut hook_entries, + &mut warnings, + &mut display_order, + config_layer_stack, + &hook_states, + policy, + ); + + for layer in config_layer_stack.layers_low_to_high() { + let (hook_source, is_managed) = hook_metadata_for_config_layer_source(&layer.name); + let policy_path = config_toml_source_path(layer); + let policy_source = HookHandlerSource { + path: &policy_path, + key_source: policy_path.display().to_string(), + source: hook_source, + is_managed, + requirement: HookRequirement::Optional, + bypass_hook_trust: false, + hook_states: &hook_states, + env: HashMap::new(), + plugin_id: None, + }; + if !policy.allows(&policy_source) { + continue; + } + let json_hooks = match layer.hooks_config_folder() { + Some(config_folder) if visited_json_hook_folders.insert(config_folder.clone()) => { + load_hooks_json(Some(config_folder.as_path()), &mut warnings) + } + _ => None, + }; + let toml_hooks = load_toml_hooks_from_layer(layer, &mut warnings); + + if let (Some((json_source_path, json_events)), Some((toml_source_path, toml_events))) = + (&json_hooks, &toml_hooks) + && !json_events.is_empty() + && !toml_events.is_empty() + { + warnings.push(format!( + "loading hooks from both {} and {}; prefer a single representation for this layer", + json_source_path.display(), + toml_source_path.display() + )); + } + + for (source_path, hook_events) in [json_hooks, toml_hooks].into_iter().flatten() { + append_hook_events( + &mut handlers, + &mut hook_entries, + &mut warnings, + &mut display_order, + HookHandlerSource { + path: &source_path, + key_source: source_path.display().to_string(), + source: hook_source, + is_managed, + requirement: HookRequirement::Optional, + bypass_hook_trust: policy.bypass_hook_trust, + hook_states: &hook_states, + env: HashMap::new(), + plugin_id: None, + }, + hook_events, + policy, + ); + } + } + } + + append_plugin_hook_sources( + &mut handlers, + &mut hook_entries, + &mut warnings, + &mut display_order, + plugin_hook_sources, + &hook_states, + policy, + ); + + DiscoveryResult { + handlers, + hook_entries, + warnings, + required_load_errors, + } +} + +fn append_managed_requirement_handlers( + handlers: &mut Vec, + hook_entries: &mut Vec, + warnings: &mut Vec, + display_order: &mut i64, + config_layer_stack: &ConfigLayerStack, + hook_states: &HashMap, + policy: HookDiscoveryPolicy, +) -> Vec { + let Some(managed_hooks) = config_layer_stack.requirements().managed_hooks.as_ref() else { + return Vec::new(); + }; + let mut required_load_errors = Vec::new(); + let source_path = managed_hooks_source_path(managed_hooks.get(), managed_hooks.source.as_ref()); + append_hook_events( + handlers, + hook_entries, + warnings, + display_order, + HookHandlerSource { + path: &source_path, + key_source: source_path.display().to_string(), + source: hook_source_for_requirement_source(managed_hooks.source.as_ref()), + is_managed: true, + requirement: HookRequirement::Required(&mut required_load_errors), + bypass_hook_trust: false, + hook_states, + env: HashMap::new(), + plugin_id: None, + }, + managed_hooks.get().hooks.clone(), + policy, + ); + required_load_errors +} + +fn append_plugin_hook_sources( + handlers: &mut Vec, + hook_entries: &mut Vec, + warnings: &mut Vec, + display_order: &mut i64, + plugin_hook_sources: Vec, + hook_states: &HashMap, + policy: HookDiscoveryPolicy, +) { + for source in plugin_hook_sources { + let PluginHookSource { + plugin_root, + plugin_id, + plugin_data_root, + source_path, + source_relative_path, + hooks, + } = source; + let mut env = HashMap::new(); + let plugin_root_value = plugin_root.display().to_string(); + let plugin_data_root_value = plugin_data_root.display().to_string(); + env.insert("PLUGIN_ROOT".to_string(), plugin_root_value.clone()); + // For OOTB compat with existing plugins that use this env var. + env.insert("CLAUDE_PLUGIN_ROOT".to_string(), plugin_root_value); + env.insert("PLUGIN_DATA".to_string(), plugin_data_root_value.clone()); + // For OOTB compat with existing plugins that use this env var. + env.insert("CLAUDE_PLUGIN_DATA".to_string(), plugin_data_root_value); + let plugin_id = plugin_id.as_key(); + append_hook_events( + handlers, + hook_entries, + warnings, + display_order, + HookHandlerSource { + path: &source_path, + key_source: crate::declarations::plugin_hook_key_source( + plugin_id.as_str(), + source_relative_path.as_str(), + ), + source: HookSource::Plugin, + is_managed: false, + requirement: HookRequirement::Optional, + bypass_hook_trust: policy.bypass_hook_trust, + hook_states, + env, + plugin_id: Some(plugin_id), + }, + hooks, + policy, + ); + } +} + +fn managed_hooks_source_path( + managed_hooks: &ManagedHooksRequirementsToml, + requirement_source: Option<&RequirementSource>, +) -> AbsolutePathBuf { + if let Some(source_path) = managed_hooks.managed_dir_for_current_platform() + && source_path.is_absolute() + && let Ok(source_path) = AbsolutePathBuf::from_absolute_path(source_path) + { + return source_path; + } + + fallback_managed_hooks_source_path(requirement_source) +} + +fn fallback_managed_hooks_source_path( + requirement_source: Option<&RequirementSource>, +) -> AbsolutePathBuf { + match requirement_source { + Some(RequirementSource::SystemRequirementsToml { file }) + | Some(RequirementSource::LegacyManagedConfigTomlFromFile { file }) => file.clone(), + Some(RequirementSource::MdmManagedPreferences { domain, key }) => { + synthetic_layer_path(&format!("/requirements.toml")) + } + Some(RequirementSource::Composite { .. }) => { + synthetic_layer_path("/requirements.toml") + } + Some(RequirementSource::EnterpriseManaged { id, name }) => { + let name = escape_xml_text(name); + let id = escape_xml_text(id); + synthetic_layer_path(&format!( + "/requirements.toml" + )) + } + Some(RequirementSource::LegacyManagedConfigTomlFromMdm) => { + synthetic_layer_path("/managed_config.toml") + } + Some(RequirementSource::Unknown) | None => { + synthetic_layer_path("/requirements.toml") + } + } +} + +fn load_hooks_json( + config_folder: Option<&Path>, + warnings: &mut Vec, +) -> Option<(AbsolutePathBuf, HookEventsToml)> { + let source_path = config_folder?.join("hooks.json"); + if !source_path.as_path().is_file() { + return None; + } + + let contents = match fs::read_to_string(source_path.as_path()) { + Ok(contents) => contents, + Err(err) => { + warnings.push(format!( + "failed to read hooks config {}: {err}", + source_path.display() + )); + return None; + } + }; + + let parsed: HooksFile = match serde_json::from_str(&contents) { + Ok(parsed) => parsed, + Err(err) => { + warnings.push(format!( + "failed to parse hooks config {}: {err}", + source_path.display() + )); + return None; + } + }; + + let source_path = AbsolutePathBuf::from_absolute_path(&source_path) + .inspect_err(|err| { + warnings.push(format!( + "failed to normalize hooks config path {}: {err}", + source_path.display() + )); + }) + .ok()?; + + (!parsed.hooks.is_empty()).then_some((source_path, parsed.hooks)) +} + +fn load_toml_hooks_from_layer( + layer: &ConfigLayerEntry, + warnings: &mut Vec, +) -> Option<(AbsolutePathBuf, HookEventsToml)> { + let source_path = config_toml_source_path(layer); + let hook_value = layer.config.get("hooks")?.clone(); + let parsed = match HookEventsToml::deserialize(hook_value) { + Ok(parsed) => parsed, + Err(err) => { + warnings.push(format!( + "failed to parse TOML hooks in {}: {err}", + source_path.display() + )); + return None; + } + }; + + (!parsed.is_empty()).then_some((source_path, parsed)) +} + +fn config_toml_source_path(layer: &ConfigLayerEntry) -> AbsolutePathBuf { + match &layer.name { + ConfigLayerSource::PackagedDefaults { file } + | ConfigLayerSource::System { file } + | ConfigLayerSource::User { file, .. } + | ConfigLayerSource::LegacyManagedConfigTomlFromFile { file } => file.clone(), + ConfigLayerSource::Project { dot_codex_folder } => layer + .hooks_config_folder() + .unwrap_or_else(|| dot_codex_folder.clone()) + .join(CONFIG_TOML_FILE), + ConfigLayerSource::Mdm { domain, key } => { + synthetic_layer_path(&format!("/{CONFIG_TOML_FILE}")) + } + ConfigLayerSource::EnterpriseManaged { id, name } => synthetic_layer_path(&format!( + "/{CONFIG_TOML_FILE}" + )), + ConfigLayerSource::LegacyManagedConfigTomlFromMdm => { + synthetic_layer_path("/managed_config.toml") + } + ConfigLayerSource::SessionFlags => synthetic_layer_path("/config.toml"), + } +} + +fn synthetic_layer_path(path: &str) -> AbsolutePathBuf { + #[cfg(windows)] + { + AbsolutePathBuf::resolve_path_against_base(path, r"C:\") + } + + #[cfg(not(windows))] + { + AbsolutePathBuf::resolve_path_against_base(path, "/") + } +} + +fn escape_xml_text(value: &str) -> String { + let mut escaped = String::with_capacity(value.len()); + for ch in value.chars() { + match ch { + '&' => escaped.push_str("&"), + '<' => escaped.push_str("<"), + '>' => escaped.push_str(">"), + '"' => escaped.push_str("""), + '\'' => escaped.push_str("'"), + _ => escaped.push(ch), + } + } + escaped +} + +fn append_hook_events( + handlers: &mut Vec, + hook_entries: &mut Vec, + warnings: &mut Vec, + display_order: &mut i64, + mut source: HookHandlerSource<'_>, + hook_events: HookEventsToml, + policy: HookDiscoveryPolicy, +) { + if !policy.allows(&source) { + return; + } + + for (event_name, groups) in hook_events.into_matcher_groups() { + append_matcher_groups( + handlers, + hook_entries, + warnings, + display_order, + &mut source, + event_name, + groups, + ); + } +} + +fn append_matcher_groups( + handlers: &mut Vec, + hook_entries: &mut Vec, + warnings: &mut Vec, + display_order: &mut i64, + source: &mut HookHandlerSource<'_>, + event_name: codex_protocol::protocol::HookEventName, + groups: Vec, +) { + for (group_index, group) in groups.into_iter().enumerate() { + let matcher = matcher_pattern_for_event(event_name, group.matcher.as_deref()); + if let Some(matcher) = matcher + && let Err(err) = validate_matcher_pattern(matcher) + { + let warning = format!( + "invalid matcher {matcher:?} in {}: {err}", + source.path.display() + ); + if group.hooks.is_empty() { + warnings.push(warning); + } else { + source.record_load_failure(warning, warnings); + } + continue; + } + for (handler_index, handler) in group.hooks.iter().cloned().enumerate() { + let normalized = match handler { + HookHandlerConfig::Command { + command, + command_windows, + timeout_sec, + r#async, + status_message, + additional_context_limit, + } => { + let command = if cfg!(windows) { + command_windows.unwrap_or(command) + } else { + command + }; + if command.trim().is_empty() { + source.record_load_failure( + format!("skipping empty hook command in {}", source.path.display()), + warnings, + ); + continue; + } + let timeout_sec = normalize_command_hook( + event_name, + timeout_sec, + source.path.as_path(), + warnings, + ); + let runs_async = r#async && event_name != HookEventName::SessionEnd; + if r#async && !runs_async { + warnings.push(format!( + "running async {} hook synchronously in {}", + hook_event_name_label(event_name), + source.path.display() + )); + } + let additional_context_limit = if matches!( + event_name, + codex_protocol::protocol::HookEventName::PreToolUse + | codex_protocol::protocol::HookEventName::PostToolUse + | codex_protocol::protocol::HookEventName::SessionStart + | codex_protocol::protocol::HookEventName::UserPromptSubmit + | codex_protocol::protocol::HookEventName::SubagentStart + ) { + additional_context_limit + } else { + if additional_context_limit.is_some() { + warnings.push(format!( + "ignoring additionalContextLimit for {event_name:?} hook in {}: this event cannot emit additionalContext", + source.path.display() + )); + } + None + }; + let normalized_additional_context_limit = additional_context_limit + .filter(|limit| *limit != DEFAULT_HOOK_OUTPUT_TOKEN_LIMIT); + let config = HookHandlerConfig::Command { + command: command.clone(), + command_windows: None, + timeout_sec: Some(timeout_sec), + r#async, + status_message: status_message.clone(), + additional_context_limit: normalized_additional_context_limit, + }; + let command = source.env.iter().fold(command, |command, (key, value)| { + command.replace(&format!("${{{key}}}"), value) + }); + NormalizedHandler { + config, + kind: ConfiguredHandlerKind::Command { + command, + env: source.env.clone(), + r#async: runs_async, + }, + timeout_sec, + status_message, + additional_context_limit, + } + } + HookHandlerConfig::McpTool { + server, + tool, + input, + timeout_sec, + status_message, + } => { + if event_name == HookEventName::SessionEnd { + source.record_load_failure( + format!( + "skipping MCP tool hook in {}: {} MCP hooks are not supported", + source.path.display(), + hook_event_name_label(event_name), + ), + warnings, + ); + continue; + } + if server.trim().is_empty() || tool.trim().is_empty() { + source.record_load_failure( + format!( + "skipping MCP tool hook in {}: server and tool must not be empty", + source.path.display() + ), + warnings, + ); + continue; + } + let timeout_sec = normalize_command_hook( + event_name, + timeout_sec, + source.path.as_path(), + warnings, + ); + let config = HookHandlerConfig::McpTool { + server: server.clone(), + tool: tool.clone(), + input: input.clone(), + timeout_sec: Some(timeout_sec), + status_message: status_message.clone(), + }; + NormalizedHandler { + config, + kind: ConfiguredHandlerKind::McpTool { + server, + tool, + input, + }, + timeout_sec, + status_message, + additional_context_limit: None, + } + } + HookHandlerConfig::Prompt {} => { + source.record_load_failure( + format!( + "skipping prompt hook in {}: prompt hooks are not supported yet", + source.path.display() + ), + warnings, + ); + continue; + } + HookHandlerConfig::Agent {} => { + source.record_load_failure( + format!( + "skipping agent hook in {}: agent hooks are not supported yet", + source.path.display() + ), + warnings, + ); + continue; + } + }; + + let NormalizedHandler { + config, + kind, + timeout_sec, + status_message, + additional_context_limit, + } = normalized; + let current_hash = hook_hash(event_name, matcher, &group, &config); + let key = crate::hook_key(&source.key_source, event_name, group_index, handler_index); + let state = source.hook_states.get(&key); + let builtin = source.plugin_id.as_deref().is_some_and(|plugin_id| { + is_allowlisted_bundled_cleanup_hook( + plugin_id, + event_name, + group.matcher.as_deref(), + &config, + /*app_connector_id*/ None, + ) + }); + let enabled = hook_enabled(source.is_managed, builtin, state); + let trusted_hash = hook_trusted_hash(source.is_managed, state); + let trust_status = + hook_trust_status(source.is_managed, builtin, ¤t_hash, trusted_hash); + let handler = match &kind { + ConfiguredHandlerKind::Command { + command, r#async, .. + } => HookListEntryHandler::Command { + command: command.clone(), + r#async: *r#async, + }, + ConfiguredHandlerKind::McpTool { server, tool, .. } => { + HookListEntryHandler::McpTool { + server: server.clone(), + tool: tool.clone(), + } + } + }; + + hook_entries.push(HookListEntry { + builtin, + key, + event_name, + handler, + matcher: matcher.map(ToOwned::to_owned), + timeout_sec, + status_message: status_message.clone(), + additional_context_limit, + source_path: source.path.clone(), + source: source.source, + plugin_id: source.plugin_id.clone(), + display_order: *display_order, + enabled, + is_managed: source.is_managed, + current_hash, + trust_status, + }); + if enabled + && (source.bypass_hook_trust + || matches!( + trust_status, + HookTrustStatus::Managed | HookTrustStatus::Trusted + )) + { + handlers.push(ConfiguredHandler { + builtin, + event_name, + matcher: matcher.map(ToOwned::to_owned), + timeout_sec, + status_message, + additional_context_limit: AdditionalContextLimit::from_config( + additional_context_limit, + ), + source_path: source.path.clone().into(), + source: source.source, + display_order: *display_order, + kind, + }); + } + *display_order += 1; + } + } +} + +/// Normalizes hook timeouts. SessionEnd and Interrupt default to one second and are capped at three +/// seconds; all other hooks keep the standard ten-minute default. +fn normalize_command_hook( + event_name: HookEventName, + timeout_sec: Option, + source_path: &Path, + warnings: &mut Vec, +) -> u64 { + match event_name { + HookEventName::SessionEnd | HookEventName::Interrupt => { + let max_timeout_sec = SESSION_END_MAX_TIMEOUT_SEC; + if timeout_sec.is_some_and(|timeout_sec| timeout_sec > max_timeout_sec) { + warnings.push(format!( + "clamping {} hook timeout to {max_timeout_sec}s in {}", + hook_event_name_label(event_name), + source_path.display() + )); + } + timeout_sec + .unwrap_or(SESSION_END_DEFAULT_TIMEOUT_SEC) + .clamp(1, max_timeout_sec) + } + _ => timeout_sec.unwrap_or(600).max(1), + } +} + +/// Hash a normalized, config-derived identity instead of source text so equivalent +/// hooks from config TOML and hooks.json converge on the same trust identity. +#[derive(Serialize)] +struct NormalizedHookIdentity { + event_name: &'static str, + #[serde(flatten)] + group: MatcherGroup, +} + +fn hook_hash( + event_name: codex_protocol::protocol::HookEventName, + matcher: Option<&str>, + group: &MatcherGroup, + normalized_handler: &HookHandlerConfig, +) -> String { + let mut group = group.clone(); + group.matcher = matcher.map(ToOwned::to_owned); + group.hooks = vec![normalized_handler.clone()]; + let identity = NormalizedHookIdentity { + event_name: crate::hook_event_key_label(event_name), + group, + }; + let Ok(value) = TomlValue::try_from(identity) else { + unreachable!("normalized hook identity should serialize to TOML"); + }; + version_for_toml(&value) +} + +fn hook_trust_status( + is_managed: bool, + is_builtin: bool, + current_hash: &str, + trusted_hash: Option<&str>, +) -> HookTrustStatus { + if is_builtin { + HookTrustStatus::Trusted + } else if is_managed { + HookTrustStatus::Managed + } else { + match trusted_hash { + Some(trusted_hash) if trusted_hash == current_hash => HookTrustStatus::Trusted, + Some(_) => HookTrustStatus::Modified, + None => HookTrustStatus::Untrusted, + } + } +} + +fn hook_enabled(is_managed: bool, is_builtin: bool, state: Option<&HookStateToml>) -> bool { + is_builtin || is_managed || state.and_then(|state| state.enabled) != Some(false) +} + +fn hook_trusted_hash(is_managed: bool, state: Option<&HookStateToml>) -> Option<&str> { + (!is_managed) + .then(|| state.and_then(|state| state.trusted_hash.as_deref())) + .flatten() +} + +fn hook_metadata_for_config_layer_source(source: &ConfigLayerSource) -> (HookSource, bool) { + match source { + ConfigLayerSource::PackagedDefaults { .. } => (HookSource::Unknown, false), + ConfigLayerSource::System { .. } => (HookSource::System, true), + ConfigLayerSource::User { .. } => (HookSource::User, false), + ConfigLayerSource::Project { .. } => (HookSource::Project, false), + ConfigLayerSource::Mdm { .. } => (HookSource::Mdm, true), + ConfigLayerSource::EnterpriseManaged { .. } => (HookSource::CloudManagedConfig, true), + ConfigLayerSource::SessionFlags => (HookSource::SessionFlags, false), + ConfigLayerSource::LegacyManagedConfigTomlFromFile { .. } => { + (HookSource::LegacyManagedConfigFile, true) + } + ConfigLayerSource::LegacyManagedConfigTomlFromMdm => { + (HookSource::LegacyManagedConfigMdm, true) + } + } +} + +fn hook_source_for_requirement_source(source: Option<&RequirementSource>) -> HookSource { + match source { + Some(RequirementSource::MdmManagedPreferences { .. }) => HookSource::Mdm, + Some(RequirementSource::SystemRequirementsToml { .. }) => HookSource::System, + Some(RequirementSource::LegacyManagedConfigTomlFromFile { .. }) => { + HookSource::LegacyManagedConfigFile + } + Some(RequirementSource::LegacyManagedConfigTomlFromMdm) => { + HookSource::LegacyManagedConfigMdm + } + Some(RequirementSource::Composite { sources }) => { + // Requirements hook composition preserves contributing sources in + // priority order, but discovery only carries one source for the + // whole merged hooks field. Use the primary contributor as the best + // available coarse attribution. + hook_source_for_requirement_source(sources.first()) + } + Some(RequirementSource::EnterpriseManaged { .. }) => HookSource::CloudRequirements, + Some(RequirementSource::Unknown) | None => HookSource::Unknown, + } +} + +#[cfg(test)] +mod tests { + use codex_config::ConfigLayerEntry; + use codex_config::ConfigLayerSource; + use codex_config::HookEventsToml; + use codex_config::RequirementSource; + use codex_protocol::protocol::HookEventName; + use codex_protocol::protocol::HookExecutionMode; + use codex_protocol::protocol::HookSource; + use codex_utils_absolute_path::AbsolutePathBuf; + use codex_utils_absolute_path::test_support::PathBufExt; + use codex_utils_absolute_path::test_support::test_path_buf; + use pretty_assertions::assert_eq; + + use super::ConfiguredHandler; + use super::ConfiguredHandlerKind; + use super::HookListEntry; + use super::HookListEntryHandler; + use super::append_matcher_groups; + use super::normalize_command_hook; + use crate::output_spill::AdditionalContextLimit; + use crate::output_spill::DEFAULT_HOOK_OUTPUT_TOKEN_LIMIT; + use codex_config::HookHandlerConfig; + use codex_config::HookStateToml; + use codex_config::MatcherGroup; + use codex_config::TomlValue; + use codex_protocol::protocol::HookTrustStatus; + + fn source_path() -> AbsolutePathBuf { + test_path_buf("/tmp/hooks.json").abs() + } + + fn hook_source() -> HookSource { + HookSource::System + } + + fn hook_handler_source<'a>( + path: &'a AbsolutePathBuf, + hook_states: &'a std::collections::HashMap, + ) -> super::HookHandlerSource<'a> { + super::HookHandlerSource { + path, + key_source: path.display().to_string(), + source: hook_source(), + is_managed: true, + requirement: super::HookRequirement::Optional, + bypass_hook_trust: false, + hook_states, + env: std::collections::HashMap::new(), + plugin_id: None, + } + } + + fn unmanaged_hook_handler_source<'a>( + path: &'a AbsolutePathBuf, + hook_states: &'a std::collections::HashMap, + bypass_hook_trust: bool, + ) -> super::HookHandlerSource<'a> { + super::HookHandlerSource { + path, + key_source: path.display().to_string(), + source: HookSource::User, + is_managed: false, + requirement: super::HookRequirement::Optional, + bypass_hook_trust, + hook_states, + env: std::collections::HashMap::new(), + plugin_id: None, + } + } + + #[test] + fn composite_requirement_hook_source_uses_primary_source() { + let source = RequirementSource::Composite { + sources: vec![ + RequirementSource::SystemRequirementsToml { + file: test_path_buf("/etc/codex/requirements.toml").abs(), + }, + RequirementSource::EnterpriseManaged { + id: "layer-1".to_string(), + name: "Engineering".to_string(), + }, + ], + }; + + assert_eq!( + super::hook_source_for_requirement_source(Some(&source)), + HookSource::System + ); + } + + #[test] + fn enterprise_managed_synthetic_path_escapes_display_fields() { + let source = RequirementSource::EnterpriseManaged { + id: "id<&>".to_string(), + name: "Name & \"Ops\"".to_string(), + }; + + let source_path = super::fallback_managed_hooks_source_path(Some(&source)); + let source_path = source_path.display().to_string(); + + assert!(source_path.contains("Name <Admin> & "Ops"")); + assert!(source_path.contains("id<&>")); + assert!(!source_path.contains("Name ")); + } + + fn command_group(matcher: Option<&str>) -> MatcherGroup { + MatcherGroup { + matcher: matcher.map(str::to_string), + hooks: vec![HookHandlerConfig::Command { + command: "echo hello".to_string(), + command_windows: None, + timeout_sec: None, + r#async: false, + status_message: None, + additional_context_limit: None, + }], + } + } + + fn command_group_with_additional_context_limit( + additional_context_limit: usize, + ) -> MatcherGroup { + MatcherGroup { + matcher: None, + hooks: vec![HookHandlerConfig::Command { + command: "echo hello".to_string(), + command_windows: None, + timeout_sec: None, + r#async: false, + status_message: None, + additional_context_limit: Some(additional_context_limit), + }], + } + } + + #[test] + fn mcp_tool_hooks_preserve_argument_templates_and_list_their_target() { + let source_path = source_path(); + let hook_states = std::collections::HashMap::new(); + let input = serde_json::from_value(serde_json::json!({ + "file_path": "${tool_input.file_path}", + "optional": "${tool_input.optional}", + })) + .expect("MCP hook input should be an object"); + let mut handlers = Vec::new(); + let mut entries = Vec::new(); + let mut warnings = Vec::new(); + let mut display_order = 0; + + append_matcher_groups( + &mut handlers, + &mut entries, + &mut warnings, + &mut display_order, + &mut hook_handler_source(&source_path, &hook_states), + HookEventName::PostToolUse, + vec![MatcherGroup { + matcher: Some("Write|Edit".to_string()), + hooks: vec![HookHandlerConfig::McpTool { + server: "security".to_string(), + tool: "scan".to_string(), + input, + timeout_sec: Some(30), + status_message: Some("Scanning file".to_string()), + }], + }], + ); + + assert!(warnings.is_empty()); + assert_eq!(handlers.len(), 1); + assert_eq!(entries.len(), 1); + assert_eq!( + entries[0].handler, + HookListEntryHandler::McpTool { + server: "security".to_string(), + tool: "scan".to_string(), + } + ); + assert!(entries[0].current_hash.starts_with("sha256:")); + } + + #[test] + fn session_end_mcp_tool_hooks_are_warned_and_skipped() { + let source_path = source_path(); + let hook_states = std::collections::HashMap::new(); + let mut handlers = Vec::new(); + let mut entries = Vec::new(); + let mut warnings = Vec::new(); + let mut display_order = 0; + + append_matcher_groups( + &mut handlers, + &mut entries, + &mut warnings, + &mut display_order, + &mut hook_handler_source(&source_path, &hook_states), + HookEventName::SessionEnd, + vec![MatcherGroup { + matcher: None, + hooks: vec![HookHandlerConfig::McpTool { + server: "security".to_string(), + tool: "scan".to_string(), + input: serde_json::Map::new(), + timeout_sec: None, + status_message: None, + }], + }], + ); + + assert!(handlers.is_empty()); + assert!(entries.is_empty()); + assert_eq!( + warnings, + vec![format!( + "skipping MCP tool hook in {}: SessionEnd MCP hooks are not supported", + source_path.display() + )] + ); + } + + #[test] + fn interrupt_mcp_tool_hooks_are_supported_and_timeout_is_clamped() { + let source_path = source_path(); + let hook_states = std::collections::HashMap::new(); + let mut handlers = Vec::new(); + let mut entries = Vec::new(); + let mut warnings = Vec::new(); + let mut display_order = 0; + + append_matcher_groups( + &mut handlers, + &mut entries, + &mut warnings, + &mut display_order, + &mut hook_handler_source(&source_path, &hook_states), + HookEventName::Interrupt, + vec![MatcherGroup { + matcher: None, + hooks: vec![ + HookHandlerConfig::McpTool { + server: "security".to_string(), + tool: "scan".to_string(), + input: serde_json::Map::new(), + timeout_sec: None, + status_message: None, + }, + HookHandlerConfig::McpTool { + server: "security".to_string(), + tool: "report".to_string(), + input: serde_json::Map::new(), + timeout_sec: Some(600), + status_message: None, + }, + ], + }], + ); + + assert_eq!( + handlers + .iter() + .map(|handler| handler.timeout_sec) + .collect::>(), + vec![1, 3] + ); + assert_eq!( + entries + .iter() + .map(|entry| entry.timeout_sec) + .collect::>(), + vec![1, 3] + ); + assert!( + entries + .iter() + .all(|entry| matches!(entry.handler, HookListEntryHandler::McpTool { .. })) + ); + assert_eq!( + warnings, + vec![format!( + "clamping Interrupt hook timeout to 3s in {}", + source_path.display() + )] + ); + } + + fn discover_command( + event_name: HookEventName, + additional_context_limit: Option, + ) -> (ConfiguredHandler, HookListEntry, Vec) { + let source_path = source_path(); + let hook_states = std::collections::HashMap::new(); + let mut handlers = Vec::new(); + let mut entries = Vec::new(); + let mut warnings = Vec::new(); + let mut display_order = 0; + append_matcher_groups( + &mut handlers, + &mut entries, + &mut warnings, + &mut display_order, + &mut hook_handler_source(&source_path, &hook_states), + event_name, + vec![match additional_context_limit { + Some(limit) => command_group_with_additional_context_limit(limit), + None => command_group(/*matcher*/ None), + }], + ); + (handlers.remove(0), entries.remove(0), warnings) + } + + #[test] + fn supported_events_retain_per_handler_additional_context_limit_and_hash_it() { + for event_name in [ + HookEventName::PreToolUse, + HookEventName::PostToolUse, + HookEventName::SessionStart, + HookEventName::UserPromptSubmit, + HookEventName::SubagentStart, + ] { + let (_, default_entry, _) = + discover_command(event_name, /*additional_context_limit*/ None); + let (explicit_default_handler, explicit_default_entry, _) = discover_command( + event_name, + /*additional_context_limit*/ Some(DEFAULT_HOOK_OUTPUT_TOKEN_LIMIT), + ); + let (custom_handler, custom_entry, _) = + discover_command(event_name, /*additional_context_limit*/ Some(20_000)); + let (unlimited_handler, unlimited_entry, _) = + discover_command(event_name, /*additional_context_limit*/ Some(0)); + + assert_eq!( + custom_handler.additional_context_limit, + AdditionalContextLimit::from_config(Some(20_000)) + ); + assert_eq!(custom_entry.additional_context_limit, Some(20_000)); + assert_ne!(default_entry.current_hash, custom_entry.current_hash); + assert_eq!( + explicit_default_handler.additional_context_limit, + AdditionalContextLimit::from_config(Some(DEFAULT_HOOK_OUTPUT_TOKEN_LIMIT)) + ); + assert_eq!( + explicit_default_entry.additional_context_limit, + Some(DEFAULT_HOOK_OUTPUT_TOKEN_LIMIT) + ); + assert_eq!( + default_entry.current_hash, + explicit_default_entry.current_hash + ); + assert_eq!( + unlimited_handler.additional_context_limit, + AdditionalContextLimit::from_config(Some(0)) + ); + assert_eq!(unlimited_entry.additional_context_limit, Some(0)); + assert_ne!(default_entry.current_hash, unlimited_entry.current_hash); + } + } + + #[test] + fn unsupported_event_warns_and_ignores_additional_context_limit() { + let source_path = source_path(); + let (handler, _, warnings) = discover_command( + HookEventName::Stop, + /*additional_context_limit*/ Some(4_096), + ); + + assert_eq!(handler.additional_context_limit, Default::default()); + assert_eq!(warnings.len(), 1); + assert!(warnings[0].contains("ignoring additionalContextLimit for Stop hook")); + assert!(warnings[0].contains(&source_path.display().to_string())); + } + + #[test] + fn user_prompt_submit_ignores_invalid_matcher_during_discovery() { + let mut handlers = Vec::new(); + let mut warnings = Vec::new(); + let mut display_order = 0; + let source_path = source_path(); + let hook_states = std::collections::HashMap::new(); + + append_matcher_groups( + &mut handlers, + &mut Vec::new(), + &mut warnings, + &mut display_order, + &mut hook_handler_source(&source_path, &hook_states), + HookEventName::UserPromptSubmit, + vec![command_group(Some("["))], + ); + + assert_eq!(warnings, Vec::::new()); + assert_eq!( + handlers, + vec![ConfiguredHandler { + builtin: false, + event_name: HookEventName::UserPromptSubmit, + matcher: None, + timeout_sec: 600, + status_message: None, + additional_context_limit: Default::default(), + source_path: source_path.clone().into(), + source: hook_source(), + display_order: 0, + kind: ConfiguredHandlerKind::Command { + command: "echo hello".to_string(), + r#async: false, + env: std::collections::HashMap::new(), + }, + }] + ); + } + + #[test] + fn pre_tool_use_keeps_valid_matcher_during_discovery() { + let mut handlers = Vec::new(); + let mut warnings = Vec::new(); + let mut display_order = 0; + let source_path = source_path(); + let hook_states = std::collections::HashMap::new(); + + append_matcher_groups( + &mut handlers, + &mut Vec::new(), + &mut warnings, + &mut display_order, + &mut hook_handler_source(&source_path, &hook_states), + HookEventName::PreToolUse, + vec![command_group(Some("^Bash$"))], + ); + + assert_eq!(warnings, Vec::::new()); + assert_eq!( + handlers, + vec![ConfiguredHandler { + builtin: false, + event_name: HookEventName::PreToolUse, + matcher: Some("^Bash$".to_string()), + timeout_sec: 600, + status_message: None, + additional_context_limit: Default::default(), + source_path: source_path.clone().into(), + source: hook_source(), + display_order: 0, + kind: ConfiguredHandlerKind::Command { + command: "echo hello".to_string(), + r#async: false, + env: std::collections::HashMap::new(), + }, + }] + ); + } + + #[test] + fn session_end_normalizes_timeout() { + let mut handlers = Vec::new(); + let mut hook_entries = Vec::new(); + let mut warnings = Vec::new(); + let mut display_order = 0; + let source_path = source_path(); + let hook_states = std::collections::HashMap::new(); + + append_matcher_groups( + &mut handlers, + &mut hook_entries, + &mut warnings, + &mut display_order, + &mut hook_handler_source(&source_path, &hook_states), + HookEventName::SessionEnd, + vec![MatcherGroup { + matcher: Some("other".to_string()), + hooks: vec![ + HookHandlerConfig::Command { + command: "echo default".to_string(), + command_windows: None, + timeout_sec: None, + r#async: false, + status_message: None, + additional_context_limit: None, + }, + HookHandlerConfig::Command { + command: "echo clamped".to_string(), + command_windows: None, + timeout_sec: Some(600), + r#async: true, + status_message: None, + additional_context_limit: None, + }, + ], + }], + ); + + assert_eq!( + handlers + .iter() + .map(|handler| handler.timeout_sec) + .collect::>(), + vec![1, 3] + ); + assert!( + handlers + .iter() + .all(ConfiguredHandler::can_apply_control_effects) + ); + assert_eq!( + handlers + .iter() + .map(|handler| handler.matcher.as_deref()) + .collect::>(), + vec![Some("other"), Some("other")] + ); + assert_eq!( + hook_entries + .iter() + .map(|entry| entry.timeout_sec) + .collect::>(), + vec![1, 3] + ); + assert!(hook_entries.iter().all(|entry| matches!( + entry.handler, + HookListEntryHandler::Command { r#async: false, .. } + ))); + assert_eq!( + hook_entries + .iter() + .map(|entry| entry.matcher.as_deref()) + .collect::>(), + vec![Some("other"), Some("other")] + ); + assert_eq!( + warnings, + vec![ + format!( + "clamping SessionEnd hook timeout to 3s in {}", + source_path.display() + ), + format!( + "running async SessionEnd hook synchronously in {}", + source_path.display() + ), + ] + ); + } + + #[test] + fn interrupt_normalizes_timeout_and_supports_async_execution() { + let mut handlers = Vec::new(); + let mut hook_entries = Vec::new(); + let mut warnings = Vec::new(); + let mut display_order = 0; + let source_path = source_path(); + let hook_states = std::collections::HashMap::new(); + + append_matcher_groups( + &mut handlers, + &mut hook_entries, + &mut warnings, + &mut display_order, + &mut hook_handler_source(&source_path, &hook_states), + HookEventName::Interrupt, + vec![MatcherGroup { + matcher: Some("ignored".to_string()), + hooks: vec![HookHandlerConfig::Command { + command: "echo interrupt".to_string(), + command_windows: None, + timeout_sec: Some(600), + r#async: true, + status_message: None, + additional_context_limit: None, + }], + }], + ); + + assert_eq!( + normalize_command_hook( + HookEventName::Interrupt, + /*timeout_sec*/ None, + source_path.as_path(), + &mut Vec::new(), + ), + 1 + ); + assert_eq!( + handlers + .iter() + .map(|handler| ( + handler.timeout_sec, + handler.matcher.as_deref(), + handler.execution_mode() + )) + .collect::>(), + vec![(3, None, HookExecutionMode::Async)] + ); + assert_eq!( + hook_entries + .iter() + .map(|entry| (entry.timeout_sec, entry.matcher.as_deref())) + .collect::>(), + vec![(3, None)] + ); + assert!(hook_entries.iter().all(|entry| matches!( + entry.handler, + HookListEntryHandler::Command { r#async: true, .. } + ))); + assert_eq!( + warnings, + vec![format!( + "clamping Interrupt hook timeout to 3s in {}", + source_path.display() + )] + ); + } + + #[test] + fn bypass_hook_trust_allows_enabled_untrusted_handlers() { + let mut handlers = Vec::new(); + let mut hook_entries = Vec::new(); + let mut warnings = Vec::new(); + let mut display_order = 0; + let source_path = source_path(); + let hook_states = std::collections::HashMap::new(); + + append_matcher_groups( + &mut handlers, + &mut hook_entries, + &mut warnings, + &mut display_order, + &mut unmanaged_hook_handler_source( + &source_path, + &hook_states, + /*bypass_hook_trust*/ true, + ), + HookEventName::PreToolUse, + vec![command_group(Some("Bash"))], + ); + + assert_eq!(warnings, Vec::::new()); + assert_eq!(handlers.len(), 1); + assert_eq!(hook_entries.len(), 1); + assert_eq!(hook_entries[0].trust_status, HookTrustStatus::Untrusted); + assert_eq!(hook_entries[0].enabled, true); + } + + #[test] + fn bypass_hook_trust_respects_disabled_handlers() { + let mut handlers = Vec::new(); + let mut hook_entries = Vec::new(); + let mut warnings = Vec::new(); + let mut display_order = 0; + let source_path = source_path(); + let hook_states = std::collections::HashMap::from([( + format!("{}:pre_tool_use:0:0", source_path.display()), + HookStateToml { + enabled: Some(false), + trusted_hash: None, + }, + )]); + + append_matcher_groups( + &mut handlers, + &mut hook_entries, + &mut warnings, + &mut display_order, + &mut unmanaged_hook_handler_source( + &source_path, + &hook_states, + /*bypass_hook_trust*/ true, + ), + HookEventName::PreToolUse, + vec![command_group(Some("Bash"))], + ); + + assert_eq!(warnings, Vec::::new()); + assert_eq!(handlers, Vec::::new()); + assert_eq!(hook_entries.len(), 1); + assert_eq!(hook_entries[0].trust_status, HookTrustStatus::Untrusted); + assert_eq!(hook_entries[0].enabled, false); + } + + #[test] + fn pre_tool_use_treats_star_matcher_as_match_all() { + let mut handlers = Vec::new(); + let mut warnings = Vec::new(); + let mut display_order = 0; + let source_path = source_path(); + let hook_states = std::collections::HashMap::new(); + + append_matcher_groups( + &mut handlers, + &mut Vec::new(), + &mut warnings, + &mut display_order, + &mut hook_handler_source(&source_path, &hook_states), + HookEventName::PreToolUse, + vec![command_group(Some("*"))], + ); + + assert_eq!(warnings, Vec::::new()); + assert_eq!(handlers.len(), 1); + assert_eq!(handlers[0].matcher.as_deref(), Some("*")); + } + + #[test] + fn post_tool_use_keeps_valid_matcher_during_discovery() { + let mut handlers = Vec::new(); + let mut warnings = Vec::new(); + let mut display_order = 0; + let source_path = source_path(); + let hook_states = std::collections::HashMap::new(); + + append_matcher_groups( + &mut handlers, + &mut Vec::new(), + &mut warnings, + &mut display_order, + &mut hook_handler_source(&source_path, &hook_states), + HookEventName::PostToolUse, + vec![command_group(Some("Edit|Write"))], + ); + + assert_eq!(warnings, Vec::::new()); + assert_eq!(handlers.len(), 1); + assert_eq!(handlers[0].event_name, HookEventName::PostToolUse); + assert_eq!(handlers[0].matcher.as_deref(), Some("Edit|Write")); + } + + #[test] + fn toml_hook_discovery_ignores_malformed_state_entries() { + let layer = ConfigLayerEntry::new( + ConfigLayerSource::User { + file: test_path_buf("/tmp/config.toml").abs(), + profile: None, + }, + config_with_malformed_state_and_session_start_hook(), + ); + let mut warnings = Vec::new(); + + let (_, hooks) = super::load_toml_hooks_from_layer(&layer, &mut warnings) + .expect("valid hook events should still load"); + + assert_eq!(warnings, Vec::::new()); + assert_eq!( + hooks, + HookEventsToml { + session_start: vec![MatcherGroup { + matcher: None, + hooks: vec![HookHandlerConfig::Command { + command: "echo hello".to_string(), + command_windows: None, + timeout_sec: None, + r#async: false, + status_message: None, + additional_context_limit: None, + }], + }], + ..Default::default() + } + ); + } + + #[test] + fn pre_tool_use_resolves_windows_command_override_during_discovery() { + let mut handlers = Vec::new(); + let mut warnings = Vec::new(); + let mut display_order = 0; + let source_path = source_path(); + let hook_states = std::collections::HashMap::new(); + + append_matcher_groups( + &mut handlers, + &mut Vec::new(), + &mut warnings, + &mut display_order, + &mut hook_handler_source(&source_path, &hook_states), + HookEventName::PreToolUse, + vec![MatcherGroup { + matcher: Some("^Bash$".to_string()), + hooks: vec![HookHandlerConfig::Command { + command: "echo unix".to_string(), + command_windows: Some("echo windows".to_string()), + timeout_sec: None, + r#async: false, + status_message: None, + additional_context_limit: None, + }], + }], + ); + + assert_eq!(warnings, Vec::::new()); + assert_eq!(handlers.len(), 1); + assert_eq!( + handlers[0].kind, + ConfiguredHandlerKind::Command { + command: if cfg!(windows) { + "echo windows" + } else { + "echo unix" + } + .to_string(), + env: std::collections::HashMap::new(), + r#async: false, + } + ); + } + + fn config_with_malformed_state_and_session_start_hook() -> TomlValue { + serde_json::from_value(serde_json::json!({ + "hooks": { + "state": { + "some_key": { + "enabled": "not a bool", + }, + }, + "SessionStart": [{ + "hooks": [{ + "type": "command", + "command": "echo hello", + }], + }], + }, + })) + .expect("config TOML should deserialize") + } + + #[test] + fn hook_metadata_for_config_layer_source_discards_source_details() { + let config_file = test_path_buf("/tmp/.codex/config.toml").abs(); + let dot_codex_folder = test_path_buf("/tmp/worktree/.codex").abs(); + + assert_eq!( + super::hook_metadata_for_config_layer_source(&ConfigLayerSource::System { + file: config_file.clone(), + }), + (HookSource::System, true), + ); + assert_eq!( + super::hook_metadata_for_config_layer_source(&ConfigLayerSource::User { + file: config_file.clone(), + profile: None, + }), + (HookSource::User, false), + ); + assert_eq!( + super::hook_metadata_for_config_layer_source(&ConfigLayerSource::Project { + dot_codex_folder + }), + (HookSource::Project, false), + ); + assert_eq!( + super::hook_metadata_for_config_layer_source(&ConfigLayerSource::Mdm { + domain: "com.openai.codex".to_string(), + key: "config".to_string(), + }), + (HookSource::Mdm, true), + ); + assert_eq!( + super::hook_metadata_for_config_layer_source(&ConfigLayerSource::EnterpriseManaged { + id: "cfg_123".to_string(), + name: "Base policy".to_string(), + }), + (HookSource::CloudManagedConfig, true), + ); + assert_eq!( + super::hook_metadata_for_config_layer_source(&ConfigLayerSource::SessionFlags), + (HookSource::SessionFlags, false), + ); + assert_eq!( + super::hook_metadata_for_config_layer_source( + &ConfigLayerSource::LegacyManagedConfigTomlFromFile { file: config_file }, + ), + (HookSource::LegacyManagedConfigFile, true), + ); + assert_eq!( + super::hook_metadata_for_config_layer_source( + &ConfigLayerSource::LegacyManagedConfigTomlFromMdm, + ), + (HookSource::LegacyManagedConfigMdm, true), + ); + } +} diff --git a/codex-rs/hooks/src/engine/dispatcher.rs b/codex-rs/hooks/src/engine/dispatcher.rs new file mode 100644 index 0000000000000000000000000000000000000000..c5ee0d2ce1da03ec07b08f9ba692088142673e99 --- /dev/null +++ b/codex-rs/hooks/src/engine/dispatcher.rs @@ -0,0 +1,656 @@ +use std::path::Path; + +use futures::StreamExt; +use futures::stream::FuturesUnordered; +use serde_json::Map; +use serde_json::Value; + +use codex_protocol::protocol::HookCompletedEvent; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookExecutionMode; +use codex_protocol::protocol::HookHandlerType; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookRunSummary; +use codex_protocol::protocol::HookScope; + +use super::ClaudeHooksEngine; +use super::ConfiguredHandler; +use super::ConfiguredHandlerKind; +use super::HandlerRunResult; +use super::HandlerSourcePath; +use super::command_runner::run_command; +use super::mcp_runner::run_mcp_tool; +use crate::events::common::matches_matcher; + +#[derive(Debug)] +pub(crate) struct ParsedHandler { + pub completed: HookCompletedEvent, + pub data: T, + pub completion_order: usize, +} + +pub(crate) fn select_handlers( + handlers: &[ConfiguredHandler], + event_name: HookEventName, + matcher_input: Option<&str>, +) -> Vec { + let matcher_inputs = matcher_input.into_iter().collect::>(); + select_handlers_for_matcher_inputs(handlers, event_name, &matcher_inputs) +} + +pub(crate) fn select_handlers_for_matcher_inputs( + handlers: &[ConfiguredHandler], + event_name: HookEventName, + matcher_inputs: &[&str], +) -> Vec { + // Check each configured handler once, even when several compatibility names + // match the same regex. A hook like `apply_patch|Write|Edit` should run a + // single time for one tool call, not once per matching alias. + handlers + .iter() + .filter(|handler| handler.event_name == event_name) + .filter(|handler| match event_name { + HookEventName::PreToolUse + | HookEventName::PermissionRequest + | HookEventName::PostToolUse + | HookEventName::SessionStart + | HookEventName::SessionEnd + | HookEventName::SubagentStart + | HookEventName::SubagentStop + | HookEventName::PreCompact + | HookEventName::PostCompact => { + if matcher_inputs.is_empty() { + matches_matcher(handler.matcher.as_deref(), /*input*/ None) + } else { + matcher_inputs + .iter() + .any(|input| matches_matcher(handler.matcher.as_deref(), Some(input))) + } + } + HookEventName::UserPromptSubmit | HookEventName::Stop | HookEventName::Interrupt => { + true + } + }) + .cloned() + .collect() +} + +pub(crate) fn running_summary(handler: &ConfiguredHandler) -> HookRunSummary { + let HandlerSourcePath::Local(source_path) = &handler.source_path else { + unreachable!("executor-scoped hooks do not produce public hook summaries"); + }; + HookRunSummary { + builtin: handler.builtin, + id: handler.run_id(), + event_name: handler.event_name, + handler_type: handler.handler_type(), + execution_mode: handler.execution_mode(), + scope: scope_for_event(handler.event_name), + source_path: source_path.clone(), + source: handler.source, + display_order: handler.display_order, + status: HookRunStatus::Running, + status_message: handler.status_message.clone(), + started_at: chrono::Utc::now().timestamp(), + completed_at: None, + duration_ms: None, + entries: Vec::new(), + } +} + +pub(crate) async fn execute_handlers( + engine: &ClaudeHooksEngine, + handlers: Vec, + input_json: String, + cwd: &Path, + turn_id: Option, + parse: fn(&ConfiguredHandler, HandlerRunResult, Option) -> ParsedHandler, +) -> Vec> { + execute_handlers_with_metadata( + engine, handlers, input_json, cwd, turn_id, /*metadata*/ None, parse, + ) + .await +} + +pub(crate) async fn execute_handlers_with_metadata( + engine: &ClaudeHooksEngine, + handlers: Vec, + input_json: String, + cwd: &Path, + turn_id: Option, + metadata: Option<&Map>, + parse: fn(&ConfiguredHandler, HandlerRunResult, Option) -> ParsedHandler, +) -> Vec> { + let mut executor_handlers = Vec::new(); + let mut pending = FuturesUnordered::new(); + for (configured_order, handler) in handlers.into_iter().enumerate() { + if matches!( + handler.source_path, + HandlerSourcePath::ExecutorScoped { .. } + ) { + executor_handlers.push(handler); + continue; + } + if handler.execution_mode() == HookExecutionMode::Async { + engine.command_runtime.schedule_async_hook( + handler, + input_json.clone(), + cwd.to_path_buf(), + turn_id.clone(), + parse, + ); + continue; + } + let input_json = input_json.clone(); + let turn_id = turn_id.clone(); + pending.push(async move { + let result = + execute_handler(engine, &handler, &input_json, cwd, /*metadata*/ None).await; + (configured_order, parse(&handler, result, turn_id)) + }); + } + + let mut completed = Vec::new(); + let mut completion_order = 0; + let mut should_stop = false; + let mut should_block = false; + while let Some((configured_order, mut parsed)) = pending.next().await { + should_stop |= parsed.completed.run.status == HookRunStatus::Stopped; + should_block |= parsed.completed.run.status == HookRunStatus::Blocked; + parsed.completion_order = completion_order; + completion_order += 1; + completed.push((configured_order, parsed)); + } + completed.sort_by_key(|(configured_order, _)| *configured_order); + + if should_stop || !should_block { + for handler in executor_handlers { + let task_engine = engine.clone(); + let input_json = input_json.clone(); + let cwd = cwd.to_path_buf(); + let metadata = metadata.cloned(); + engine.command_runtime.schedule_async_task(async move { + let result = + execute_handler(&task_engine, &handler, &input_json, &cwd, metadata.as_ref()) + .await; + if let Some(error) = result.error { + tracing::warn!( + source_path = %handler.source_path, + %error, + "executor-scoped hook failed" + ); + } + }); + } + } + + completed.into_iter().map(|(_, parsed)| parsed).collect() +} + +async fn execute_handler( + engine: &ClaudeHooksEngine, + handler: &ConfiguredHandler, + input_json: &str, + cwd: &Path, + metadata: Option<&Map>, +) -> HandlerRunResult { + match &handler.kind { + ConfiguredHandlerKind::Command { command, env, .. } => { + run_command( + &engine.command_runtime, + handler, + command, + env, + input_json, + cwd, + ) + .await + } + ConfiguredHandlerKind::McpTool { + server, + tool, + input, + } => { + run_mcp_tool( + engine.mcp_executor.as_ref(), + handler, + server, + tool, + input, + input_json, + metadata, + ) + .await + } + } +} + +pub(crate) fn completed_summary( + handler: &ConfiguredHandler, + run_result: &HandlerRunResult, + status: HookRunStatus, + entries: Vec, +) -> HookRunSummary { + let HandlerSourcePath::Local(source_path) = &handler.source_path else { + unreachable!("executor-scoped hooks do not produce public hook summaries"); + }; + HookRunSummary { + builtin: handler.builtin, + id: handler.run_id(), + event_name: handler.event_name, + handler_type: handler.handler_type(), + execution_mode: handler.execution_mode(), + scope: scope_for_event(handler.event_name), + source_path: source_path.clone(), + source: handler.source, + display_order: handler.display_order, + status, + status_message: handler.status_message.clone(), + started_at: run_result.started_at, + completed_at: Some(run_result.completed_at), + duration_ms: Some(run_result.duration_ms), + entries, + } +} + +pub(crate) fn scope_for_event(event_name: HookEventName) -> HookScope { + match event_name { + HookEventName::SessionStart | HookEventName::SessionEnd | HookEventName::SubagentStart => { + HookScope::Thread + } + HookEventName::PreToolUse + | HookEventName::PermissionRequest + | HookEventName::PostToolUse + | HookEventName::PreCompact + | HookEventName::PostCompact + | HookEventName::UserPromptSubmit + | HookEventName::SubagentStop + | HookEventName::Stop + | HookEventName::Interrupt => HookScope::Turn, + } +} + +pub(crate) fn hook_event_name_label(event_name: HookEventName) -> &'static str { + match event_name { + HookEventName::PreToolUse => "PreToolUse", + HookEventName::PermissionRequest => "PermissionRequest", + HookEventName::PostToolUse => "PostToolUse", + HookEventName::PreCompact => "PreCompact", + HookEventName::PostCompact => "PostCompact", + HookEventName::SessionStart => "SessionStart", + HookEventName::SessionEnd => "SessionEnd", + HookEventName::UserPromptSubmit => "UserPromptSubmit", + HookEventName::SubagentStart => "SubagentStart", + HookEventName::SubagentStop => "SubagentStop", + HookEventName::Stop => "Stop", + HookEventName::Interrupt => "Interrupt", + } +} + +/// Returns the canonical label for a hook execution mode. +pub fn hook_execution_mode_label(mode: HookExecutionMode) -> &'static str { + match mode { + HookExecutionMode::Sync => "sync", + HookExecutionMode::Async => "async", + } +} + +/// Returns the canonical label for a hook handler type. +pub fn hook_handler_type_label(handler_type: HookHandlerType) -> &'static str { + match handler_type { + HookHandlerType::Command => "command", + HookHandlerType::McpTool => "mcp_tool", + HookHandlerType::Prompt => "prompt", + HookHandlerType::Agent => "agent", + } +} + +pub(crate) fn hook_scope_label(scope: HookScope) -> &'static str { + match scope { + HookScope::Thread => "thread", + HookScope::Turn => "turn", + } +} + +pub(crate) fn hook_source_label(source: codex_protocol::protocol::HookSource) -> &'static str { + match source { + codex_protocol::protocol::HookSource::System => "system", + codex_protocol::protocol::HookSource::User => "user", + codex_protocol::protocol::HookSource::Project => "project", + codex_protocol::protocol::HookSource::Mdm => "mdm", + codex_protocol::protocol::HookSource::SessionFlags => "session_flags", + codex_protocol::protocol::HookSource::Plugin => "plugin", + codex_protocol::protocol::HookSource::CloudRequirements => "cloud_requirements", + codex_protocol::protocol::HookSource::CloudManagedConfig => "cloud_managed_config", + codex_protocol::protocol::HookSource::LegacyManagedConfigFile => { + "legacy_managed_config_file" + } + codex_protocol::protocol::HookSource::LegacyManagedConfigMdm => "legacy_managed_config_mdm", + codex_protocol::protocol::HookSource::Unknown => "unknown", + } +} + +#[cfg(test)] +mod tests { + use codex_protocol::protocol::HookEventName; + use codex_protocol::protocol::HookSource; + use codex_utils_absolute_path::test_support::PathBufExt; + use codex_utils_absolute_path::test_support::test_path_buf; + use pretty_assertions::assert_eq; + + use super::ConfiguredHandler; + use super::ConfiguredHandlerKind; + use super::select_handlers; + use super::select_handlers_for_matcher_inputs; + + fn make_handler( + event_name: HookEventName, + matcher: Option<&str>, + command: &str, + display_order: i64, + ) -> ConfiguredHandler { + ConfiguredHandler { + builtin: false, + event_name, + matcher: matcher.map(str::to_owned), + timeout_sec: 5, + status_message: None, + additional_context_limit: Default::default(), + source_path: test_path_buf("/tmp/hooks.json").abs().into(), + source: HookSource::User, + display_order, + kind: ConfiguredHandlerKind::Command { + command: command.to_string(), + r#async: false, + env: std::collections::HashMap::new(), + }, + } + } + + #[test] + fn select_handlers_keeps_duplicate_stop_handlers() { + let handlers = vec![ + make_handler( + HookEventName::Stop, + /*matcher*/ None, + "echo same", + /*display_order*/ 0, + ), + make_handler( + HookEventName::Stop, + /*matcher*/ None, + "echo same", + /*display_order*/ 1, + ), + ]; + + let selected = select_handlers(&handlers, HookEventName::Stop, /*matcher_input*/ None); + + assert_eq!(selected.len(), 2); + assert_eq!(selected[0].display_order, 0); + assert_eq!(selected[1].display_order, 1); + } + + #[test] + fn select_handlers_ignores_interrupt_matchers() { + let handlers = vec![ + make_handler( + HookEventName::Interrupt, + Some("^interrupted$"), + "echo first", + /*display_order*/ 0, + ), + make_handler( + HookEventName::Interrupt, + /*matcher*/ None, + "echo second", + /*display_order*/ 1, + ), + ]; + + let selected = select_handlers( + &handlers, + HookEventName::Interrupt, + /*matcher_input*/ None, + ); + + assert_eq!(selected.len(), 2); + assert_eq!(selected[0].display_order, 0); + assert_eq!(selected[1].display_order, 1); + } + + #[test] + fn select_handlers_keeps_overlapping_session_start_matchers() { + let handlers = vec![ + make_handler( + HookEventName::SessionStart, + Some("start.*"), + "echo same", + /*display_order*/ 0, + ), + make_handler( + HookEventName::SessionStart, + Some("^startup$"), + "echo same", + /*display_order*/ 1, + ), + ]; + + let selected = select_handlers(&handlers, HookEventName::SessionStart, Some("startup")); + + assert_eq!(selected.len(), 2); + assert_eq!(selected[0].display_order, 0); + assert_eq!(selected[1].display_order, 1); + } + + #[test] + fn compact_hooks_match_trigger() { + let handlers = vec![ + make_handler( + HookEventName::PreCompact, + Some("manual"), + "echo manual", + /*display_order*/ 0, + ), + make_handler( + HookEventName::PreCompact, + Some("auto"), + "echo auto", + /*display_order*/ 1, + ), + ]; + + let selected = select_handlers(&handlers, HookEventName::PreCompact, Some("manual")); + + assert_eq!(selected.len(), 1); + assert_eq!(selected[0].display_order, 0); + } + + #[test] + fn pre_tool_use_matches_tool_name() { + let handlers = vec![ + make_handler( + HookEventName::PreToolUse, + Some("^Bash$"), + "echo same", + /*display_order*/ 0, + ), + make_handler( + HookEventName::PreToolUse, + Some("^Edit$"), + "echo same", + /*display_order*/ 1, + ), + ]; + + let selected = select_handlers(&handlers, HookEventName::PreToolUse, Some("Bash")); + + assert_eq!(selected.len(), 1); + assert_eq!(selected[0].display_order, 0); + } + + #[test] + fn post_tool_use_matches_tool_name() { + let handlers = vec![ + make_handler( + HookEventName::PostToolUse, + Some("^Bash$"), + "echo same", + /*display_order*/ 0, + ), + make_handler( + HookEventName::PostToolUse, + Some("^Edit$"), + "echo same", + /*display_order*/ 1, + ), + ]; + + let selected = select_handlers(&handlers, HookEventName::PostToolUse, Some("Bash")); + + assert_eq!(selected.len(), 1); + assert_eq!(selected[0].display_order, 0); + } + + #[test] + fn pre_tool_use_star_matcher_matches_all_tools() { + let handlers = vec![ + make_handler( + HookEventName::PreToolUse, + Some("*"), + "echo same", + /*display_order*/ 0, + ), + make_handler( + HookEventName::PreToolUse, + Some("^Edit$"), + "echo same", + /*display_order*/ 1, + ), + ]; + + let selected = select_handlers(&handlers, HookEventName::PreToolUse, Some("Bash")); + + assert_eq!(selected.len(), 1); + assert_eq!(selected[0].display_order, 0); + } + + #[test] + fn pre_tool_use_regex_alternation_matches_each_tool_name() { + let handlers = vec![make_handler( + HookEventName::PreToolUse, + Some("Edit|Write"), + "echo same", + /*display_order*/ 0, + )]; + + let selected_edit = select_handlers(&handlers, HookEventName::PreToolUse, Some("Edit")); + let selected_write = select_handlers(&handlers, HookEventName::PreToolUse, Some("Write")); + let selected_bash = select_handlers(&handlers, HookEventName::PreToolUse, Some("Bash")); + + assert_eq!(selected_edit.len(), 1); + assert_eq!(selected_write.len(), 1); + assert_eq!(selected_bash.len(), 0); + } + + #[test] + fn pre_tool_use_aliases_match_once_per_handler() { + let handlers = vec![ + make_handler( + HookEventName::PreToolUse, + Some("^apply_patch$"), + "echo apply_patch", + /*display_order*/ 0, + ), + make_handler( + HookEventName::PreToolUse, + Some("^Write$"), + "echo write", + /*display_order*/ 1, + ), + make_handler( + HookEventName::PreToolUse, + Some("^Edit$"), + "echo edit", + /*display_order*/ 2, + ), + make_handler( + HookEventName::PreToolUse, + Some("apply_patch|Write|Edit"), + "echo combined", + /*display_order*/ 3, + ), + ]; + + let selected = select_handlers_for_matcher_inputs( + &handlers, + HookEventName::PreToolUse, + &["apply_patch", "Write", "Edit"], + ); + + assert_eq!(selected.len(), 4); + assert_eq!( + selected + .iter() + .map(|handler| handler.display_order) + .collect::>(), + vec![0, 1, 2, 3], + ); + } + + #[test] + fn user_prompt_submit_ignores_matcher() { + let handlers = vec![ + make_handler( + HookEventName::UserPromptSubmit, + Some("^hello"), + "echo first", + /*display_order*/ 0, + ), + make_handler( + HookEventName::UserPromptSubmit, + Some("["), + "echo second", + /*display_order*/ 1, + ), + ]; + + let selected = select_handlers( + &handlers, + HookEventName::UserPromptSubmit, + /*matcher_input*/ None, + ); + + assert_eq!(selected.len(), 2); + assert_eq!(selected[0].display_order, 0); + assert_eq!(selected[1].display_order, 1); + } + + #[test] + fn select_handlers_preserves_declaration_order() { + let handlers = vec![ + make_handler( + HookEventName::Stop, + /*matcher*/ None, + "first", + /*display_order*/ 0, + ), + make_handler( + HookEventName::Stop, + /*matcher*/ None, + "second", + /*display_order*/ 1, + ), + make_handler( + HookEventName::Stop, + /*matcher*/ None, + "third", + /*display_order*/ 2, + ), + ]; + + let selected = select_handlers(&handlers, HookEventName::Stop, /*matcher_input*/ None); + + assert_eq!(selected, handlers); + } +} diff --git a/codex-rs/hooks/src/engine/mcp_runner.rs b/codex-rs/hooks/src/engine/mcp_runner.rs new file mode 100644 index 0000000000000000000000000000000000000000..7a8da53e4d2b49ef66f2d14afb0e9fd45c8eeb4a --- /dev/null +++ b/codex-rs/hooks/src/engine/mcp_runner.rs @@ -0,0 +1,166 @@ +use std::time::Duration; +use std::time::Instant; + +use anyhow::Context; +use anyhow::Result; +use regex::Regex; +use serde_json::Map; +use serde_json::Value; + +use super::ConfiguredHandler; +use super::HandlerRunResult; +use super::HandlerSourcePath; +use crate::mcp::HookMcpCall; +use crate::mcp::HookMcpExecutor; + +/// Expands an MCP argument template against this hook event and invokes the configured tool. +#[tracing::instrument( + name = "codex.hooks.mcp_tool", + level = "trace", + skip_all, + fields( + hook.server = server, + hook.tool = tool, + hook.timeout_sec = handler.timeout_sec, + ) +)] +pub(crate) async fn run_mcp_tool( + executor: &dyn HookMcpExecutor, + handler: &ConfiguredHandler, + server: &str, + tool: &str, + argument_template: &Map, + hook_event_json: &str, + metadata: Option<&Map>, +) -> HandlerRunResult { + let started_at = chrono::Utc::now().timestamp(); + let started = Instant::now(); + let result = async { + let hook_event: Value = + serde_json::from_str(hook_event_json).context("failed to parse hook event input")?; + let input = expand_mcp_argument_template(argument_template, &hook_event)?; + let (environment_id, mut call_metadata) = match &handler.source_path { + HandlerSourcePath::Local(_) => (None, None), + HandlerSourcePath::ExecutorScoped { + environment_id, + mcp_environment_id, + mcp_metadata, + .. + } => ( + Some( + mcp_environment_id + .as_ref() + .unwrap_or(environment_id) + .clone(), + ), + mcp_metadata.as_deref().cloned(), + ), + }; + if let Some(metadata) = metadata { + call_metadata + .get_or_insert_with(Map::new) + .extend(metadata.clone()); + } + executor + .execute(HookMcpCall { + server: server.to_string(), + tool: tool.to_string(), + environment_id, + metadata: call_metadata, + input, + timeout: Duration::from_secs(handler.timeout_sec), + }) + .await + } + .await; + + let (exit_code, stdout, error) = match result { + Ok(output) => (Some(0), output, None), + Err(error) => (None, String::new(), Some(error.to_string())), + }; + HandlerRunResult { + started_at, + completed_at: chrono::Utc::now().timestamp(), + duration_ms: started.elapsed().as_millis().try_into().unwrap_or(i64::MAX), + exit_code, + stdout, + stderr: String::new(), + error, + } +} + +/// Recursively substitutes `${field.nested}` placeholders using values from a hook event. +/// +/// A complete placeholder preserves its JSON type; a placeholder embedded in surrounding +/// text is rendered as a string. Missing fields fail the hook instead of passing unresolved +/// arguments to the server. For example, the template `{"count":"${tool_input.count}"}` +/// becomes `{"count":3}` when the event contains `{"tool_input":{"count":3}}`. +pub(crate) fn expand_mcp_argument_template( + argument_template: &Map, + hook_event: &Value, +) -> Result> { + argument_template + .iter() + .map(|(key, value)| Ok((key.clone(), resolve_value(value, hook_event)?))) + .collect() +} + +fn resolve_value(value: &Value, hook_event: &Value) -> Result { + match value { + Value::Object(input) => Ok(Value::Object(expand_mcp_argument_template( + input, hook_event, + )?)), + Value::Array(items) => items + .iter() + .map(|item| resolve_value(item, hook_event)) + .collect::>>() + .map(Value::Array), + Value::String(text) => resolve_string(text, hook_event), + _ => Ok(value.clone()), + } +} + +fn resolve_string(text: &str, hook_event: &Value) -> Result { + let pattern = Regex::new(r"\$\{([^{}]+)\}")?; + let captures = pattern.captures_iter(text).collect::>(); + if captures.is_empty() { + return Ok(Value::String(text.to_string())); + } + + if let [capture] = captures.as_slice() + && let Some(placeholder) = capture.get(0) + && placeholder.start() == 0 + && placeholder.end() == text.len() + { + return resolve_path(hook_event, &capture[1]).cloned(); + } + + let mut resolved = String::new(); + let mut previous_end = 0; + for capture in captures { + let Some(placeholder) = capture.get(0) else { + continue; + }; + resolved.push_str(&text[previous_end..placeholder.start()]); + let value = resolve_path(hook_event, &capture[1])?; + match value { + Value::String(value) => resolved.push_str(value), + _ => resolved.push_str(&serde_json::to_string(value)?), + } + previous_end = placeholder.end(); + } + resolved.push_str(&text[previous_end..]); + Ok(Value::String(resolved)) +} + +fn resolve_path<'a>(hook_event: &'a Value, path: &str) -> Result<&'a Value> { + path.split('.').try_fold(hook_event, |value, field| { + value + .get(field) + .with_context(|| format!("hook input placeholder `${{{path}}}` was not found")) + }) +} + +#[cfg(test)] +#[path = "mcp_runner_tests.rs"] +mod tests; diff --git a/codex-rs/hooks/src/engine/mcp_runner_tests.rs b/codex-rs/hooks/src/engine/mcp_runner_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..a16613ea071d1a58525470cf10d656086ed5dcd4 --- /dev/null +++ b/codex-rs/hooks/src/engine/mcp_runner_tests.rs @@ -0,0 +1,141 @@ +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; + +use futures::FutureExt; +use futures::future::BoxFuture; +use pretty_assertions::assert_eq; +use serde_json::Map; +use serde_json::Value; +use serde_json::json; + +use super::expand_mcp_argument_template; +use super::run_mcp_tool; +use crate::engine::ConfiguredHandler; +use crate::engine::ConfiguredHandlerKind; +use crate::mcp::HookMcpCall; +use crate::mcp::HookMcpExecutor; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookSource; +use codex_utils_absolute_path::test_support::PathBufExt; +use codex_utils_absolute_path::test_support::test_path_buf; + +struct RecordingExecutor { + calls: Arc>>, + output: String, +} + +impl HookMcpExecutor for RecordingExecutor { + fn execute(&self, call: HookMcpCall) -> BoxFuture<'_, anyhow::Result> { + async move { + self.calls.lock().expect("lock calls").push(call); + Ok(self.output.clone()) + } + .boxed() + } +} + +#[test] +fn placeholders_preserve_json_types_and_expand_nested_inputs() { + let event_input = json!({ + "tool_input": { + "file_path": "/tmp/example.rs", + "count": 3, + "optional": null, + "metadata": { "language": "rust" }, + }, + }); + let input = serde_json::from_value::>(json!({ + "path": "${tool_input.file_path}", + "message": "scan ${tool_input.file_path}", + "count": "${tool_input.count}", + "optional": "${tool_input.optional}", + "nested": ["${tool_input.metadata}", { "literal": true }], + })) + .expect("object input"); + + assert_eq!( + Value::Object( + expand_mcp_argument_template(&input, &event_input).expect("expand MCP arguments") + ), + json!({ + "path": "/tmp/example.rs", + "message": "scan /tmp/example.rs", + "count": 3, + "optional": null, + "nested": [{ "language": "rust" }, { "literal": true }], + }) + ); +} + +#[test] +fn missing_placeholder_fails_without_passing_unresolved_input() { + let input = serde_json::from_value::>(json!({ + "path": "${tool_input.missing}", + })) + .expect("object input"); + + let error = expand_mcp_argument_template(&input, &json!({ "tool_input": {} })) + .expect_err("missing placeholder should fail"); + + assert_eq!( + error.to_string(), + "hook input placeholder `${tool_input.missing}` was not found" + ); +} + +#[tokio::test] +async fn mcp_tool_results_use_command_hook_output_contract() { + let configured_input = serde_json::from_value::>(json!({ + "file_path": "${tool_input.file_path}", + })) + .expect("object input"); + let handler = ConfiguredHandler { + builtin: false, + event_name: HookEventName::PostToolUse, + matcher: None, + timeout_sec: 30, + status_message: None, + additional_context_limit: Default::default(), + source_path: test_path_buf("/tmp/hooks.json").abs().into(), + source: HookSource::User, + display_order: 0, + kind: ConfiguredHandlerKind::McpTool { + server: "security".to_string(), + tool: "scan".to_string(), + input: configured_input.clone(), + }, + }; + let calls = Arc::new(Mutex::new(Vec::new())); + let executor = RecordingExecutor { + calls: Arc::clone(&calls), + output: r#"{"decision":"block","reason":"unsafe file"}"#.to_string(), + }; + + let result = run_mcp_tool( + &executor, + &handler, + "security", + "scan", + &configured_input, + r#"{"tool_input":{"file_path":"/tmp/example.rs"}}"#, + /*metadata*/ None, + ) + .await; + + assert_eq!(result.exit_code, Some(0)); + assert_eq!(result.stdout, executor.output); + assert_eq!(result.error, None); + assert_eq!( + *calls.lock().expect("lock calls"), + vec![HookMcpCall { + server: "security".to_string(), + tool: "scan".to_string(), + environment_id: None, + metadata: None, + input: serde_json::from_value(json!({ "file_path": "/tmp/example.rs" })) + .expect("object input"), + timeout: Duration::from_secs(30), + }] + ); +} diff --git a/codex-rs/hooks/src/engine/mod.rs b/codex-rs/hooks/src/engine/mod.rs new file mode 100644 index 0000000000000000000000000000000000000000..deaf29679a8ce5b35af474d4309b52a8a3a72ff7 --- /dev/null +++ b/codex-rs/hooks/src/engine/mod.rs @@ -0,0 +1,492 @@ +pub(crate) mod command_runner; +pub(crate) mod discovery; +pub(crate) mod dispatcher; +pub(crate) mod mcp_runner; +pub(crate) mod output_parser; +pub(crate) mod schema_loader; + +use crate::events::compact::PostCompactRequest; +use crate::events::compact::PreCompactOutcome; +use crate::events::compact::PreCompactRequest; +use crate::events::compact::StatelessHookOutcome; +use crate::events::interrupt::InterruptOutcome; +use crate::events::interrupt::InterruptRequest; +use crate::events::permission_request::PermissionRequestOutcome; +use crate::events::permission_request::PermissionRequestRequest; +use crate::events::post_tool_use::PostToolUseOutcome; +use crate::events::post_tool_use::PostToolUseRequest; +use crate::events::pre_tool_use::PreToolUseOutcome; +use crate::events::pre_tool_use::PreToolUseRequest; +use crate::events::session_end::SessionEndOutcome; +use crate::events::session_end::SessionEndRequest; +use crate::events::session_start::SessionStartOutcome; +use crate::events::session_start::SessionStartRequest; +use crate::events::stop::StopOutcome; +use crate::events::stop::StopRequest; +use crate::events::user_prompt_submit::UserPromptSubmitOutcome; +use crate::events::user_prompt_submit::UserPromptSubmitRequest; +use crate::mcp::HookMcpExecutor; +use crate::output_spill::AdditionalContextLimit; +use codex_config::ConfigLayerStack; +use codex_config::HookHandlerConfig; +use codex_plugin::ExecutorPluginHookSource; +use codex_plugin::PluginHookSource; +use codex_plugin::PluginId; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookExecutionMode; +use codex_protocol::protocol::HookHandlerType; +use codex_protocol::protocol::HookRunSummary; +use codex_protocol::protocol::HookSource; +use codex_protocol::protocol::HookTrustStatus; +use codex_utils_absolute_path::AbsolutePathBuf; +use codex_utils_path_uri::PathUri; +use serde_json::Map; +use serde_json::Value; +use std::collections::HashMap; +use std::collections::HashSet; +use std::sync::Arc; +use std::time::Duration; + +use command_runner::CommandHookRuntime; + +#[derive(Debug, Clone)] +pub(crate) struct CommandShell { + pub program: String, + pub args: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct ConfiguredHandler { + /// Internally admitted cleanup hook, enabled independently of per-hook state. + pub builtin: bool, + pub event_name: codex_protocol::protocol::HookEventName, + pub matcher: Option, + pub timeout_sec: u64, + pub status_message: Option, + pub additional_context_limit: AdditionalContextLimit, + pub source_path: HandlerSourcePath, + pub source: HookSource, + pub display_order: i64, + pub kind: ConfiguredHandlerKind, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum HandlerSourcePath { + Local(AbsolutePathBuf), + /// Executor-scoped handlers are currently excluded from user-visible hook reporting + /// (events, summary, telemetry). Their handlers are always executed async. + /// + /// TODO: With CCA, all hooks will be executor-scoped, so user visibility + /// (participation in lifecycle events and summaries) and execution behavior + /// (non-blocking) will need to be determined independently. + ExecutorScoped { + plugin_id: PluginId, + environment_id: String, + mcp_environment_id: Option, + mcp_metadata: Option>>, + manifest_path: PathUri, + source_relative_path: String, + }, +} + +impl From for HandlerSourcePath { + fn from(path: AbsolutePathBuf) -> Self { + Self::Local(path) + } +} + +impl std::fmt::Display for HandlerSourcePath { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Local(path) => write!(formatter, "{}", path.display()), + Self::ExecutorScoped { + environment_id, + manifest_path, + source_relative_path, + .. + } => write!( + formatter, + "{environment_id}:{manifest_path}:{source_relative_path}" + ), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum ConfiguredHandlerKind { + Command { + command: String, + env: HashMap, + r#async: bool, + }, + McpTool { + server: String, + tool: String, + input: serde_json::Map, + }, +} + +#[derive(Debug)] +pub(crate) struct HandlerRunResult { + pub started_at: i64, + pub completed_at: i64, + pub duration_ms: i64, + pub exit_code: Option, + pub stdout: String, + pub stderr: String, + pub error: Option, +} + +impl ConfiguredHandler { + pub(crate) fn execution_mode(&self) -> HookExecutionMode { + if matches!(self.source_path, HandlerSourcePath::ExecutorScoped { .. }) { + return HookExecutionMode::Async; + } + + match self.kind { + ConfiguredHandlerKind::Command { r#async: true, .. } => HookExecutionMode::Async, + ConfiguredHandlerKind::Command { r#async: false, .. } + | ConfiguredHandlerKind::McpTool { .. } => HookExecutionMode::Sync, + } + } + + /// Only synchronous hooks can apply control effects. + pub(crate) fn can_apply_control_effects(&self) -> bool { + self.execution_mode() == HookExecutionMode::Sync + } + + pub fn run_id(&self) -> String { + format!( + "{}:{}:{}", + self.event_name_label(), + self.display_order, + self.source_path + ) + } + + fn event_name_label(&self) -> &'static str { + match self.event_name { + codex_protocol::protocol::HookEventName::PreToolUse => "pre-tool-use", + codex_protocol::protocol::HookEventName::PermissionRequest => "permission-request", + codex_protocol::protocol::HookEventName::PostToolUse => "post-tool-use", + codex_protocol::protocol::HookEventName::PreCompact => "pre-compact", + codex_protocol::protocol::HookEventName::PostCompact => "post-compact", + codex_protocol::protocol::HookEventName::SessionStart => "session-start", + codex_protocol::protocol::HookEventName::SessionEnd => "session-end", + codex_protocol::protocol::HookEventName::UserPromptSubmit => "user-prompt-submit", + codex_protocol::protocol::HookEventName::SubagentStart => "subagent-start", + codex_protocol::protocol::HookEventName::SubagentStop => "subagent-stop", + codex_protocol::protocol::HookEventName::Stop => "stop", + codex_protocol::protocol::HookEventName::Interrupt => "interrupt", + } + } + + fn handler_type(&self) -> HookHandlerType { + match &self.kind { + ConfiguredHandlerKind::Command { .. } => HookHandlerType::Command, + ConfiguredHandlerKind::McpTool { .. } => HookHandlerType::McpTool, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum HookListEntryHandler { + Command { command: String, r#async: bool }, + McpTool { server: String, tool: String }, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct HookListEntry { + /// Builtin hooks remain available internally but are omitted from the public hooks list. + pub builtin: bool, + pub key: String, + pub event_name: HookEventName, + pub handler: HookListEntryHandler, + pub matcher: Option, + pub timeout_sec: u64, + pub status_message: Option, + pub additional_context_limit: Option, + pub source_path: AbsolutePathBuf, + pub source: HookSource, + pub plugin_id: Option, + pub display_order: i64, + pub enabled: bool, + pub is_managed: bool, + pub current_hash: String, + pub trust_status: HookTrustStatus, +} + +#[derive(Clone)] +pub(crate) struct ClaudeHooksEngine { + pub(crate) handlers: Vec, + warnings: Vec, + required_load_errors: Vec, + pub(crate) command_runtime: CommandHookRuntime, + pub(crate) mcp_executor: Arc, +} + +impl ClaudeHooksEngine { + pub(crate) fn new( + enabled: bool, + bypass_hook_trust: bool, + config_layer_stack: Option<&ConfigLayerStack>, + plugin_hook_sources: Vec, + plugin_hook_load_warnings: Vec, + command_runtime: CommandHookRuntime, + mcp_executor: Arc, + ) -> Self { + if !enabled && plugin_hook_sources.is_empty() { + return Self { + handlers: Vec::new(), + warnings: Vec::new(), + required_load_errors: Vec::new(), + command_runtime, + mcp_executor, + }; + } + + let _ = schema_loader::generated_hook_schemas(); + let mut discovered = discovery::discover_handlers( + config_layer_stack, + plugin_hook_sources, + plugin_hook_load_warnings, + bypass_hook_trust, + ); + if !enabled { + discovered.handlers.retain(|handler| handler.builtin); + // Disabled ordinary hooks must not emit warnings or reject session startup. + discovered.warnings.clear(); + discovered.required_load_errors.clear(); + } + Self { + handlers: discovered.handlers, + warnings: discovered.warnings, + required_load_errors: discovered.required_load_errors, + command_runtime, + mcp_executor, + } + } + + pub(crate) fn warnings(&self) -> &[String] { + &self.warnings + } + + pub(crate) fn set_executor_hooks(&mut self, executor_hooks: Vec) { + self.handlers.retain(|handler| { + !matches!( + handler.source_path, + HandlerSourcePath::ExecutorScoped { .. } + ) + }); + + let mut display_order = self + .handlers + .iter() + .map(|handler| handler.display_order) + .max() + .map_or(0, |display_order| display_order.saturating_add(1)); + let mut seen_targets = HashSet::new(); + for source in executor_hooks { + for (event_name, groups) in source.hooks.into_matcher_groups() { + let Some(handler) = groups.into_iter().flat_map(|group| group.hooks).next() else { + continue; + }; + let HookHandlerConfig::McpTool { + server, + tool, + input, + timeout_sec, + status_message, + } = handler + else { + unreachable!("allowlisted executor handler must be an MCP tool"); + }; + + // Bundled plugins can share a cleanup target; run it once per event + // and MCP environment. + let target = ( + std::mem::discriminant(&event_name), + source + .mcp_environment_id + .as_ref() + .unwrap_or(&source.environment_id) + .clone(), + server.clone(), + tool.clone(), + ); + if !seen_targets.insert(target) { + continue; + } + self.handlers.push(ConfiguredHandler { + builtin: true, + event_name, + matcher: None, + timeout_sec: timeout_sec.unwrap_or(5).max(1), + status_message, + additional_context_limit: Default::default(), + source_path: HandlerSourcePath::ExecutorScoped { + plugin_id: source.plugin_id.clone(), + environment_id: source.environment_id.clone(), + mcp_environment_id: source.mcp_environment_id.clone(), + mcp_metadata: source.mcp_metadata.clone().map(Box::new), + manifest_path: source.manifest_path.clone(), + source_relative_path: source.source_relative_path.clone(), + }, + source: HookSource::Plugin, + display_order, + kind: ConfiguredHandlerKind::McpTool { + server, + tool, + input, + }, + }); + display_order = display_order.saturating_add(1); + } + } + } + + pub(crate) fn required_load_errors(&self) -> &[String] { + &self.required_load_errors + } + + pub(crate) fn preview_session_start( + &self, + request: &SessionStartRequest, + ) -> Vec { + crate::events::session_start::preview(&self.handlers, request) + } + + pub(crate) fn preview_pre_tool_use(&self, request: &PreToolUseRequest) -> Vec { + crate::events::pre_tool_use::preview(&self.handlers, request) + } + + pub(crate) fn preview_permission_request( + &self, + request: &PermissionRequestRequest, + ) -> Vec { + crate::events::permission_request::preview(&self.handlers, request) + } + + pub(crate) fn max_permission_request_timeout(&self) -> Duration { + Duration::from_secs( + self.handlers + .iter() + .filter(|handler| { + handler.event_name == HookEventName::PermissionRequest + && handler.can_apply_control_effects() + }) + .map(|handler| handler.timeout_sec) + .max() + .unwrap_or_default(), + ) + } + + pub(crate) fn preview_post_tool_use( + &self, + request: &PostToolUseRequest, + ) -> Vec { + crate::events::post_tool_use::preview(&self.handlers, request) + } + + pub(crate) async fn run_session_start( + &self, + request: SessionStartRequest, + turn_id: Option, + ) -> SessionStartOutcome { + crate::events::session_start::run(self, request, turn_id).await + } + + pub(crate) async fn run_pre_tool_use(&self, request: PreToolUseRequest) -> PreToolUseOutcome { + crate::events::pre_tool_use::run(self, request).await + } + + pub(crate) async fn run_permission_request( + &self, + request: PermissionRequestRequest, + ) -> PermissionRequestOutcome { + crate::events::permission_request::run(self, request).await + } + + pub(crate) async fn run_post_tool_use( + &self, + request: PostToolUseRequest, + ) -> PostToolUseOutcome { + let mut outcome = crate::events::post_tool_use::run(self, request).await; + if let Some(feedback_message) = outcome.feedback_message.take() { + outcome.feedback_message = Some( + self.command_runtime + .output_spiller() + .maybe_spill_text(feedback_message) + .await, + ); + } + outcome + } + + pub(crate) fn preview_pre_compact(&self, request: &PreCompactRequest) -> Vec { + crate::events::compact::preview_pre(&self.handlers, request) + } + + pub(crate) async fn run_pre_compact(&self, request: PreCompactRequest) -> PreCompactOutcome { + crate::events::compact::run_pre(self, request).await + } + + pub(crate) fn preview_post_compact(&self, request: &PostCompactRequest) -> Vec { + crate::events::compact::preview_post(&self.handlers, request) + } + + pub(crate) async fn run_post_compact( + &self, + request: PostCompactRequest, + ) -> StatelessHookOutcome { + crate::events::compact::run_post(self, request).await + } + + pub(crate) fn preview_user_prompt_submit( + &self, + request: &UserPromptSubmitRequest, + ) -> Vec { + crate::events::user_prompt_submit::preview(&self.handlers, request) + } + + pub(crate) async fn run_user_prompt_submit( + &self, + request: UserPromptSubmitRequest, + ) -> UserPromptSubmitOutcome { + crate::events::user_prompt_submit::run(self, request).await + } + + pub(crate) fn preview_stop(&self, request: &StopRequest) -> Vec { + crate::events::stop::preview(&self.handlers, request) + } + + pub(crate) fn preview_session_end(&self) -> Vec { + crate::events::session_end::preview(&self.handlers) + } + + pub(crate) async fn run_session_end(&self, request: SessionEndRequest) -> SessionEndOutcome { + crate::events::session_end::run(self, request).await + } + + pub(crate) async fn run_stop(&self, request: StopRequest) -> StopOutcome { + let mut outcome = crate::events::stop::run(self, request).await; + outcome.continuation_fragments = self + .command_runtime + .output_spiller() + .maybe_spill_prompt_fragments(outcome.continuation_fragments) + .await; + outcome + } + + pub(crate) fn preview_interrupt(&self) -> Vec { + crate::events::interrupt::preview(&self.handlers) + } + + pub(crate) async fn run_interrupt(&self, request: InterruptRequest) -> InterruptOutcome { + crate::events::interrupt::run(self, request).await + } +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; diff --git a/codex-rs/hooks/src/engine/mod_tests.rs b/codex-rs/hooks/src/engine/mod_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..a185e45029ffa6f88c647e862d6bae396661fb22 --- /dev/null +++ b/codex-rs/hooks/src/engine/mod_tests.rs @@ -0,0 +1,2889 @@ +use std::collections::HashMap; +use std::fs; +use std::path::Path; +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; + +use codex_config::AbsolutePathBuf; +use codex_config::ConfigLayerEntry; +use codex_config::ConfigLayerSource; +use codex_config::ConfigLayerStack; +use codex_config::ConfigRequirements; +use codex_config::ConfigRequirementsToml; +use codex_config::Constrained; +use codex_config::ConstrainedWithSource; +use codex_config::HookEventsToml; +use codex_config::HookHandlerConfig; +use codex_config::ManagedHooksRequirementsToml; +use codex_config::MatcherGroup; +use codex_config::RequirementSource; +use codex_config::Sourced; +use codex_config::TomlValue; +use codex_plugin::ExecutorPluginHookSource; +use codex_plugin::PluginHookSource; +use codex_plugin::PluginId; +use codex_protocol::ThreadId; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookExecutionMode; +use codex_protocol::protocol::HookHandlerType; +use codex_protocol::protocol::HookOutputEntry; +use codex_protocol::protocol::HookOutputEntryKind; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookSource; +use codex_protocol::protocol::HookTrustStatus; +use futures::FutureExt; +use futures::future::BoxFuture; +use pretty_assertions::assert_eq; +use tempfile::tempdir; +use tokio::sync::Notify; + +use super::ClaudeHooksEngine; +use super::CommandHookRuntime; +use super::CommandShell; +use super::ConfiguredHandler; +use super::ConfiguredHandlerKind; +use super::HandlerSourcePath; +use super::HookListEntryHandler; +use crate::events::interrupt::InterruptRequest; +use crate::events::pre_tool_use::PreToolUseRequest; +use crate::events::stop::StopHookTarget; +use crate::events::stop::StopRequest; +use crate::mcp::HookMcpCall; +use crate::mcp::HookMcpExecutor; + +fn cwd() -> AbsolutePathBuf { + AbsolutePathBuf::current_dir().expect("current dir") +} + +fn command_runtime(shell: CommandShell) -> CommandHookRuntime { + let (result_sender, _result_receiver) = async_channel::unbounded(); + CommandHookRuntime::new( + shell, + Arc::new(std::env::vars_os().collect()), + ThreadId::new(), + result_sender, + ) +} + +pub(crate) fn mcp_executor() -> Arc { + Arc::new(StaticMcpExecutor { + calls: Arc::new(Mutex::new(Vec::new())), + output: String::new(), + outputs_by_tool: HashMap::new(), + }) +} + +#[test] +fn permission_request_timeout_only_counts_synchronous_handlers() { + let mut engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + /*config_layer_stack*/ None, + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + let command = "echo synchronous permission hook"; + let synchronous_handler = ConfiguredHandler { + builtin: false, + event_name: HookEventName::PermissionRequest, + matcher: None, + timeout_sec: 5, + status_message: None, + additional_context_limit: Default::default(), + source_path: cwd().join("hooks.json").into(), + source: HookSource::User, + display_order: 0, + kind: ConfiguredHandlerKind::Command { + command: command.to_string(), + r#async: false, + env: HashMap::new(), + }, + }; + let asynchronous_handler = ConfiguredHandler { + timeout_sec: 600, + kind: ConfiguredHandlerKind::Command { + command: command.to_string(), + r#async: true, + env: HashMap::new(), + }, + ..synchronous_handler.clone() + }; + + engine.handlers = vec![synchronous_handler, asynchronous_handler.clone()]; + assert_eq!( + engine.max_permission_request_timeout(), + Duration::from_secs(5) + ); + + engine.handlers = vec![asynchronous_handler]; + assert_eq!(engine.max_permission_request_timeout(), Duration::ZERO); +} + +fn managed_hooks_for_current_platform( + managed_dir: impl AsRef, + hooks: HookEventsToml, +) -> ManagedHooksRequirementsToml { + let managed_dir = managed_dir.as_ref().to_path_buf(); + ManagedHooksRequirementsToml { + managed_dir: if cfg!(windows) { + None + } else { + Some(managed_dir.clone()) + }, + windows_managed_dir: if cfg!(windows) { + Some(managed_dir) + } else { + None + }, + hooks, + } +} + +fn pre_tool_use_hook_events(command: impl Into) -> HookEventsToml { + HookEventsToml { + pre_tool_use: vec![MatcherGroup { + matcher: Some("^Bash$".to_string()), + hooks: vec![HookHandlerConfig::Command { + command: command.into(), + command_windows: None, + timeout_sec: Some(10), + r#async: false, + status_message: Some("checking".to_string()), + additional_context_limit: None, + }], + }], + ..Default::default() + } +} + +fn config_toml_with_pre_tool_use(command: &str) -> TomlValue { + let mut config_toml = TomlValue::Table(Default::default()); + let TomlValue::Table(config_table) = &mut config_toml else { + unreachable!("config TOML root should be a table"); + }; + let mut hooks_table = TomlValue::Table(Default::default()); + let TomlValue::Table(hooks_entries) = &mut hooks_table else { + unreachable!("hooks entry should be a table"); + }; + let mut pre_tool_use_group = TomlValue::Table(Default::default()); + let TomlValue::Table(pre_tool_use_group_entries) = &mut pre_tool_use_group else { + unreachable!("PreToolUse group should be a table"); + }; + pre_tool_use_group_entries.insert( + "matcher".to_string(), + TomlValue::String("^Bash$".to_string()), + ); + let mut handler = TomlValue::Table(Default::default()); + let TomlValue::Table(handler_entries) = &mut handler else { + unreachable!("PreToolUse handler should be a table"); + }; + handler_entries.insert("type".to_string(), TomlValue::String("command".to_string())); + handler_entries.insert( + "command".to_string(), + TomlValue::String(command.to_string()), + ); + handler_entries.insert("timeout".to_string(), TomlValue::Integer(10)); + handler_entries.insert( + "statusMessage".to_string(), + TomlValue::String("checking".to_string()), + ); + pre_tool_use_group_entries.insert("hooks".to_string(), TomlValue::Array(vec![handler])); + hooks_entries.insert( + "PreToolUse".to_string(), + TomlValue::Array(vec![pre_tool_use_group]), + ); + config_table.insert("hooks".to_string(), hooks_table); + config_toml +} + +fn requirements_with_managed_hooks_only( + allow_managed_hooks_only: bool, + managed_hooks: Option, +) -> (ConfigRequirements, ConfigRequirementsToml) { + ( + ConfigRequirements { + allow_managed_hooks_only: Some(Sourced::new( + allow_managed_hooks_only, + RequirementSource::LegacyManagedConfigTomlFromMdm, + )), + managed_hooks: managed_hooks.clone().map(|hooks| { + ConstrainedWithSource::new( + Constrained::allow_any(hooks), + Some(RequirementSource::LegacyManagedConfigTomlFromMdm), + ) + }), + ..ConfigRequirements::default() + }, + ConfigRequirementsToml { + allow_managed_hooks_only: Some(allow_managed_hooks_only), + hooks: managed_hooks, + ..ConfigRequirementsToml::default() + }, + ) +} + +fn required_hooks_stack( + managed_hooks: ManagedHooksRequirementsToml, + source: RequirementSource, +) -> ConfigLayerStack { + ConfigLayerStack::new( + Vec::new(), + ConfigRequirements { + managed_hooks: Some(ConstrainedWithSource::new( + Constrained::allow_any(managed_hooks.clone()), + Some(source), + )), + ..ConfigRequirements::default() + }, + ConfigRequirementsToml { + hooks: Some(managed_hooks), + ..ConfigRequirementsToml::default() + }, + ) + .expect("config layer stack") +} + +#[test] +fn required_managed_hooks_allow_disabled_hooks_feature() { + let temp = tempdir().expect("create temp dir"); + let managed_hooks = + managed_hooks_for_current_platform(temp.path(), pre_tool_use_hook_events(" ")); + let config_layer_stack = required_hooks_stack( + managed_hooks, + RequirementSource::LegacyManagedConfigTomlFromMdm, + ); + + for plugin_hook_sources in [ + Vec::new(), + vec![bundled_cleanup_source( + "browser@openai-bundled", + "node_repl", + "Stop", + )], + ] { + let (hooks, _result_receiver) = crate::Hooks::new( + crate::HooksConfig { + feature_enabled: false, + config_layer_stack: Some(config_layer_stack.clone()), + plugin_hook_sources, + ..Default::default() + }, + ThreadId::new(), + mcp_executor(), + ) + .expect("disabled hooks feature should not enforce managed requirements hooks"); + + assert!(hooks.startup_warnings().is_empty()); + } +} + +#[test] +fn required_managed_hooks_reject_invalid_matchers() { + let temp = tempdir().expect("create temp dir"); + let mut events = pre_tool_use_hook_events("echo managed"); + events.pre_tool_use[0].matcher = Some("[".to_string()); + let config_layer_stack = required_hooks_stack( + managed_hooks_for_current_platform(temp.path(), events), + RequirementSource::LegacyManagedConfigTomlFromMdm, + ); + + let error = crate::Hooks::new( + crate::HooksConfig { + feature_enabled: true, + config_layer_stack: Some(config_layer_stack), + ..Default::default() + }, + ThreadId::new(), + mcp_executor(), + ) + .err() + .expect("invalid required managed matcher should reject startup"); + + assert!(error.to_string().contains("invalid matcher")); +} + +#[test] +fn required_managed_hooks_allow_invalid_matchers_without_handlers() { + let temp = tempdir().expect("create temp dir"); + let mut events = pre_tool_use_hook_events("echo managed"); + events.pre_tool_use.push(MatcherGroup { + matcher: Some("[".to_string()), + hooks: Vec::new(), + }); + let config_layer_stack = required_hooks_stack( + managed_hooks_for_current_platform(temp.path(), events), + RequirementSource::LegacyManagedConfigTomlFromMdm, + ); + + let (hooks, _result_receiver) = crate::Hooks::new( + crate::HooksConfig { + feature_enabled: true, + config_layer_stack: Some(config_layer_stack), + ..Default::default() + }, + ThreadId::new(), + mcp_executor(), + ) + .expect("an empty matcher group should not prevent required managed hooks from loading"); + + assert_eq!(hooks.startup_warnings().len(), 1); + assert!(hooks.startup_warnings()[0].contains("invalid matcher")); +} + +#[test] +fn required_managed_hooks_reject_empty_commands() { + let temp = tempdir().expect("create temp dir"); + let config_layer_stack = required_hooks_stack( + managed_hooks_for_current_platform(temp.path(), pre_tool_use_hook_events(" ")), + RequirementSource::LegacyManagedConfigTomlFromMdm, + ); + + let error = crate::Hooks::new( + crate::HooksConfig { + feature_enabled: true, + config_layer_stack: Some(config_layer_stack), + ..Default::default() + }, + ThreadId::new(), + mcp_executor(), + ) + .err() + .expect("empty required managed command should reject startup"); + + assert!(error.to_string().contains("skipping empty hook command")); +} + +#[test] +fn required_managed_hooks_reject_unsupported_handler_types() { + let temp = tempdir().expect("create temp dir"); + let events = HookEventsToml { + pre_tool_use: vec![MatcherGroup { + matcher: Some("^Bash$".to_string()), + hooks: vec![HookHandlerConfig::Prompt {}], + }], + ..Default::default() + }; + let config_layer_stack = required_hooks_stack( + managed_hooks_for_current_platform(temp.path(), events), + RequirementSource::LegacyManagedConfigTomlFromMdm, + ); + + let error = crate::Hooks::new( + crate::HooksConfig { + feature_enabled: true, + config_layer_stack: Some(config_layer_stack), + ..Default::default() + }, + ThreadId::new(), + mcp_executor(), + ) + .err() + .expect("unsupported required managed handler should reject startup"); + + assert!( + error + .to_string() + .contains("prompt hooks are not supported yet") + ); +} + +#[test] +fn required_managed_mcp_hooks_reject_empty_targets() { + let temp = tempdir().expect("create temp dir"); + let events = HookEventsToml { + pre_tool_use: vec![MatcherGroup { + matcher: Some("^Bash$".to_string()), + hooks: vec![HookHandlerConfig::McpTool { + server: "policy".to_string(), + tool: " ".to_string(), + input: Default::default(), + timeout_sec: None, + status_message: None, + }], + }], + ..Default::default() + }; + let config_layer_stack = required_hooks_stack( + managed_hooks_for_current_platform(temp.path(), events), + RequirementSource::LegacyManagedConfigTomlFromMdm, + ); + + let error = crate::Hooks::new( + crate::HooksConfig { + feature_enabled: true, + config_layer_stack: Some(config_layer_stack), + ..Default::default() + }, + ThreadId::new(), + mcp_executor(), + ) + .err() + .expect("invalid required managed MCP hook should reject startup"); + + assert!( + error + .to_string() + .contains("server and tool must not be empty") + ); +} + +#[test] +fn required_managed_session_end_mcp_hooks_reject_startup() { + let temp = tempdir().expect("create temp dir"); + let events = HookEventsToml { + session_end: vec![MatcherGroup { + matcher: None, + hooks: vec![HookHandlerConfig::McpTool { + server: "policy".to_string(), + tool: "check".to_string(), + input: Default::default(), + timeout_sec: None, + status_message: None, + }], + }], + ..Default::default() + }; + let config_layer_stack = required_hooks_stack( + managed_hooks_for_current_platform(temp.path(), events), + RequirementSource::LegacyManagedConfigTomlFromMdm, + ); + + let error = crate::Hooks::new( + crate::HooksConfig { + feature_enabled: true, + config_layer_stack: Some(config_layer_stack), + ..Default::default() + }, + ThreadId::new(), + mcp_executor(), + ) + .err() + .expect("required managed SessionEnd MCP hook should reject startup"); + + assert!( + error + .to_string() + .contains("SessionEnd MCP hooks are not supported") + ); +} + +#[test] +fn required_managed_hooks_with_unknown_source_still_reject_discovery_failures() { + let temp = tempdir().expect("create temp dir"); + let config_layer_stack = required_hooks_stack( + managed_hooks_for_current_platform(temp.path(), pre_tool_use_hook_events("")), + RequirementSource::Unknown, + ); + + let error = crate::Hooks::new( + crate::HooksConfig { + feature_enabled: true, + config_layer_stack: Some(config_layer_stack), + ..Default::default() + }, + ThreadId::new(), + mcp_executor(), + ) + .err() + .expect("unknown-source managed requirements hook should still reject startup"); + + assert!(error.to_string().contains("skipping empty hook command")); +} + +#[test] +fn valid_required_managed_hooks_allow_startup() { + let temp = tempdir().expect("create temp dir"); + let config_layer_stack = required_hooks_stack( + managed_hooks_for_current_platform(temp.path(), pre_tool_use_hook_events("echo managed")), + RequirementSource::LegacyManagedConfigTomlFromMdm, + ); + + let (hooks, _result_receiver) = crate::Hooks::new( + crate::HooksConfig { + feature_enabled: true, + config_layer_stack: Some(config_layer_stack), + ..Default::default() + }, + ThreadId::new(), + mcp_executor(), + ) + .expect("valid managed requirements hook should allow startup"); + + assert!(hooks.startup_warnings().is_empty()); +} + +#[test] +fn managed_config_layer_hook_failures_remain_startup_warnings() { + let temp = tempdir().expect("create temp dir"); + let config_path = + AbsolutePathBuf::try_from(temp.path().join("config.toml")).expect("absolute config path"); + let config_layer_stack = ConfigLayerStack::new( + vec![ConfigLayerEntry::new( + ConfigLayerSource::System { file: config_path }, + config_toml_with_pre_tool_use(" "), + )], + ConfigRequirements::default(), + ConfigRequirementsToml::default(), + ) + .expect("config layer stack"); + + let (hooks, _result_receiver) = crate::Hooks::new( + crate::HooksConfig { + feature_enabled: true, + config_layer_stack: Some(config_layer_stack), + ..Default::default() + }, + ThreadId::new(), + mcp_executor(), + ) + .expect("managed config layer hooks should remain optional"); + + assert_eq!(hooks.startup_warnings().len(), 1); + assert!(hooks.startup_warnings()[0].contains("skipping empty hook command")); +} + +#[tokio::test] +async fn requirements_managed_hooks_execute_from_managed_dir() { + let temp = tempdir().expect("create temp dir"); + let managed_dir = + AbsolutePathBuf::try_from(temp.path().join("managed-hooks")).expect("absolute path"); + fs::create_dir_all(managed_dir.as_path()).expect("create managed hooks dir"); + let script_path = managed_dir.join("pre_tool_use.py"); + let log_path = managed_dir.join("pre_tool_use_log.jsonl"); + fs::write( + script_path.as_path(), + format!( + r#"import json +from pathlib import Path +import sys + +payload = json.load(sys.stdin) +with Path(r"{log_path}").open("a", encoding="utf-8") as handle: + handle.write(json.dumps(payload) + "\n") +"#, + log_path = log_path.display(), + ), + ) + .expect("write managed hook script"); + + let managed_hooks = managed_hooks_for_current_platform( + managed_dir.clone(), + HookEventsToml { + pre_tool_use: vec![MatcherGroup { + matcher: Some("^Bash$".to_string()), + hooks: vec![HookHandlerConfig::Command { + command: format!("python3 {}", script_path.display()), + command_windows: None, + timeout_sec: Some(10), + r#async: false, + status_message: Some("checking".to_string()), + additional_context_limit: None, + }], + }], + ..Default::default() + }, + ); + let config_layer_stack = ConfigLayerStack::new( + Vec::new(), + ConfigRequirements { + managed_hooks: Some(ConstrainedWithSource::new( + Constrained::allow_any(managed_hooks.clone()), + Some(RequirementSource::LegacyManagedConfigTomlFromMdm), + )), + ..ConfigRequirements::default() + }, + ConfigRequirementsToml { + hooks: Some(managed_hooks), + ..ConfigRequirementsToml::default() + }, + ) + .expect("config layer stack"); + + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + assert!(engine.warnings().is_empty()); + assert_eq!(engine.handlers.len(), 1); + assert_eq!( + engine.handlers[0].source, + HookSource::LegacyManagedConfigMdm + ); + let listed = crate::list_hooks(crate::HooksConfig { + legacy_notify_argv: None, + feature_enabled: true, + bypass_hook_trust: false, + config_layer_stack: Some(config_layer_stack.clone()), + plugin_hook_sources: Vec::new(), + plugin_hook_load_warnings: Vec::new(), + shell_program: None, + shell_args: Vec::new(), + }); + assert!(listed.hooks[0].is_managed); + let cwd = cwd(); + let preview = engine.preview_pre_tool_use(&PreToolUseRequest { + session_id: ThreadId::new(), + turn_id: "turn-1".to_string(), + subagent: None, + cwd: cwd.clone(), + transcript_path: None, + model: "gpt-test".to_string(), + permission_mode: "default".to_string(), + tool_name: "Bash".to_string(), + matcher_aliases: Vec::new(), + tool_use_id: "tool-1".to_string(), + tool_input: serde_json::json!({ "command": "echo hello" }), + }); + assert_eq!(preview.len(), 1); + assert_eq!(preview[0].source_path, managed_dir); + + let outcome = engine + .run_pre_tool_use(PreToolUseRequest { + session_id: ThreadId::new(), + turn_id: "turn-1".to_string(), + subagent: None, + cwd, + transcript_path: None, + model: "gpt-test".to_string(), + permission_mode: "default".to_string(), + tool_name: "Bash".to_string(), + matcher_aliases: Vec::new(), + tool_use_id: "tool-1".to_string(), + tool_input: serde_json::json!({ "command": "echo hello" }), + }) + .await; + + assert!(!outcome.should_block); + let log_contents = fs::read_to_string(log_path).expect("read managed hook log"); + assert!(log_contents.contains("\"hook_event_name\": \"PreToolUse\"")); +} + +#[tokio::test] +async fn requirements_managed_hooks_execute_windows_command_override() { + let temp = tempdir().expect("create temp dir"); + let managed_dir = + AbsolutePathBuf::try_from(temp.path().join("managed-hooks")).expect("absolute path"); + fs::create_dir_all(managed_dir.as_path()).expect("create managed hooks dir"); + + let managed_hooks = managed_hooks_for_current_platform( + managed_dir, + HookEventsToml { + pre_tool_use: vec![MatcherGroup { + matcher: Some("^Bash$".to_string()), + hooks: vec![HookHandlerConfig::Command { + command: "exit 17".to_string(), + command_windows: Some("exit /B 19".to_string()), + timeout_sec: Some(10), + r#async: false, + status_message: Some("checking".to_string()), + additional_context_limit: None, + }], + }], + ..Default::default() + }, + ); + let config_layer_stack = ConfigLayerStack::new( + Vec::new(), + ConfigRequirements { + managed_hooks: Some(ConstrainedWithSource::new( + Constrained::allow_any(managed_hooks.clone()), + Some(RequirementSource::LegacyManagedConfigTomlFromMdm), + )), + ..ConfigRequirements::default() + }, + ConfigRequirementsToml { + hooks: Some(managed_hooks), + ..ConfigRequirementsToml::default() + }, + ) + .expect("config layer stack"); + + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + let outcome = engine + .run_pre_tool_use(PreToolUseRequest { + session_id: ThreadId::new(), + turn_id: "turn-1".to_string(), + subagent: None, + cwd: cwd(), + transcript_path: None, + model: "gpt-test".to_string(), + permission_mode: "default".to_string(), + tool_name: "Bash".to_string(), + matcher_aliases: Vec::new(), + tool_use_id: "tool-1".to_string(), + tool_input: serde_json::json!({ "command": "echo hello" }), + }) + .await; + + assert!(!outcome.should_block); + let expected_exit_code = if cfg!(windows) { 19 } else { 17 }; + assert_eq!(outcome.hook_events.len(), 1); + assert_eq!(outcome.hook_events[0].run.status, HookRunStatus::Failed); + assert_eq!( + outcome.hook_events[0].run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: format!("hook exited with code {expected_exit_code}"), + }] + ); +} + +#[test] +fn unknown_requirement_source_hooks_stay_managed() { + let temp = tempdir().expect("create temp dir"); + let managed_dir = + AbsolutePathBuf::try_from(temp.path().join("managed-hooks")).expect("absolute path"); + fs::create_dir_all(managed_dir.as_path()).expect("create managed hooks dir"); + let managed_hooks = managed_hooks_for_current_platform( + managed_dir, + HookEventsToml { + pre_tool_use: vec![MatcherGroup { + matcher: Some("^Bash$".to_string()), + hooks: vec![HookHandlerConfig::Command { + command: "python3 /tmp/managed.py".to_string(), + command_windows: None, + timeout_sec: Some(10), + r#async: false, + status_message: Some("checking".to_string()), + additional_context_limit: None, + }], + }], + ..Default::default() + }, + ); + let config_layer_stack = ConfigLayerStack::new( + Vec::new(), + ConfigRequirements { + managed_hooks: Some(ConstrainedWithSource::new( + Constrained::allow_any(managed_hooks.clone()), + Some(RequirementSource::Unknown), + )), + ..ConfigRequirements::default() + }, + ConfigRequirementsToml { + hooks: Some(managed_hooks), + ..ConfigRequirementsToml::default() + }, + ) + .expect("config layer stack"); + + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + assert_eq!(engine.handlers.len(), 1); + assert_eq!(engine.handlers[0].source, HookSource::Unknown); + let discovered = super::discovery::discover_handlers( + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + /*bypass_hook_trust*/ false, + ); + assert_eq!(discovered.hook_entries.len(), 1); + assert_eq!(discovered.hook_entries[0].source, HookSource::Unknown); + assert_eq!(discovered.hook_entries[0].enabled, true); + assert_eq!(discovered.hook_entries[0].is_managed, true); + assert_eq!( + discovered.hook_entries[0].trust_status, + HookTrustStatus::Managed + ); +} + +#[test] +fn user_disablement_filters_non_managed_hooks_but_not_managed_hooks() { + let temp = tempdir().expect("create temp dir"); + let managed_dir = + AbsolutePathBuf::try_from(temp.path().join("managed-hooks")).expect("absolute path"); + fs::create_dir_all(managed_dir.as_path()).expect("create managed hooks dir"); + let managed_hooks = managed_hooks_for_current_platform( + managed_dir.clone(), + HookEventsToml { + pre_tool_use: vec![MatcherGroup { + matcher: Some("^Bash$".to_string()), + hooks: vec![HookHandlerConfig::Command { + command: "python3 /tmp/managed.py".to_string(), + command_windows: None, + timeout_sec: Some(10), + r#async: false, + status_message: Some("checking".to_string()), + additional_context_limit: None, + }], + }], + ..Default::default() + }, + ); + let config_path = + AbsolutePathBuf::try_from(temp.path().join("config.toml")).expect("absolute path"); + let managed_disabled_key = format!("{}:pre_tool_use:0:0", managed_dir.display()); + let user_disabled_key = format!("{}:pre_tool_use:0:0", config_path.display()); + let user_config = config_with_pre_tool_use_hook_and_states( + "python3 /tmp/user.py", + [&managed_disabled_key, &user_disabled_key], + ); + let config_layer_stack = ConfigLayerStack::new( + vec![ConfigLayerEntry::new( + ConfigLayerSource::User { + file: config_path, + profile: None, + }, + user_config, + )], + ConfigRequirements { + managed_hooks: Some(ConstrainedWithSource::new( + Constrained::allow_any(managed_hooks.clone()), + Some(RequirementSource::LegacyManagedConfigTomlFromMdm), + )), + ..ConfigRequirements::default() + }, + ConfigRequirementsToml { + hooks: Some(managed_hooks), + ..ConfigRequirementsToml::default() + }, + ) + .expect("config layer stack"); + + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + assert_eq!(engine.handlers.len(), 1); + assert_eq!( + engine.handlers[0].source, + HookSource::LegacyManagedConfigMdm + ); + let discovered = super::discovery::discover_handlers( + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + /*bypass_hook_trust*/ false, + ); + assert_eq!(discovered.hook_entries.len(), 2); + assert_eq!(discovered.hook_entries[0].key, managed_disabled_key); + assert_eq!(discovered.hook_entries[0].enabled, true); + assert!(discovered.hook_entries[0].is_managed); + assert_eq!( + discovered.hook_entries[0].trust_status, + HookTrustStatus::Managed + ); + assert_eq!(discovered.hook_entries[1].key, user_disabled_key); + assert_eq!(discovered.hook_entries[1].enabled, false); + assert!(!discovered.hook_entries[1].is_managed); +} + +#[test] +fn user_disablement_does_not_filter_managed_layer_hooks() { + let temp = tempdir().expect("create temp dir"); + let managed_config_path = + AbsolutePathBuf::try_from(temp.path().join("managed_config.toml")).expect("absolute path"); + let user_config_path = + AbsolutePathBuf::try_from(temp.path().join("config.toml")).expect("absolute path"); + let managed_key = format!("{}:pre_tool_use:0:0", managed_config_path.display()); + + let config_layer_stack = ConfigLayerStack::new( + vec![ + ConfigLayerEntry::new( + ConfigLayerSource::User { + file: user_config_path, + profile: None, + }, + config_with_hook_state(&managed_key, /*enabled*/ false), + ), + ConfigLayerEntry::new( + ConfigLayerSource::LegacyManagedConfigTomlFromFile { + file: managed_config_path, + }, + config_with_pre_tool_use_hook("python3 /tmp/managed-layer.py"), + ), + ], + ConfigRequirements::default(), + ConfigRequirementsToml::default(), + ) + .expect("config layer stack"); + + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + assert_eq!(engine.handlers.len(), 1); + assert_eq!( + engine.handlers[0].source, + HookSource::LegacyManagedConfigFile + ); + let discovered = super::discovery::discover_handlers( + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + /*bypass_hook_trust*/ false, + ); + assert_eq!(discovered.hook_entries.len(), 1); + assert_eq!(discovered.hook_entries[0].key, managed_key); + assert_eq!(discovered.hook_entries[0].enabled, true); + assert!(discovered.hook_entries[0].is_managed); + assert_eq!( + discovered.hook_entries[0].trust_status, + HookTrustStatus::Managed + ); +} + +fn config_with_hook_state(key: &str, enabled: bool) -> TomlValue { + serde_json::from_value(serde_json::json!({ + "hooks": { + "state": { + (key): { + "enabled": enabled, + }, + }, + }, + })) + .expect("config TOML should deserialize") +} + +fn config_with_pre_tool_use_hook_and_states( + command: &str, + disabled_keys: [&str; N], +) -> TomlValue { + let state = disabled_keys + .into_iter() + .map(|key| (key.to_string(), serde_json::json!({ "enabled": false }))) + .collect::>(); + serde_json::from_value(serde_json::json!({ + "hooks": { + "state": state, + "PreToolUse": [{ + "hooks": [{ + "type": "command", + "command": command, + }], + }], + }, + })) + .expect("config TOML should deserialize") +} + +fn config_with_pre_tool_use_hook(command: &str) -> TomlValue { + serde_json::from_value(serde_json::json!({ + "hooks": { + "PreToolUse": [{ + "hooks": [{ + "type": "command", + "command": command, + }], + }], + }, + })) + .expect("config TOML should deserialize") +} + +fn trusted_plugin_hook_stack( + config_path: AbsolutePathBuf, + plugin_hook_sources: &[PluginHookSource], +) -> ConfigLayerStack { + let discovered = super::discovery::discover_handlers( + /*config_layer_stack*/ None, + plugin_hook_sources.to_vec(), + Vec::new(), + /*bypass_hook_trust*/ false, + ); + let state = discovered + .hook_entries + .into_iter() + .map(|entry| { + ( + entry.key, + serde_json::json!({ + "trusted_hash": entry.current_hash, + }), + ) + }) + .collect::>(); + let config = serde_json::from_value(serde_json::json!({ + "hooks": { + "state": state, + }, + })) + .expect("config TOML should deserialize"); + + ConfigLayerStack::new( + vec![ConfigLayerEntry::new( + ConfigLayerSource::User { + file: config_path, + profile: None, + }, + config, + )], + ConfigRequirements::default(), + ConfigRequirementsToml::default(), + ) + .expect("config layer stack") +} + +#[test] +fn requirements_managed_hooks_load_when_managed_dir_is_missing() { + let temp = tempdir().expect("create temp dir"); + let missing_dir = temp.path().join("missing-managed-hooks"); + let managed_hooks = managed_hooks_for_current_platform( + missing_dir.clone(), + HookEventsToml { + pre_tool_use: vec![MatcherGroup { + matcher: Some("^Bash$".to_string()), + hooks: vec![HookHandlerConfig::Command { + command: "echo hi".to_string(), + command_windows: None, + timeout_sec: Some(10), + r#async: false, + status_message: Some("checking".to_string()), + additional_context_limit: None, + }], + }], + ..Default::default() + }, + ); + let config_layer_stack = ConfigLayerStack::new( + Vec::new(), + ConfigRequirements { + managed_hooks: Some(ConstrainedWithSource::new( + Constrained::allow_any(managed_hooks.clone()), + Some(RequirementSource::LegacyManagedConfigTomlFromMdm), + )), + ..ConfigRequirements::default() + }, + ConfigRequirementsToml { + hooks: Some(managed_hooks), + ..ConfigRequirementsToml::default() + }, + ) + .expect("config layer stack"); + + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + assert!(engine.warnings().is_empty()); + let cwd = cwd(); + let preview = engine.preview_pre_tool_use(&PreToolUseRequest { + session_id: ThreadId::new(), + turn_id: "turn-1".to_string(), + subagent: None, + cwd, + transcript_path: None, + model: "gpt-test".to_string(), + permission_mode: "default".to_string(), + tool_name: "Bash".to_string(), + matcher_aliases: Vec::new(), + tool_use_id: "tool-1".to_string(), + tool_input: serde_json::json!({ "command": "echo hello" }), + }); + assert_eq!(preview.len(), 1); + assert_eq!( + engine.handlers[0].kind, + ConfiguredHandlerKind::Command { + command: "echo hi".to_string(), + r#async: false, + env: HashMap::new(), + } + ); + assert_eq!( + engine.handlers[0].source_path, + AbsolutePathBuf::try_from(missing_dir) + .expect("absolute missing dir") + .into() + ); +} + +#[test] +fn allow_managed_hooks_only_false_keeps_unmanaged_hooks() { + let temp = tempdir().expect("create temp dir"); + let config_path = + AbsolutePathBuf::try_from(temp.path().join("config.toml")).expect("absolute config path"); + let (requirements, requirements_toml) = requirements_with_managed_hooks_only( + /*allow_managed_hooks_only*/ false, /*managed_hooks*/ None, + ); + let config_layer_stack = ConfigLayerStack::new( + vec![ConfigLayerEntry::new( + ConfigLayerSource::User { + file: config_path, + profile: None, + }, + config_toml_with_pre_tool_use("python3 /tmp/user-hook.py"), + )], + requirements, + requirements_toml, + ) + .expect("config layer stack"); + + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + assert!(engine.warnings().is_empty()); + assert!(engine.handlers.is_empty()); + let discovered = super::discovery::discover_handlers( + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + /*bypass_hook_trust*/ false, + ); + assert_eq!(discovered.hook_entries.len(), 1); + assert!(!discovered.hook_entries[0].is_managed); + assert_eq!( + discovered.hook_entries[0].handler, + HookListEntryHandler::Command { + command: "python3 /tmp/user-hook.py".to_string(), + r#async: false, + } + ); +} + +#[test] +fn allow_managed_hooks_only_in_config_toml_does_not_enable_policy() { + let temp = tempdir().expect("create temp dir"); + let config_path = + AbsolutePathBuf::try_from(temp.path().join("config.toml")).expect("absolute config path"); + let mut config_toml = config_toml_with_pre_tool_use("python3 /tmp/user-hook.py"); + let TomlValue::Table(config_table) = &mut config_toml else { + unreachable!("config TOML root should be a table"); + }; + config_table.insert( + "allow_managed_hooks_only".to_string(), + TomlValue::Boolean(true), + ); + let config_layer_stack = ConfigLayerStack::new( + vec![ConfigLayerEntry::new( + ConfigLayerSource::User { + file: config_path, + profile: None, + }, + config_toml, + )], + ConfigRequirements::default(), + ConfigRequirementsToml::default(), + ) + .expect("config layer stack"); + + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + assert!(engine.warnings().is_empty()); + assert!(engine.handlers.is_empty()); + let discovered = super::discovery::discover_handlers( + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + /*bypass_hook_trust*/ false, + ); + assert_eq!(discovered.hook_entries.len(), 1); + assert!(!discovered.hook_entries[0].is_managed); + assert_eq!( + discovered.hook_entries[0].handler, + HookListEntryHandler::Command { + command: "python3 /tmp/user-hook.py".to_string(), + r#async: false, + } + ); +} + +#[test] +fn allow_managed_hooks_only_skips_unmanaged_json_and_toml_hooks() { + let temp = tempdir().expect("create temp dir"); + let config_path = + AbsolutePathBuf::try_from(temp.path().join("config.toml")).expect("absolute config path"); + let hooks_json_path = + AbsolutePathBuf::try_from(temp.path().join("hooks.json")).expect("absolute hooks path"); + fs::write( + hooks_json_path.as_path(), + r#"{ + "hooks": { + "PreToolUse": [ + { + "matcher": "^Bash$", + "hooks": [ + { + "type": "command", + "command": "python3 /tmp/json-hook.py" + } + ] + } + ] + } + }"#, + ) + .expect("write hooks.json"); + let (requirements, requirements_toml) = requirements_with_managed_hooks_only( + /*allow_managed_hooks_only*/ true, /*managed_hooks*/ None, + ); + let config_layer_stack = ConfigLayerStack::new( + vec![ConfigLayerEntry::new( + ConfigLayerSource::User { + file: config_path, + profile: None, + }, + config_toml_with_pre_tool_use("python3 /tmp/toml-hook.py"), + )], + requirements, + requirements_toml, + ) + .expect("config layer stack"); + + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + assert!(engine.handlers.is_empty()); + assert!(engine.warnings().is_empty()); +} + +#[test] +fn allow_managed_hooks_only_skips_unmanaged_plugin_hooks() { + let temp = tempdir().expect("create temp dir"); + let plugin_root = + AbsolutePathBuf::try_from(temp.path().join("demo-plugin")).expect("plugin root"); + let plugin_data_root = + AbsolutePathBuf::try_from(temp.path().join("plugin-data")).expect("plugin data root"); + let source_path = plugin_root.join("hooks/hooks.json"); + let plugin_id = PluginId::parse("demo-plugin@test-marketplace").expect("plugin id"); + let plugin_hook_sources = vec![PluginHookSource { + plugin_id, + plugin_root, + plugin_data_root, + source_path, + source_relative_path: "hooks/hooks.json".to_string(), + hooks: pre_tool_use_hook_events("python3 /tmp/plugin-hook.py"), + }]; + let (requirements, requirements_toml) = requirements_with_managed_hooks_only( + /*allow_managed_hooks_only*/ true, /*managed_hooks*/ None, + ); + let config_layer_stack = ConfigLayerStack::new(Vec::new(), requirements, requirements_toml) + .expect("config layer stack"); + + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + Some(&config_layer_stack), + plugin_hook_sources, + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + assert!(engine.handlers.is_empty()); + assert!(engine.warnings().is_empty()); +} + +#[test] +fn allow_managed_hooks_only_keeps_managed_requirement_and_config_layer_hooks() { + let temp = tempdir().expect("create temp dir"); + let managed_dir = + AbsolutePathBuf::try_from(temp.path().join("managed-hooks")).expect("absolute path"); + fs::create_dir_all(managed_dir.as_path()).expect("create managed hooks dir"); + let system_config_path = + AbsolutePathBuf::try_from(temp.path().join("system").join("config.toml")) + .expect("absolute system config path"); + let system_parent = system_config_path + .as_path() + .parent() + .expect("system config parent"); + fs::create_dir_all(system_parent).expect("create system config dir"); + let legacy_config_path = AbsolutePathBuf::try_from(temp.path().join("managed_config.toml")) + .expect("absolute legacy config path"); + + let managed_hooks = managed_hooks_for_current_platform( + managed_dir, + pre_tool_use_hook_events("python3 /tmp/requirements-hook.py"), + ); + let (requirements, requirements_toml) = requirements_with_managed_hooks_only( + /*allow_managed_hooks_only*/ true, + Some(managed_hooks), + ); + let config_layer_stack = ConfigLayerStack::new( + vec![ + ConfigLayerEntry::new( + ConfigLayerSource::Mdm { + domain: "com.openai.codex".to_string(), + key: "config".to_string(), + }, + config_toml_with_pre_tool_use("python3 /tmp/mdm-hook.py"), + ), + ConfigLayerEntry::new( + ConfigLayerSource::System { + file: system_config_path, + }, + config_toml_with_pre_tool_use("python3 /tmp/system-hook.py"), + ), + ConfigLayerEntry::new( + ConfigLayerSource::LegacyManagedConfigTomlFromFile { + file: legacy_config_path, + }, + config_toml_with_pre_tool_use("python3 /tmp/legacy-file-hook.py"), + ), + ConfigLayerEntry::new( + ConfigLayerSource::LegacyManagedConfigTomlFromMdm, + config_toml_with_pre_tool_use("python3 /tmp/legacy-mdm-hook.py"), + ), + ], + requirements, + requirements_toml, + ) + .expect("config layer stack"); + + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + assert!(engine.warnings().is_empty()); + assert_eq!( + engine + .handlers + .iter() + .map(|handler| match &handler.kind { + ConfiguredHandlerKind::Command { command, .. } => Some(command.as_str()), + ConfiguredHandlerKind::McpTool { .. } => None, + }) + .collect::>(), + vec![ + Some("python3 /tmp/requirements-hook.py"), + Some("python3 /tmp/mdm-hook.py"), + Some("python3 /tmp/system-hook.py"), + Some("python3 /tmp/legacy-file-hook.py"), + Some("python3 /tmp/legacy-mdm-hook.py"), + ] + ); + let discovered = super::discovery::discover_handlers( + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + /*bypass_hook_trust*/ false, + ); + assert!(discovered.hook_entries.iter().all(|entry| entry.is_managed)); +} + +#[test] +fn discovers_hooks_from_json_and_toml_in_the_same_layer() { + let temp = tempdir().expect("create temp dir"); + let config_path = + AbsolutePathBuf::try_from(temp.path().join("config.toml")).expect("absolute config path"); + let hooks_json_path = + AbsolutePathBuf::try_from(temp.path().join("hooks.json")).expect("absolute hooks path"); + fs::write( + hooks_json_path.as_path(), + r#"{ + "hooks": { + "PreToolUse": [ + { + "matcher": "^Bash$", + "hooks": [ + { + "type": "command", + "command": "python3 /tmp/json-hook.py" + } + ] + } + ] + } + }"#, + ) + .expect("write hooks.json"); + let mut config_toml = TomlValue::Table(Default::default()); + let TomlValue::Table(config_table) = &mut config_toml else { + unreachable!("config TOML root should be a table"); + }; + let mut hooks_table = TomlValue::Table(Default::default()); + let TomlValue::Table(hooks_entries) = &mut hooks_table else { + unreachable!("hooks entry should be a table"); + }; + let mut pre_tool_use_group = TomlValue::Table(Default::default()); + let TomlValue::Table(pre_tool_use_group_entries) = &mut pre_tool_use_group else { + unreachable!("PreToolUse group should be a table"); + }; + pre_tool_use_group_entries.insert( + "matcher".to_string(), + TomlValue::String("^Bash$".to_string()), + ); + pre_tool_use_group_entries.insert( + "hooks".to_string(), + TomlValue::Array(vec![TomlValue::Table(Default::default())]), + ); + let Some(TomlValue::Array(hooks_array)) = pre_tool_use_group_entries.get_mut("hooks") else { + unreachable!("PreToolUse hooks should be an array"); + }; + let Some(TomlValue::Table(handler_entries)) = hooks_array.first_mut() else { + unreachable!("PreToolUse handler should be a table"); + }; + handler_entries.insert("type".to_string(), TomlValue::String("command".to_string())); + handler_entries.insert( + "command".to_string(), + TomlValue::String("python3 /tmp/toml-hook.py".to_string()), + ); + hooks_entries.insert( + "PreToolUse".to_string(), + TomlValue::Array(vec![pre_tool_use_group]), + ); + config_table.insert("hooks".to_string(), hooks_table); + let config_layer_stack = ConfigLayerStack::new( + vec![ConfigLayerEntry::new( + ConfigLayerSource::System { + file: config_path.clone(), + }, + config_toml, + )], + ConfigRequirements::default(), + ConfigRequirementsToml::default(), + ) + .expect("config layer stack"); + + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + assert!(engine.warnings().iter().any(|warning| { + warning.contains("loading hooks from both") + && warning.contains(&hooks_json_path.display().to_string()) + && warning.contains(&config_path.display().to_string()) + })); + + let cwd = cwd(); + let preview = engine.preview_pre_tool_use(&PreToolUseRequest { + session_id: ThreadId::new(), + turn_id: "turn-1".to_string(), + subagent: None, + cwd, + transcript_path: None, + model: "gpt-test".to_string(), + permission_mode: "default".to_string(), + tool_name: "Bash".to_string(), + matcher_aliases: Vec::new(), + tool_use_id: "tool-1".to_string(), + tool_input: serde_json::json!({ "command": "echo hello" }), + }); + assert_eq!(preview.len(), 2); + assert_eq!( + engine + .handlers + .iter() + .map(|handler| handler.source) + .collect::>(), + vec![HookSource::System, HookSource::System] + ); + assert_eq!(preview[0].source_path, hooks_json_path); + assert_eq!(preview[1].source_path, config_path); +} + +#[test] +fn profile_user_layers_load_shared_hooks_json_once() { + let temp = tempdir().expect("create temp dir"); + let config_path = + AbsolutePathBuf::try_from(temp.path().join("config.toml")).expect("absolute config path"); + let profile_path = AbsolutePathBuf::try_from(temp.path().join("work.config.toml")) + .expect("absolute profile path"); + let hooks_json_path = + AbsolutePathBuf::try_from(temp.path().join("hooks.json")).expect("absolute hooks path"); + fs::write( + hooks_json_path.as_path(), + r#"{ + "hooks": { + "PreToolUse": [ + { + "matcher": "^Bash$", + "hooks": [ + { + "type": "command", + "command": "python3 /tmp/json-hook.py" + } + ] + } + ] + } + }"#, + ) + .expect("write hooks.json"); + let config_layer_stack = ConfigLayerStack::new( + vec![ + ConfigLayerEntry::new( + ConfigLayerSource::User { + file: config_path, + profile: None, + }, + TomlValue::Table(Default::default()), + ), + ConfigLayerEntry::new( + ConfigLayerSource::User { + file: profile_path, + profile: Some("work".to_string()), + }, + TomlValue::Table(Default::default()), + ), + ], + ConfigRequirements::default(), + ConfigRequirementsToml::default(), + ) + .expect("config layer stack"); + + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ true, + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + assert!(engine.warnings().is_empty()); + assert_eq!(engine.handlers.len(), 1); + let preview = engine.preview_pre_tool_use(&PreToolUseRequest { + session_id: ThreadId::new(), + turn_id: "turn-1".to_string(), + subagent: None, + cwd: cwd(), + transcript_path: None, + model: "gpt-test".to_string(), + permission_mode: "default".to_string(), + tool_name: "Bash".to_string(), + matcher_aliases: Vec::new(), + tool_use_id: "tool-1".to_string(), + tool_input: serde_json::json!({ "command": "echo hello" }), + }); + assert_eq!(preview.len(), 1); + assert_eq!(preview[0].source_path, hooks_json_path); + + let listed = crate::list_hooks(crate::HooksConfig { + feature_enabled: true, + bypass_hook_trust: true, + config_layer_stack: Some(config_layer_stack), + ..Default::default() + }); + assert!(listed.warnings.is_empty()); + assert_eq!(listed.hooks.len(), 1); + assert_eq!(listed.hooks[0].source_path, hooks_json_path); +} + +#[test] +fn malformed_hooks_json_is_reported_as_startup_warning() { + let temp = tempdir().expect("create temp dir"); + let config_path = + AbsolutePathBuf::try_from(temp.path().join("config.toml")).expect("absolute config path"); + let hooks_json_path = + AbsolutePathBuf::try_from(temp.path().join("hooks.json")).expect("absolute hooks path"); + fs::write( + hooks_json_path.as_path(), + r#"{ + "SessionStart": [ + { + "hooks": [ + { + "type": "command", + "command": "python3 /tmp/session-start.py" + } + ] + } + ] + }"#, + ) + .expect("write hooks.json"); + let config_layer_stack = ConfigLayerStack::new( + vec![ConfigLayerEntry::new( + ConfigLayerSource::System { file: config_path }, + TomlValue::Table(Default::default()), + )], + ConfigRequirements::default(), + ConfigRequirementsToml::default(), + ) + .expect("config layer stack"); + + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + assert!(engine.handlers.is_empty()); + assert_eq!(engine.warnings().len(), 1); + assert!(engine.warnings()[0].contains("failed to parse hooks config")); + assert!( + engine.warnings()[0].contains(&hooks_json_path.display().to_string()), + "warning should identify the malformed file: {}", + engine.warnings()[0] + ); + assert!(engine.warnings()[0].contains("unknown field `SessionStart`")); +} + +fn bundled_cleanup_source(plugin_id: &str, server: &str, event: &str) -> PluginHookSource { + let plugin_root = cwd().join("bundled-plugin"); + PluginHookSource { + plugin_id: PluginId::parse(plugin_id).expect("plugin ID"), + plugin_data_root: plugin_root.join("data"), + source_path: plugin_root.join(".codex-plugin/plugin.json"), + source_relative_path: "plugin.json#hooks[0]".to_string(), + plugin_root, + hooks: serde_json::from_value(serde_json::json!({ + (event): [{"hooks": [{ + "type": "mcp_tool", + "server": server, + "tool": "turn_ended", + "input": {"turn_id": "${turn_id}"}, + }]}], + })) + .expect("cleanup hooks"), + } +} + +#[test] +fn bundled_cleanup_hooks_are_trusted_without_saved_hashes() { + for (plugin_id, server) in [ + ("browser@openai-bundled", "node_repl"), + ("chrome@openai-bundled", "node_repl"), + ("chrome-dev@openai-bundled", "node_repl"), + ("chrome-internal@openai-bundled", "node_repl"), + ("computer-use@openai-bundled", "node_repl"), + ("unified-computer-use@openai-bundled", "cua_repl"), + ] { + for event in ["Stop", "Interrupt", "SubagentStop"] { + let discovered = super::discovery::discover_handlers( + /*config_layer_stack*/ None, + vec![bundled_cleanup_source(plugin_id, server, event)], + Vec::new(), + /*bypass_hook_trust*/ false, + ); + assert_eq!(discovered.handlers.len(), 1, "{plugin_id} {event}"); + assert_eq!( + discovered + .hook_entries + .iter() + .map(|entry| ( + entry.plugin_id.as_deref(), + entry.trust_status, + entry.enabled, + entry.is_managed, + entry.builtin, + )) + .collect::>(), + vec![(Some(plugin_id), HookTrustStatus::Trusted, true, false, true)], + "{plugin_id} {event}" + ); + assert!(discovered.handlers[0].builtin, "{plugin_id} {event}"); + } + } +} + +#[test] +fn bundled_cleanup_trust_does_not_extend_to_other_handlers() { + let mut source = bundled_cleanup_source("browser@openai-bundled", "node_repl", "Stop"); + let cleanup = source.hooks.stop[0].hooks[0].clone(); + source.hooks.stop[0].hooks.extend([ + HookHandlerConfig::McpTool { + server: "node_repl".to_string(), + tool: "evaluate".to_string(), + input: Default::default(), + timeout_sec: None, + status_message: None, + }, + HookHandlerConfig::McpTool { + server: "other_server".to_string(), + tool: "turn_ended".to_string(), + input: Default::default(), + timeout_sec: None, + status_message: None, + }, + HookHandlerConfig::Command { + command: "echo cleanup".to_string(), + command_windows: None, + timeout_sec: None, + r#async: false, + status_message: None, + additional_context_limit: None, + }, + ]); + source.hooks.stop.push(MatcherGroup { + matcher: Some("Bash".to_string()), + hooks: vec![cleanup], + }); + let discovered = super::discovery::discover_handlers( + /*config_layer_stack*/ None, + vec![source], + Vec::new(), + /*bypass_hook_trust*/ false, + ); + assert_eq!(discovered.handlers.len(), 1); + assert_eq!( + discovered + .hook_entries + .iter() + .map(|entry| (entry.trust_status, entry.builtin)) + .collect::>(), + vec![ + (HookTrustStatus::Trusted, true), + (HookTrustStatus::Untrusted, false), + (HookTrustStatus::Untrusted, false), + (HookTrustStatus::Untrusted, false), + (HookTrustStatus::Untrusted, false), + ] + ); +} + +#[test] +fn bundled_cleanup_trust_requires_matching_plugin_server_and_event() { + for (plugin_id, server, event) in [ + ("other@openai-bundled", "node_repl", "Stop"), + ("browser@other", "node_repl", "Stop"), + ("browser@openai-bundled-alpha", "node_repl", "Stop"), + ("browser@openai-bundled", "node_repl", "PreToolUse"), + ("browser@openai-bundled", "cua_repl", "Stop"), + ("unified-computer-use@openai-bundled", "node_repl", "Stop"), + ("unified-computer-use@other", "cua_repl", "Stop"), + ( + "unified-computer-use@openai-bundled", + "cua_repl", + "PreToolUse", + ), + ] { + let discovered = super::discovery::discover_handlers( + /*config_layer_stack*/ None, + vec![bundled_cleanup_source(plugin_id, server, event)], + Vec::new(), + /*bypass_hook_trust*/ false, + ); + assert!(discovered.handlers.is_empty(), "{plugin_id} {event}"); + assert_eq!(discovered.hook_entries.len(), 1); + assert_eq!( + discovered.hook_entries[0].trust_status, + HookTrustStatus::Untrusted + ); + } +} + +#[test] +fn local_hosted_app_cleanup_hooks_require_saved_trust() { + let mut source = bundled_cleanup_source("browser@openai-curated-remote", "codex_apps", "Stop"); + source.hooks.stop[0].hooks[0] = HookHandlerConfig::McpTool { + server: "codex_apps".to_string(), + tool: "browser.turn_ended".to_string(), + input: Default::default(), + timeout_sec: None, + status_message: None, + }; + let discovered = super::discovery::discover_handlers( + /*config_layer_stack*/ None, + vec![source], + Vec::new(), + /*bypass_hook_trust*/ false, + ); + assert!(discovered.handlers.is_empty()); + assert_eq!( + discovered + .hook_entries + .iter() + .map(|entry| entry.trust_status) + .collect::>(), + vec![HookTrustStatus::Untrusted] + ); +} + +#[test] +fn builtin_cleanup_ignores_disablement_but_preserves_managed_only_policy() { + let source = bundled_cleanup_source("browser@openai-bundled", "node_repl", "Stop"); + let discovered = super::discovery::discover_handlers( + /*config_layer_stack*/ None, + vec![source.clone()], + Vec::new(), + /*bypass_hook_trust*/ false, + ); + let expected_entry = discovered.hook_entries[0].clone(); + let disabled_stack = ConfigLayerStack::new( + vec![ConfigLayerEntry::new( + ConfigLayerSource::User { + file: cwd().join("config.toml"), + profile: None, + }, + config_with_hook_state(&expected_entry.key, /*enabled*/ false), + )], + ConfigRequirements::default(), + ConfigRequirementsToml::default(), + ) + .expect("disabled hook config"); + let disabled = super::discovery::discover_handlers( + Some(&disabled_stack), + vec![source.clone()], + Vec::new(), + /*bypass_hook_trust*/ false, + ); + assert_eq!(disabled.handlers, discovered.handlers); + assert_eq!(disabled.hook_entries, vec![expected_entry]); + + let (requirements, requirements_toml) = requirements_with_managed_hooks_only( + /*allow_managed_hooks_only*/ true, /*managed_hooks*/ None, + ); + let managed_stack = ConfigLayerStack::new(Vec::new(), requirements, requirements_toml) + .expect("managed-only config"); + let managed_only = super::discovery::discover_handlers( + Some(&managed_stack), + vec![source.clone()], + Vec::new(), + /*bypass_hook_trust*/ false, + ); + assert!(managed_only.handlers.is_empty()); + assert!(managed_only.hook_entries.is_empty()); + + for (stack, expected) in [ + (&disabled_stack, discovered.handlers), + (&managed_stack, Vec::new()), + ] { + let feature_disabled = ClaudeHooksEngine::new( + /*enabled*/ false, + /*bypass_hook_trust*/ false, + Some(stack), + vec![source.clone()], + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + assert_eq!(feature_disabled.handlers, expected); + } +} + +#[test] +fn disabled_hooks_feature_keeps_builtin_cleanup_but_not_trusted_plugin_hooks() { + let mut source = bundled_cleanup_source("browser@openai-bundled", "node_repl", "Stop"); + source.hooks.stop[0].hooks.push(HookHandlerConfig::Command { + command: "echo ordinary hook".to_string(), + command_windows: None, + timeout_sec: None, + r#async: false, + status_message: None, + additional_context_limit: None, + }); + let stack = trusted_plugin_hook_stack(cwd().join("config.toml"), &[source.clone()]); + let discovered = super::discovery::discover_handlers( + Some(&stack), + vec![source.clone()], + Vec::new(), + /*bypass_hook_trust*/ false, + ); + assert_eq!(discovered.handlers.len(), 2); + + let engine = ClaudeHooksEngine::new( + /*enabled*/ false, + /*bypass_hook_trust*/ false, + Some(&stack), + vec![source], + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + assert_eq!(engine.handlers, vec![discovered.handlers[0].clone()]); +} + +#[tokio::test] +async fn plugin_hook_sources_run_with_plugin_env_and_plugin_source() { + let temp = tempdir().expect("create temp dir"); + let plugin_root = + AbsolutePathBuf::try_from(temp.path().join("demo-plugin")).expect("plugin root"); + let plugin_data_root = + AbsolutePathBuf::try_from(temp.path().join("plugin-data")).expect("plugin data root"); + fs::create_dir_all(plugin_root.join("hooks")).expect("create hooks dir"); + let source_path = plugin_root.join("hooks/hooks.json"); + let script_path = plugin_root.join("hooks/write_env.py"); + fs::write( + script_path.as_path(), + r#"import json +import os +print(json.dumps({ + "systemMessage": json.dumps({ + "plugin": os.environ.get("PLUGIN_ROOT"), + "claude": os.environ.get("CLAUDE_PLUGIN_ROOT"), + }) +})) +"#, + ) + .expect("write hook script"); + let plugin_id = PluginId::parse("demo-plugin@test-marketplace").expect("plugin id"); + let plugin_hook_sources = vec![PluginHookSource { + plugin_id, + plugin_root: plugin_root.clone(), + plugin_data_root: plugin_data_root.clone(), + source_path: source_path.clone(), + source_relative_path: "hooks/hooks.json".to_string(), + hooks: HookEventsToml { + pre_tool_use: vec![MatcherGroup { + matcher: Some("Bash".to_string()), + hooks: vec![HookHandlerConfig::Command { + command: format!("python3 {}", script_path.display()), + command_windows: None, + timeout_sec: Some(10), + r#async: false, + status_message: None, + additional_context_limit: None, + }], + }], + ..Default::default() + }, + }]; + let config_layer_stack = trusted_plugin_hook_stack( + AbsolutePathBuf::try_from(temp.path().join("config.toml")).expect("absolute config path"), + &plugin_hook_sources, + ); + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + Some(&config_layer_stack), + plugin_hook_sources.clone(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + let preview = engine.preview_pre_tool_use(&PreToolUseRequest { + session_id: ThreadId::new(), + turn_id: "turn-1".to_string(), + subagent: None, + cwd: cwd(), + transcript_path: None, + model: "gpt-test".to_string(), + permission_mode: "default".to_string(), + tool_name: "Bash".to_string(), + matcher_aliases: Vec::new(), + tool_use_id: "tool-1".to_string(), + tool_input: serde_json::json!({ "command": "echo hello" }), + }); + assert_eq!(preview.len(), 1); + assert_eq!(preview[0].source, HookSource::Plugin); + assert_eq!(preview[0].source_path, source_path); + let listed = crate::list_hooks(crate::HooksConfig { + legacy_notify_argv: None, + feature_enabled: true, + bypass_hook_trust: false, + config_layer_stack: None, + plugin_hook_sources, + plugin_hook_load_warnings: Vec::new(), + shell_program: None, + shell_args: Vec::new(), + }); + assert_eq!( + listed.hooks[0].plugin_id.as_deref(), + Some("demo-plugin@test-marketplace") + ); + + let outcome = engine + .run_pre_tool_use(PreToolUseRequest { + session_id: ThreadId::new(), + turn_id: "turn-1".to_string(), + subagent: None, + cwd: cwd(), + transcript_path: None, + model: "gpt-test".to_string(), + permission_mode: "default".to_string(), + tool_name: "Bash".to_string(), + matcher_aliases: Vec::new(), + tool_use_id: "tool-1".to_string(), + tool_input: serde_json::json!({ "command": "echo hello" }), + }) + .await; + + assert_eq!(outcome.hook_events.len(), 1); + assert_eq!(outcome.hook_events[0].run.source, HookSource::Plugin); + assert_eq!( + outcome.hook_events[0].run.status, + HookRunStatus::Completed, + "hook entries: {:#?}", + outcome.hook_events[0].run.entries + ); + assert_eq!(outcome.hook_events[0].run.entries.len(), 1); + assert_eq!( + outcome.hook_events[0].run.entries[0].kind, + HookOutputEntryKind::Warning + ); + let logged: serde_json::Value = + serde_json::from_str(&outcome.hook_events[0].run.entries[0].text) + .expect("parse env payload"); + assert_eq!( + logged, + serde_json::json!({ + "plugin": plugin_root.display().to_string(), + "claude": plugin_root.display().to_string(), + }) + ); +} + +#[test] +fn plugin_hook_sources_expand_plugin_placeholders() { + let temp = tempdir().expect("create temp dir"); + let plugin_root = + AbsolutePathBuf::try_from(temp.path().join("demo-plugin")).expect("plugin root"); + let plugin_data_root = + AbsolutePathBuf::try_from(temp.path().join("plugin-data")).expect("plugin data root"); + let source_path = plugin_root.join("hooks/hooks.json"); + let plugin_id = PluginId::parse("demo-plugin@test-marketplace").expect("plugin id"); + let plugin_hook_sources = vec![PluginHookSource { + plugin_id, + plugin_root: plugin_root.clone(), + plugin_data_root: plugin_data_root.clone(), + source_path, + source_relative_path: "hooks/hooks.json".to_string(), + hooks: HookEventsToml { + pre_tool_use: vec![MatcherGroup { + matcher: Some("Bash".to_string()), + hooks: vec![HookHandlerConfig::Command { + command: + "run ${PLUGIN_ROOT} ${CLAUDE_PLUGIN_ROOT} ${PLUGIN_DATA} ${CLAUDE_PLUGIN_DATA}" + .to_string(), + command_windows: None, + timeout_sec: Some(5), + r#async: false, + status_message: None, + additional_context_limit: None, + }], + }], + ..Default::default() + }, + }]; + let config_layer_stack = trusted_plugin_hook_stack( + AbsolutePathBuf::try_from(temp.path().join("config.toml")).expect("absolute config path"), + &plugin_hook_sources, + ); + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + Some(&config_layer_stack), + plugin_hook_sources, + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + assert_eq!( + engine.handlers[0].kind, + ConfiguredHandlerKind::Command { + command: format!( + "run {} {} {} {}", + plugin_root.display(), + plugin_root.display(), + plugin_data_root.display(), + plugin_data_root.display() + ), + r#async: false, + env: HashMap::from([ + ("PLUGIN_ROOT".to_string(), plugin_root.display().to_string()), + ( + "CLAUDE_PLUGIN_ROOT".to_string(), + plugin_root.display().to_string(), + ), + ( + "PLUGIN_DATA".to_string(), + plugin_data_root.display().to_string(), + ), + ( + "CLAUDE_PLUGIN_DATA".to_string(), + plugin_data_root.display().to_string(), + ), + ]), + } + ); +} + +#[test] +fn plugin_hook_load_warnings_are_startup_warnings() { + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + /*config_layer_stack*/ None, + Vec::new(), + vec!["failed plugin hook".to_string()], + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + mcp_executor(), + ); + + assert_eq!(engine.warnings(), &["failed plugin hook".to_string()]); +} + +struct StaticMcpExecutor { + calls: Arc>>, + output: String, + outputs_by_tool: HashMap, +} + +impl HookMcpExecutor for StaticMcpExecutor { + fn execute(&self, call: HookMcpCall) -> BoxFuture<'_, anyhow::Result> { + async move { + let output = self + .outputs_by_tool + .get(&call.tool) + .unwrap_or(&self.output) + .clone(); + self.calls.lock().expect("lock MCP calls").push(call); + Ok(output) + } + .boxed() + } +} + +fn executor_stop_hook_fixture() -> ( + ClaudeHooksEngine, + Arc>>, + StopRequest, + HookMcpCall, + ExecutorPluginHookSource, +) { + let calls = Arc::new(Mutex::new(Vec::new())); + let executor = Arc::new(StaticMcpExecutor { + calls: Arc::clone(&calls), + output: r#"{"decision":"block","reason":"keep going"}"#.to_string(), + outputs_by_tool: HashMap::from([( + "terminate".to_string(), + r#"{"continue":false,"stopReason":"done"}"#.to_string(), + )]), + }); + let mut engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ false, + /*config_layer_stack*/ None, + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + executor, + ); + let source = ExecutorPluginHookSource { + plugin_id: PluginId::parse("computer-use@openai-bundled").expect("valid plugin ID"), + environment_id: "executor-a".to_string(), + mcp_environment_id: None, + mcp_metadata: None, + plugin_root: "file:///plugins/computer-use" + .parse() + .expect("valid plugin root URI"), + manifest_path: "file:///plugins/computer-use/.codex-plugin/plugin.json" + .parse() + .expect("valid plugin manifest URI"), + source_relative_path: ".codex-plugin/plugin.json#hooks[0]".to_string(), + hooks: HookEventsToml { + stop: vec![MatcherGroup { + matcher: None, + hooks: vec![HookHandlerConfig::McpTool { + server: "node_repl".to_string(), + tool: "turn_ended".to_string(), + input: serde_json::from_value(serde_json::json!({ + "turn_id": "${turn_id}", + })) + .expect("executor hook input"), + timeout_sec: None, + status_message: None, + }], + }], + ..Default::default() + }, + }; + engine.set_executor_hooks(vec![source.clone()]); + assert_eq!( + engine.handlers, + vec![ConfiguredHandler { + builtin: true, + event_name: HookEventName::Stop, + matcher: None, + timeout_sec: 5, + status_message: None, + additional_context_limit: Default::default(), + source_path: HandlerSourcePath::ExecutorScoped { + plugin_id: PluginId::parse("computer-use@openai-bundled").expect("valid plugin ID"), + environment_id: "executor-a".to_string(), + mcp_environment_id: None, + mcp_metadata: None, + manifest_path: "file:///plugins/computer-use/.codex-plugin/plugin.json" + .parse() + .expect("valid plugin manifest URI"), + source_relative_path: ".codex-plugin/plugin.json#hooks[0]".to_string(), + }, + source: HookSource::Plugin, + display_order: 0, + kind: ConfiguredHandlerKind::McpTool { + server: "node_repl".to_string(), + tool: "turn_ended".to_string(), + input: serde_json::from_value(serde_json::json!({ + "turn_id": "${turn_id}", + })) + .expect("executor hook input"), + }, + }] + ); + assert_eq!( + engine.handlers[0].execution_mode(), + HookExecutionMode::Async + ); + assert!(!engine.handlers[0].can_apply_control_effects()); + let request_metadata = Some(serde_json::Map::from_iter([( + "x-codex-turn-metadata".to_string(), + serde_json::json!({ "turn_id": "turn-1" }), + )])); + let request = StopRequest { + session_id: ThreadId::new(), + turn_id: "turn-1".to_string(), + cwd: cwd(), + transcript_path: None, + model: "test-model".to_string(), + permission_mode: "default".to_string(), + request_metadata: request_metadata.clone(), + stop_hook_active: false, + last_assistant_message: None, + target: StopHookTarget::Stop, + }; + + let expected_executor_call = HookMcpCall { + server: "node_repl".to_string(), + tool: "turn_ended".to_string(), + environment_id: Some("executor-a".to_string()), + metadata: request_metadata, + input: serde_json::from_value(serde_json::json!({ "turn_id": "turn-1" })) + .expect("expanded executor hook input"), + timeout: Duration::from_secs(5), + }; + + (engine, calls, request, expected_executor_call, source) +} + +async fn wait_for_mcp_calls(calls: &Arc>>, count: usize) { + tokio::time::timeout(Duration::from_secs(1), async { + while calls.lock().expect("lock MCP calls").len() < count { + tokio::task::yield_now().await; + } + }) + .await + .expect("executor hook should complete in the background"); +} + +#[tokio::test] +async fn executor_stop_hooks_run_unless_regular_hooks_block_without_stopping() { + let (mut engine, calls, request, expected_executor_call, _source) = + executor_stop_hook_fixture(); + + assert_eq!(engine.preview_stop(&request), Vec::new()); + let outcome = engine.run_stop(request.clone()).await; + wait_for_mcp_calls(&calls, /*count*/ 1).await; + assert!(!outcome.should_block); + assert!(outcome.hook_events.is_empty()); + assert_eq!( + *calls.lock().expect("lock MCP calls"), + vec![expected_executor_call.clone()] + ); + + engine.handlers.push(ConfiguredHandler { + builtin: false, + event_name: HookEventName::Stop, + matcher: None, + timeout_sec: 30, + status_message: None, + additional_context_limit: Default::default(), + source_path: cwd().join("hooks.json").into(), + source: HookSource::User, + display_order: 0, + kind: ConfiguredHandlerKind::McpTool { + server: "security".to_string(), + tool: "check".to_string(), + input: Default::default(), + }, + }); + + let outcome = engine.run_stop(request.clone()).await; + assert!(outcome.should_block); + assert_eq!(outcome.hook_events.len(), 1); + assert_eq!( + *calls.lock().expect("lock MCP calls"), + vec![ + expected_executor_call.clone(), + HookMcpCall { + server: "security".to_string(), + tool: "check".to_string(), + environment_id: None, + metadata: None, + input: Default::default(), + timeout: Duration::from_secs(30), + }, + ] + ); + + engine.handlers.push(ConfiguredHandler { + builtin: false, + event_name: HookEventName::Stop, + matcher: None, + timeout_sec: 30, + status_message: None, + additional_context_limit: Default::default(), + source_path: cwd().join("hooks.json").into(), + source: HookSource::User, + display_order: 1, + kind: ConfiguredHandlerKind::McpTool { + server: "security".to_string(), + tool: "terminate".to_string(), + input: Default::default(), + }, + }); + + let outcome = engine.run_stop(request).await; + wait_for_mcp_calls(&calls, /*count*/ 5).await; + assert!(outcome.should_stop); + assert!(!outcome.should_block); + assert_eq!(outcome.hook_events.len(), 2); + assert_eq!( + *calls.lock().expect("lock MCP calls"), + vec![ + expected_executor_call.clone(), + HookMcpCall { + server: "security".to_string(), + tool: "check".to_string(), + environment_id: None, + metadata: None, + input: Default::default(), + timeout: Duration::from_secs(30), + }, + HookMcpCall { + server: "security".to_string(), + tool: "check".to_string(), + environment_id: None, + metadata: None, + input: Default::default(), + timeout: Duration::from_secs(30), + }, + HookMcpCall { + server: "security".to_string(), + tool: "terminate".to_string(), + environment_id: None, + metadata: None, + input: Default::default(), + timeout: Duration::from_secs(30), + }, + expected_executor_call, + ] + ); +} + +#[test] +fn executor_stop_hooks_register_each_target_environment_once() { + let (mut engine, _, _, _, first_source) = executor_stop_hook_fixture(); + let mut expected_handlers = engine.handlers.clone(); + let mut duplicate_target = first_source.clone(); + duplicate_target.plugin_id = PluginId::parse("chrome@openai-bundled").expect("valid plugin ID"); + let mut second_source = first_source.clone(); + second_source.environment_id = "executor-b".to_string(); + let mut second_handler = expected_handlers[0].clone(); + second_handler.display_order = 1; + let HandlerSourcePath::ExecutorScoped { environment_id, .. } = &mut second_handler.source_path + else { + panic!("executor Stop handler should have an executor source"); + }; + *environment_id = second_source.environment_id.clone(); + expected_handlers.push(second_handler); + engine.set_executor_hooks(vec![first_source, duplicate_target, second_source]); + + assert_eq!(engine.handlers, expected_handlers); +} + +#[tokio::test] +async fn executor_interrupt_hooks_register_and_run() { + let (mut engine, calls, stop_request, expected_stop_call, mut source) = + executor_stop_hook_fixture(); + source.hooks.subagent_stop = source.hooks.stop.clone(); + source.hooks.interrupt = source.hooks.stop.clone(); + engine.set_executor_hooks(vec![source]); + + assert_eq!( + engine + .handlers + .iter() + .map(|handler| handler.event_name) + .collect::>(), + vec![ + HookEventName::SubagentStop, + HookEventName::Stop, + HookEventName::Interrupt, + ] + ); + assert_eq!(engine.preview_interrupt(), Vec::new()); + + let outcome = engine + .run_interrupt(InterruptRequest { + session_id: stop_request.session_id, + turn_id: stop_request.turn_id, + cwd: stop_request.cwd, + transcript_path: stop_request.transcript_path, + model: stop_request.model, + permission_mode: stop_request.permission_mode, + request_metadata: stop_request.request_metadata, + }) + .await; + wait_for_mcp_calls(&calls, /*count*/ 1).await; + + assert!(outcome.hook_events.is_empty()); + assert_eq!( + *calls.lock().expect("lock MCP calls"), + vec![expected_stop_call] + ); +} + +#[test] +fn executor_hooks_register_events_from_each_environment() { + let (mut engine, _, _, _, mut first_source) = executor_stop_hook_fixture(); + first_source.hooks.interrupt = std::mem::take(&mut first_source.hooks.stop); + engine.set_executor_hooks(vec![first_source.clone()]); + let mut expected_handlers = engine.handlers.clone(); + let mut second_source = first_source.clone(); + second_source.environment_id = "executor-b".to_string(); + second_source.hooks.stop = std::mem::take(&mut second_source.hooks.interrupt); + let mut second_handler = expected_handlers[0].clone(); + second_handler.event_name = HookEventName::Stop; + second_handler.display_order = 1; + let HandlerSourcePath::ExecutorScoped { environment_id, .. } = &mut second_handler.source_path + else { + panic!("executor Stop handler should have an executor source"); + }; + *environment_id = second_source.environment_id.clone(); + expected_handlers.push(second_handler); + + engine.set_executor_hooks(vec![first_source, second_source]); + + assert_eq!(engine.handlers, expected_handlers); +} + +#[tokio::test] +async fn executor_stop_hooks_use_an_admitted_mcp_environment() { + let (mut engine, calls, request, mut expected_call, mut source) = executor_stop_hook_fixture(); + source.mcp_environment_id = Some("local".to_string()); + engine.set_executor_hooks(vec![source]); + + engine.run_stop(request).await; + wait_for_mcp_calls(&calls, /*count*/ 1).await; + + expected_call.environment_id = Some("local".to_string()); + assert_eq!(*calls.lock().expect("lock MCP calls"), vec![expected_call]); +} + +#[tokio::test] +async fn memory_consolidation_stop_preserves_policy_and_executor_cleanup() { + for (source, runs_policy) in [ + (HookSource::User, false), + (HookSource::Project, false), + (HookSource::SessionFlags, false), + (HookSource::Plugin, false), + (HookSource::System, true), + (HookSource::Mdm, true), + (HookSource::CloudRequirements, true), + (HookSource::CloudManagedConfig, true), + (HookSource::LegacyManagedConfigFile, true), + (HookSource::LegacyManagedConfigMdm, true), + (HookSource::Unknown, true), + ] { + let (mut engine, calls, mut request, expected_executor_call, _source) = + executor_stop_hook_fixture(); + request.target = StopHookTarget::MemoryConsolidation; + let policy_call = HookMcpCall { + server: "security".to_string(), + tool: "check".to_string(), + environment_id: None, + metadata: None, + input: Default::default(), + timeout: Duration::from_secs(5), + }; + let mut handler = engine.handlers[0].clone(); + handler.source_path = cwd().join("hooks.json").into(); + handler.source = source; + handler.kind = ConfiguredHandlerKind::McpTool { + server: policy_call.server.clone(), + tool: policy_call.tool.clone(), + input: policy_call.input.clone(), + }; + engine.handlers.push(handler); + assert_eq!( + engine.preview_stop(&request).len(), + usize::from(runs_policy) + ); + + let outcome = engine.run_stop(request).await; + let expected_calls = if runs_policy { + vec![policy_call, expected_executor_call] + } else { + vec![expected_executor_call] + }; + wait_for_mcp_calls(&calls, expected_calls.len()).await; + assert_eq!(*calls.lock().expect("lock MCP calls"), expected_calls); + assert_eq!(outcome.should_block, runs_policy); + } +} + +#[tokio::test] +async fn executor_stop_hooks_do_not_delay_stop_completion() { + struct BlockingMcpExecutor { + started: Arc, + release: Arc, + } + + impl HookMcpExecutor for BlockingMcpExecutor { + fn execute(&self, _call: HookMcpCall) -> BoxFuture<'_, anyhow::Result> { + async move { + self.started.notify_one(); + self.release.notified().await; + Ok(String::new()) + } + .boxed() + } + } + + let (mut engine, _, request, _, _) = executor_stop_hook_fixture(); + let started = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + engine.mcp_executor = Arc::new(BlockingMcpExecutor { + started: Arc::clone(&started), + release: Arc::clone(&release), + }); + + let outcome = tokio::time::timeout(Duration::from_millis(100), engine.run_stop(request)) + .await + .expect("executor hook must not delay Stop completion"); + assert!(!outcome.should_block); + tokio::time::timeout(Duration::from_secs(1), started.notified()) + .await + .expect("executor hook should start in the background"); + release.notify_one(); +} + +#[tokio::test] +async fn mcp_tool_hooks_expand_event_input_and_apply_pre_tool_decisions() { + let temp = tempdir().expect("create temp dir"); + let config_path = + AbsolutePathBuf::try_from(temp.path().join("config.toml")).expect("absolute config path"); + fs::write( + temp.path().join("hooks.json"), + serde_json::json!({ + "hooks": { + "PreToolUse": [{ + "matcher": "Bash", + "hooks": [{ + "type": "mcp_tool", + "server": "security", + "tool": "scan", + "input": { "command": "${tool_input.command}" }, + "timeout": 20, + }], + }], + }, + }) + .to_string(), + ) + .expect("write MCP hooks.json"); + let config_layer_stack = ConfigLayerStack::new( + vec![ConfigLayerEntry::new( + ConfigLayerSource::User { + file: config_path, + profile: None, + }, + TomlValue::Table(Default::default()), + )], + ConfigRequirements::default(), + ConfigRequirementsToml::default(), + ) + .expect("config layer stack"); + + let request = PreToolUseRequest { + session_id: ThreadId::new(), + turn_id: "turn-1".to_string(), + subagent: None, + cwd: cwd(), + transcript_path: None, + model: "gpt-test".to_string(), + permission_mode: "default".to_string(), + tool_name: "Bash".to_string(), + matcher_aliases: Vec::new(), + tool_use_id: "tool-1".to_string(), + tool_input: serde_json::json!({ "command": "rm important.txt" }), + }; + let calls = Arc::new(Mutex::new(Vec::new())); + let executor = StaticMcpExecutor { + calls: Arc::clone(&calls), + output: serde_json::json!({ + "hookSpecificOutput": { + "hookEventName": "PreToolUse", + "permissionDecision": "deny", + "permissionDecisionReason": "blocked by MCP scanner", + }, + }) + .to_string(), + outputs_by_tool: HashMap::new(), + }; + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ true, + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + Arc::new(executor), + ); + let outcome = engine.run_pre_tool_use(request).await; + + assert!(outcome.should_block); + assert_eq!( + outcome.block_reason.as_deref(), + Some("blocked by MCP scanner") + ); + assert_eq!( + outcome.hook_events[0].run.handler_type, + HookHandlerType::McpTool + ); + assert_eq!( + outcome.hook_events[0].run.execution_mode, + codex_protocol::protocol::HookExecutionMode::Sync + ); + assert_eq!( + *calls.lock().expect("lock MCP calls"), + vec![HookMcpCall { + server: "security".to_string(), + tool: "scan".to_string(), + environment_id: None, + metadata: None, + input: serde_json::from_value(serde_json::json!({ + "command": "rm important.txt", + })) + .expect("object input"), + timeout: Duration::from_secs(20), + }] + ); +} + +#[tokio::test] +async fn mcp_interrupt_hooks_expand_event_input_and_bound_timeout() { + let temp = tempdir().expect("create temp dir"); + let config_path = + AbsolutePathBuf::try_from(temp.path().join("config.toml")).expect("absolute config path"); + fs::write( + temp.path().join("hooks.json"), + serde_json::json!({ + "hooks": { + "Interrupt": [{ + "hooks": [{ + "type": "mcp_tool", + "server": "security", + "tool": "notify", + "input": { + "event": "${hook_event_name}", + "turn_id": "${turn_id}", + "permission_mode": "${permission_mode}", + }, + "timeout": 20, + }], + }], + }, + }) + .to_string(), + ) + .expect("write MCP Interrupt hooks.json"); + let config_layer_stack = ConfigLayerStack::new( + vec![ConfigLayerEntry::new( + ConfigLayerSource::User { + file: config_path, + profile: None, + }, + TomlValue::Table(Default::default()), + )], + ConfigRequirements::default(), + ConfigRequirementsToml::default(), + ) + .expect("config layer stack"); + + let calls = Arc::new(Mutex::new(Vec::new())); + let executor = StaticMcpExecutor { + calls: Arc::clone(&calls), + output: serde_json::json!({ + "systemMessage": "interrupt observed", + }) + .to_string(), + outputs_by_tool: HashMap::new(), + }; + let engine = ClaudeHooksEngine::new( + /*enabled*/ true, + /*bypass_hook_trust*/ true, + Some(&config_layer_stack), + Vec::new(), + Vec::new(), + command_runtime(CommandShell { + program: String::new(), + args: Vec::new(), + }), + Arc::new(executor), + ); + let outcome = engine + .run_interrupt(InterruptRequest { + session_id: ThreadId::new(), + turn_id: "turn-1".to_string(), + cwd: cwd(), + transcript_path: None, + model: "gpt-test".to_string(), + permission_mode: "default".to_string(), + request_metadata: None, + }) + .await; + + assert_eq!(outcome.hook_events.len(), 1); + assert_eq!( + outcome.hook_events[0].run.handler_type, + HookHandlerType::McpTool + ); + assert_eq!(outcome.hook_events[0].run.status, HookRunStatus::Completed); + assert_eq!( + outcome.hook_events[0].run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Warning, + text: "interrupt observed".to_string(), + }] + ); + assert_eq!( + *calls.lock().expect("lock MCP calls"), + vec![HookMcpCall { + server: "security".to_string(), + tool: "notify".to_string(), + environment_id: None, + metadata: None, + input: serde_json::from_value(serde_json::json!({ + "event": "Interrupt", + "turn_id": "turn-1", + "permission_mode": "default", + })) + .expect("object input"), + timeout: Duration::from_secs(3), + }] + ); +} diff --git a/codex-rs/hooks/src/engine/output_parser.rs b/codex-rs/hooks/src/engine/output_parser.rs new file mode 100644 index 0000000000000000000000000000000000000000..aaefa46dc10217a5e9259b9f8ff010bfee0f7c36 --- /dev/null +++ b/codex-rs/hooks/src/engine/output_parser.rs @@ -0,0 +1,617 @@ +#[derive(Debug, Clone)] +pub(crate) struct UniversalOutput { + pub continue_processing: bool, + pub stop_reason: Option, + pub suppress_output: bool, + pub system_message: Option, +} + +#[derive(Debug, Clone)] +pub(crate) struct SessionStartOutput { + pub universal: UniversalOutput, + pub additional_context: Option, +} + +#[derive(Debug, Clone)] +pub(crate) struct PreToolUseOutput { + pub universal: UniversalOutput, + pub block_reason: Option, + pub additional_context: Option, + pub updated_input: Option, + pub invalid_reason: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum PermissionRequestDecision { + Allow, + Deny { message: String }, +} + +#[derive(Debug, Clone)] +pub(crate) struct PermissionRequestOutput { + pub universal: UniversalOutput, + pub decision: Option, + pub invalid_reason: Option, +} + +#[derive(Debug, Clone)] +pub(crate) struct PostToolUseOutput { + pub universal: UniversalOutput, + pub should_block: bool, + pub reason: Option, + pub invalid_block_reason: Option, + pub additional_context: Option, + pub invalid_reason: Option, +} + +#[derive(Debug, Clone)] +pub(crate) struct UserPromptSubmitOutput { + pub universal: UniversalOutput, + pub should_block: bool, + pub reason: Option, + pub invalid_block_reason: Option, + pub additional_context: Option, +} + +#[derive(Debug, Clone)] +pub(crate) struct StopOutput { + pub universal: UniversalOutput, + pub should_block: bool, + pub reason: Option, + pub invalid_block_reason: Option, +} + +#[derive(Debug, Clone)] +pub(crate) struct StatelessHookOutput { + pub universal: UniversalOutput, + pub invalid_reason: Option, +} + +#[derive(Debug, Clone)] +pub(crate) struct InterruptOutput { + pub system_message: Option, +} + +use crate::schema::BlockDecisionWire; +use crate::schema::HookUniversalOutputWire; +use crate::schema::InterruptCommandOutputWire; +use crate::schema::PermissionRequestBehaviorWire; +use crate::schema::PermissionRequestCommandOutputWire; +use crate::schema::PermissionRequestDecisionWire; +use crate::schema::PostCompactCommandOutputWire; +use crate::schema::PostToolUseCommandOutputWire; +use crate::schema::PreCompactCommandOutputWire; +use crate::schema::PreToolUseCommandOutputWire; +use crate::schema::PreToolUseDecisionWire; +use crate::schema::PreToolUsePermissionDecisionWire; +use crate::schema::SessionStartCommandOutputWire; +use crate::schema::StopCommandOutputWire; +use crate::schema::SubagentStartCommandOutputWire; +use crate::schema::SubagentStopCommandOutputWire; +use crate::schema::UserPromptSubmitCommandOutputWire; + +pub(crate) fn parse_session_start(stdout: &str) -> Option { + let wire: SessionStartCommandOutputWire = parse_json(stdout)?; + Some(session_start_output( + wire.universal, + wire.hook_specific_output + .and_then(|output| output.additional_context), + )) +} + +pub(crate) fn parse_subagent_start(stdout: &str) -> Option { + let wire: SubagentStartCommandOutputWire = parse_json(stdout)?; + Some(session_start_output( + wire.universal, + wire.hook_specific_output + .and_then(|output| output.additional_context), + )) +} + +fn session_start_output( + universal: HookUniversalOutputWire, + additional_context: Option, +) -> SessionStartOutput { + SessionStartOutput { + universal: UniversalOutput::from(universal), + additional_context, + } +} + +pub(crate) fn parse_pre_tool_use(stdout: &str) -> Option { + let PreToolUseCommandOutputWire { + universal: universal_wire, + decision, + reason, + hook_specific_output, + } = parse_json(stdout)?; + let universal = UniversalOutput::from(universal_wire); + let hook_specific_output = hook_specific_output.as_ref(); + let additional_context = + hook_specific_output.and_then(|output| output.additional_context.clone()); + let use_hook_specific_decision = hook_specific_output.is_some_and(|output| { + output.permission_decision.is_some() + || output.permission_decision_reason.is_some() + || output.updated_input.is_some() + }); + let invalid_reason = unsupported_pre_tool_use_universal(&universal).or_else(|| { + if use_hook_specific_decision { + hook_specific_output.and_then(unsupported_pre_tool_use_hook_specific_output) + } else { + unsupported_pre_tool_use_legacy_decision(decision.as_ref(), reason.as_deref()) + } + }); + let block_reason = if invalid_reason.is_none() { + if use_hook_specific_decision { + hook_specific_output.and_then(|output| match output.permission_decision { + Some(PreToolUsePermissionDecisionWire::Deny) => output + .permission_decision_reason + .as_deref() + .and_then(trimmed_reason), + _ => None, + }) + } else { + match decision.as_ref() { + Some(PreToolUseDecisionWire::Block) => reason.as_deref().and_then(trimmed_reason), + Some(PreToolUseDecisionWire::Approve) | None => None, + } + } + } else { + None + }; + let updated_input = if invalid_reason.is_none() { + hook_specific_output.and_then(|output| { + matches!( + output.permission_decision, + Some(PreToolUsePermissionDecisionWire::Allow) + ) + .then(|| output.updated_input.clone()) + .flatten() + }) + } else { + None + }; + + Some(PreToolUseOutput { + universal, + block_reason, + additional_context, + updated_input, + invalid_reason, + }) +} + +pub(crate) fn parse_permission_request(stdout: &str) -> Option { + let wire: PermissionRequestCommandOutputWire = parse_json(stdout)?; + let universal = UniversalOutput::from(wire.universal); + let hook_specific_output = wire.hook_specific_output.as_ref(); + let decision = hook_specific_output.and_then(|output| output.decision.as_ref()); + let invalid_reason = unsupported_permission_request_universal(&universal).or_else(|| { + hook_specific_output.and_then(|output| { + unsupported_permission_request_hook_specific_output(output.decision.as_ref()) + }) + }); + let decision = if invalid_reason.is_none() { + decision.map(permission_request_decision) + } else { + None + }; + + Some(PermissionRequestOutput { + universal, + decision, + invalid_reason, + }) +} + +pub(crate) fn parse_post_tool_use(stdout: &str) -> Option { + let wire: PostToolUseCommandOutputWire = parse_json(stdout)?; + let universal = UniversalOutput::from(wire.universal); + let invalid_reason = unsupported_post_tool_use_universal(&universal).or_else(|| { + wire.hook_specific_output + .as_ref() + .and_then(unsupported_post_tool_use_hook_specific_output) + }); + let should_block = matches!(wire.decision, Some(BlockDecisionWire::Block)); + let invalid_block_reason = if should_block + && match wire.reason.as_deref() { + Some(reason) => reason.trim().is_empty(), + None => true, + } { + Some(invalid_block_message("PostToolUse")) + } else if !should_block && universal.continue_processing && wire.reason.is_some() { + Some("PostToolUse hook returned reason without decision".to_string()) + } else { + None + }; + let additional_context = wire + .hook_specific_output + .and_then(|output| output.additional_context); + + Some(PostToolUseOutput { + universal, + should_block: should_block && invalid_reason.is_none() && invalid_block_reason.is_none(), + reason: wire.reason, + invalid_block_reason, + additional_context, + invalid_reason, + }) +} + +pub(crate) fn parse_pre_compact(stdout: &str) -> Option { + let wire: PreCompactCommandOutputWire = parse_json(stdout)?; + let universal = UniversalOutput::from(wire.universal); + Some(StatelessHookOutput { + universal, + invalid_reason: None, + }) +} + +pub(crate) fn parse_post_compact(stdout: &str) -> Option { + let wire: PostCompactCommandOutputWire = parse_json(stdout)?; + let universal = UniversalOutput::from(wire.universal); + Some(StatelessHookOutput { + universal, + invalid_reason: None, + }) +} + +pub(crate) fn parse_interrupt(stdout: &str) -> Option { + let wire: InterruptCommandOutputWire = parse_json(stdout)?; + Some(InterruptOutput { + system_message: wire.system_message, + }) +} + +pub(crate) fn parse_user_prompt_submit(stdout: &str) -> Option { + let wire: UserPromptSubmitCommandOutputWire = parse_json(stdout)?; + let should_block = matches!(wire.decision, Some(BlockDecisionWire::Block)); + let invalid_block_reason = if should_block + && match wire.reason.as_deref() { + Some(reason) => reason.trim().is_empty(), + None => true, + } { + Some(invalid_block_message("UserPromptSubmit")) + } else { + None + }; + let additional_context = wire + .hook_specific_output + .and_then(|output| output.additional_context); + Some(UserPromptSubmitOutput { + universal: UniversalOutput::from(wire.universal), + should_block: should_block && invalid_block_reason.is_none(), + reason: wire.reason, + invalid_block_reason, + additional_context, + }) +} + +pub(crate) fn parse_stop(stdout: &str) -> Option { + let wire: StopCommandOutputWire = parse_json(stdout)?; + Some(stop_output( + wire.universal, + wire.decision, + wire.reason, + "Stop", + )) +} + +pub(crate) fn parse_subagent_stop(stdout: &str) -> Option { + let wire: SubagentStopCommandOutputWire = parse_json(stdout)?; + Some(stop_output( + wire.universal, + wire.decision, + wire.reason, + "SubagentStop", + )) +} + +fn stop_output( + universal: HookUniversalOutputWire, + decision: Option, + reason: Option, + event_name: &str, +) -> StopOutput { + let should_block = matches!(decision, Some(BlockDecisionWire::Block)); + let invalid_block_reason = if should_block + && match reason.as_deref() { + Some(reason) => reason.trim().is_empty(), + None => true, + } { + Some(invalid_block_message(event_name)) + } else { + None + }; + StopOutput { + universal: UniversalOutput::from(universal), + should_block: should_block && invalid_block_reason.is_none(), + reason, + invalid_block_reason, + } +} + +impl From for UniversalOutput { + fn from(value: HookUniversalOutputWire) -> Self { + Self { + continue_processing: value.r#continue, + stop_reason: value.stop_reason, + suppress_output: value.suppress_output, + system_message: value.system_message, + } + } +} + +fn parse_json(stdout: &str) -> Option +where + T: for<'de> serde::Deserialize<'de>, +{ + let trimmed = stdout.trim(); + if trimmed.is_empty() { + return None; + } + let value: serde_json::Value = serde_json::from_str(trimmed).ok()?; + if !value.is_object() { + return None; + } + serde_json::from_value(value).ok() +} + +pub(crate) fn looks_like_json(stdout: &str) -> bool { + let trimmed = stdout.trim_start(); + trimmed.starts_with('{') || trimmed.starts_with('[') +} + +fn invalid_block_message(event_name: &str) -> String { + format!("{event_name} hook returned decision:block without a non-empty reason") +} + +fn unsupported_pre_tool_use_universal(universal: &UniversalOutput) -> Option { + if !universal.continue_processing { + Some("PreToolUse hook returned unsupported continue:false".to_string()) + } else if universal.stop_reason.is_some() { + Some("PreToolUse hook returned unsupported stopReason".to_string()) + } else if universal.suppress_output { + Some("PreToolUse hook returned unsupported suppressOutput".to_string()) + } else { + None + } +} + +fn unsupported_permission_request_universal(universal: &UniversalOutput) -> Option { + if !universal.continue_processing { + Some("PermissionRequest hook returned unsupported continue:false".to_string()) + } else if universal.stop_reason.is_some() { + Some("PermissionRequest hook returned unsupported stopReason".to_string()) + } else if universal.suppress_output { + Some("PermissionRequest hook returned unsupported suppressOutput".to_string()) + } else { + None + } +} + +fn unsupported_post_tool_use_universal(universal: &UniversalOutput) -> Option { + if universal.suppress_output { + Some("PostToolUse hook returned unsupported suppressOutput".to_string()) + } else { + None + } +} + +fn unsupported_permission_request_hook_specific_output( + decision: Option<&PermissionRequestDecisionWire>, +) -> Option { + let decision = decision?; + if decision.updated_input.is_some() { + Some("PermissionRequest hook returned unsupported updatedInput".to_string()) + } else if decision.updated_permissions.is_some() { + Some("PermissionRequest hook returned unsupported updatedPermissions".to_string()) + } else if decision.interrupt { + Some("PermissionRequest hook returned unsupported interrupt:true".to_string()) + } else { + None + } +} + +fn permission_request_decision( + decision: &PermissionRequestDecisionWire, +) -> PermissionRequestDecision { + match decision.behavior { + PermissionRequestBehaviorWire::Allow => PermissionRequestDecision::Allow, + PermissionRequestBehaviorWire::Deny => PermissionRequestDecision::Deny { + message: decision + .message + .as_deref() + .and_then(trimmed_reason) + .unwrap_or_else(|| "PermissionRequest hook denied approval".to_string()), + }, + } +} + +fn unsupported_post_tool_use_hook_specific_output( + output: &crate::schema::PostToolUseHookSpecificOutputWire, +) -> Option { + if output.updated_mcp_tool_output.is_some() { + Some("PostToolUse hook returned unsupported updatedMCPToolOutput".to_string()) + } else { + None + } +} + +fn unsupported_pre_tool_use_hook_specific_output( + output: &crate::schema::PreToolUseHookSpecificOutputWire, +) -> Option { + if output.updated_input.is_some() + && !matches!( + output.permission_decision, + Some(PreToolUsePermissionDecisionWire::Allow) + ) + { + Some("PreToolUse hook returned updatedInput without permissionDecision:allow".to_string()) + } else { + match output.permission_decision { + Some(PreToolUsePermissionDecisionWire::Allow) => { + output.updated_input.is_none().then(|| { + "PreToolUse hook returned unsupported permissionDecision:allow".to_string() + }) + } + Some(PreToolUsePermissionDecisionWire::Ask) => { + Some("PreToolUse hook returned unsupported permissionDecision:ask".to_string()) + } + Some(PreToolUsePermissionDecisionWire::Deny) => { + if output + .permission_decision_reason + .as_deref() + .and_then(trimmed_reason) + .is_none() + { + Some(invalid_pre_tool_use_reason_message()) + } else { + None + } + } + None => { + if output.permission_decision_reason.is_some() { + Some("PreToolUse hook returned permissionDecisionReason without permissionDecision".to_string()) + } else { + None + } + } + } + } +} + +fn unsupported_pre_tool_use_legacy_decision( + decision: Option<&PreToolUseDecisionWire>, + reason: Option<&str>, +) -> Option { + match decision { + Some(PreToolUseDecisionWire::Approve) => { + Some("PreToolUse hook returned unsupported decision:approve".to_string()) + } + Some(PreToolUseDecisionWire::Block) => { + if reason.and_then(trimmed_reason).is_none() { + Some(invalid_block_message("PreToolUse")) + } else { + None + } + } + None => { + if reason.is_some() { + Some("PreToolUse hook returned reason without decision".to_string()) + } else { + None + } + } + } +} + +fn invalid_pre_tool_use_reason_message() -> String { + "PreToolUse hook returned permissionDecision:deny without a non-empty permissionDecisionReason" + .to_string() +} + +fn trimmed_reason(reason: &str) -> Option { + let trimmed = reason.trim(); + if trimmed.is_empty() { + None + } else { + Some(trimmed.to_string()) + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + use serde_json::json; + + use super::parse_permission_request; + use super::parse_user_prompt_submit; + + #[test] + fn structured_output_rejects_invalid_shapes_and_types() { + for stdout in [ + "[]", + r#"{"systemMessage":123}"#, + r#"{"hookSpecificOutput":{"hookEventName":"UserPromptSubmit","additionalContext":123}}"#, + r#"{"hookSpecificOutput":{"additionalContext":"missing event name"}}"#, + r#"{"hookSpecificOutput":{"hookEventName":"UserPromptSubmit","updatedInput":{}}}"#, + r#"{"unexpectedField":true}"#, + "{", + ] { + assert!( + parse_user_prompt_submit(stdout).is_none(), + "invalid structured output should fail: {stdout}", + ); + } + } + + #[test] + fn permission_request_rejects_reserved_updated_input_field() { + let parsed = parse_permission_request( + &json!({ + "continue": true, + "hookSpecificOutput": { + "hookEventName": "PermissionRequest", + "decision": { + "behavior": "allow", + "updatedInput": {} + } + } + }) + .to_string(), + ) + .expect("permission request hook output should parse"); + + assert_eq!( + parsed.invalid_reason, + Some("PermissionRequest hook returned unsupported updatedInput".to_string()) + ); + } + + #[test] + fn permission_request_rejects_reserved_updated_permissions_field() { + let parsed = parse_permission_request( + &json!({ + "continue": true, + "hookSpecificOutput": { + "hookEventName": "PermissionRequest", + "decision": { + "behavior": "allow", + "updatedPermissions": {} + } + } + }) + .to_string(), + ) + .expect("permission request hook output should parse"); + + assert_eq!( + parsed.invalid_reason, + Some("PermissionRequest hook returned unsupported updatedPermissions".to_string()) + ); + } + + #[test] + fn permission_request_rejects_reserved_interrupt_field() { + let parsed = parse_permission_request( + &json!({ + "continue": true, + "hookSpecificOutput": { + "hookEventName": "PermissionRequest", + "decision": { + "behavior": "allow", + "interrupt": true + } + } + }) + .to_string(), + ) + .expect("permission request hook output should parse"); + + assert_eq!( + parsed.invalid_reason, + Some("PermissionRequest hook returned unsupported interrupt:true".to_string()) + ); + } +} diff --git a/codex-rs/hooks/src/engine/schema_loader.rs b/codex-rs/hooks/src/engine/schema_loader.rs new file mode 100644 index 0000000000000000000000000000000000000000..7fce9a3c435ca0a43a4798c77d7bd61399e34845 --- /dev/null +++ b/codex-rs/hooks/src/engine/schema_loader.rs @@ -0,0 +1,168 @@ +use std::sync::OnceLock; + +use serde_json::Value; + +#[allow(dead_code)] +pub(crate) struct GeneratedHookSchemas { + pub post_tool_use_command_input: Value, + pub post_tool_use_command_output: Value, + pub permission_request_command_input: Value, + pub permission_request_command_output: Value, + pub post_compact_command_input: Value, + pub post_compact_command_output: Value, + pub pre_tool_use_command_input: Value, + pub pre_tool_use_command_output: Value, + pub pre_compact_command_input: Value, + pub pre_compact_command_output: Value, + pub session_start_command_input: Value, + pub session_start_command_output: Value, + pub session_end_command_input: Value, + pub subagent_start_command_input: Value, + pub subagent_start_command_output: Value, + pub subagent_stop_command_input: Value, + pub subagent_stop_command_output: Value, + pub user_prompt_submit_command_input: Value, + pub user_prompt_submit_command_output: Value, + pub stop_command_input: Value, + pub stop_command_output: Value, + pub interrupt_command_input: Value, + pub interrupt_command_output: Value, +} + +pub(crate) fn generated_hook_schemas() -> &'static GeneratedHookSchemas { + static SCHEMAS: OnceLock = OnceLock::new(); + SCHEMAS.get_or_init(|| GeneratedHookSchemas { + post_tool_use_command_input: parse_json_schema( + "post-tool-use.command.input", + include_str!("../../schema/generated/post-tool-use.command.input.schema.json"), + ), + post_tool_use_command_output: parse_json_schema( + "post-tool-use.command.output", + include_str!("../../schema/generated/post-tool-use.command.output.schema.json"), + ), + permission_request_command_input: parse_json_schema( + "permission-request.command.input", + include_str!("../../schema/generated/permission-request.command.input.schema.json"), + ), + permission_request_command_output: parse_json_schema( + "permission-request.command.output", + include_str!("../../schema/generated/permission-request.command.output.schema.json"), + ), + post_compact_command_input: parse_json_schema( + "post-compact.command.input", + include_str!("../../schema/generated/post-compact.command.input.schema.json"), + ), + post_compact_command_output: parse_json_schema( + "post-compact.command.output", + include_str!("../../schema/generated/post-compact.command.output.schema.json"), + ), + pre_tool_use_command_input: parse_json_schema( + "pre-tool-use.command.input", + include_str!("../../schema/generated/pre-tool-use.command.input.schema.json"), + ), + pre_tool_use_command_output: parse_json_schema( + "pre-tool-use.command.output", + include_str!("../../schema/generated/pre-tool-use.command.output.schema.json"), + ), + pre_compact_command_input: parse_json_schema( + "pre-compact.command.input", + include_str!("../../schema/generated/pre-compact.command.input.schema.json"), + ), + pre_compact_command_output: parse_json_schema( + "pre-compact.command.output", + include_str!("../../schema/generated/pre-compact.command.output.schema.json"), + ), + session_start_command_input: parse_json_schema( + "session-start.command.input", + include_str!("../../schema/generated/session-start.command.input.schema.json"), + ), + session_start_command_output: parse_json_schema( + "session-start.command.output", + include_str!("../../schema/generated/session-start.command.output.schema.json"), + ), + session_end_command_input: parse_json_schema( + "session-end.command.input", + include_str!("../../schema/generated/session-end.command.input.schema.json"), + ), + subagent_start_command_input: parse_json_schema( + "subagent-start.command.input", + include_str!("../../schema/generated/subagent-start.command.input.schema.json"), + ), + subagent_start_command_output: parse_json_schema( + "subagent-start.command.output", + include_str!("../../schema/generated/subagent-start.command.output.schema.json"), + ), + subagent_stop_command_input: parse_json_schema( + "subagent-stop.command.input", + include_str!("../../schema/generated/subagent-stop.command.input.schema.json"), + ), + subagent_stop_command_output: parse_json_schema( + "subagent-stop.command.output", + include_str!("../../schema/generated/subagent-stop.command.output.schema.json"), + ), + user_prompt_submit_command_input: parse_json_schema( + "user-prompt-submit.command.input", + include_str!("../../schema/generated/user-prompt-submit.command.input.schema.json"), + ), + user_prompt_submit_command_output: parse_json_schema( + "user-prompt-submit.command.output", + include_str!("../../schema/generated/user-prompt-submit.command.output.schema.json"), + ), + stop_command_input: parse_json_schema( + "stop.command.input", + include_str!("../../schema/generated/stop.command.input.schema.json"), + ), + stop_command_output: parse_json_schema( + "stop.command.output", + include_str!("../../schema/generated/stop.command.output.schema.json"), + ), + interrupt_command_input: parse_json_schema( + "interrupt.command.input", + include_str!("../../schema/generated/interrupt.command.input.schema.json"), + ), + interrupt_command_output: parse_json_schema( + "interrupt.command.output", + include_str!("../../schema/generated/interrupt.command.output.schema.json"), + ), + }) +} + +fn parse_json_schema(name: &str, schema: &str) -> Value { + serde_json::from_str(schema) + .unwrap_or_else(|err| panic!("invalid generated hooks schema {name}: {err}")) +} + +#[cfg(test)] +mod tests { + use super::generated_hook_schemas; + use pretty_assertions::assert_eq; + + #[test] + fn loads_generated_hook_schemas() { + let schemas = generated_hook_schemas(); + + assert_eq!(schemas.post_tool_use_command_input["type"], "object"); + assert_eq!(schemas.post_tool_use_command_output["type"], "object"); + assert_eq!(schemas.permission_request_command_input["type"], "object"); + assert_eq!(schemas.permission_request_command_output["type"], "object"); + assert_eq!(schemas.post_compact_command_input["type"], "object"); + assert_eq!(schemas.post_compact_command_output["type"], "object"); + assert_eq!(schemas.pre_tool_use_command_input["type"], "object"); + assert_eq!(schemas.pre_tool_use_command_output["type"], "object"); + assert_eq!(schemas.pre_compact_command_input["type"], "object"); + assert_eq!(schemas.pre_compact_command_output["type"], "object"); + assert_eq!(schemas.session_start_command_input["type"], "object"); + assert_eq!(schemas.session_start_command_output["type"], "object"); + assert_eq!(schemas.session_end_command_input["type"], "object"); + assert_eq!(schemas.subagent_start_command_input["type"], "object"); + assert_eq!(schemas.subagent_start_command_output["type"], "object"); + assert_eq!(schemas.subagent_stop_command_input["type"], "object"); + assert_eq!(schemas.subagent_stop_command_output["type"], "object"); + assert_eq!(schemas.user_prompt_submit_command_input["type"], "object"); + assert_eq!(schemas.user_prompt_submit_command_output["type"], "object"); + assert_eq!(schemas.stop_command_input["type"], "object"); + assert_eq!(schemas.stop_command_output["type"], "object"); + assert_eq!(schemas.interrupt_command_input["type"], "object"); + assert_eq!(schemas.interrupt_command_output["type"], "object"); + } +} diff --git a/codex-rs/hooks/src/events/common.rs b/codex-rs/hooks/src/events/common.rs new file mode 100644 index 0000000000000000000000000000000000000000..7b4c4dcab31e9e5747594c39ec582fb0c49f4095 --- /dev/null +++ b/codex-rs/hooks/src/events/common.rs @@ -0,0 +1,306 @@ +use codex_protocol::protocol::HookCompletedEvent; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookOutputEntry; +use codex_protocol::protocol::HookOutputEntryKind; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookRunSummary; + +use crate::engine::ConfiguredHandler; +use crate::engine::HandlerSourcePath; +use crate::engine::dispatcher; +use crate::output_spill::AdditionalContext; + +/// Identifies a thread-spawned subagent when a normal hook runs inside it. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SubagentHookContext { + pub agent_id: String, + pub agent_type: String, +} + +pub(crate) fn join_text_chunks(chunks: Vec) -> Option { + if chunks.is_empty() { + None + } else { + Some(chunks.join("\n\n")) + } +} + +pub(crate) fn trimmed_non_empty(text: &str) -> Option { + let trimmed = text.trim(); + if trimmed.is_empty() { + None + } else { + Some(trimmed.to_string()) + } +} + +pub(crate) fn append_additional_context( + entries: &mut Vec, + additional_contexts_for_model: &mut Vec, + handler: &ConfiguredHandler, + additional_context: String, +) { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Context, + text: additional_context.clone(), + }); + additional_contexts_for_model.push(AdditionalContext { + text: additional_context, + limit: handler.additional_context_limit, + }); +} + +pub(crate) fn flatten_additional_contexts<'a>( + additional_contexts: impl IntoIterator, +) -> Vec { + additional_contexts + .into_iter() + .flat_map(|chunk| chunk.iter().cloned()) + .collect() +} + +pub(crate) fn serialization_failure_hook_events( + handlers: Vec, + turn_id: Option, + error_message: String, +) -> Vec { + handlers + .into_iter() + .filter(|handler| matches!(handler.source_path, HandlerSourcePath::Local(_))) + .map(|handler| { + let mut run = dispatcher::running_summary(&handler); + run.status = HookRunStatus::Failed; + run.completed_at = Some(run.started_at); + run.duration_ms = Some(0); + run.entries = vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: error_message.clone(), + }]; + HookCompletedEvent { + turn_id: turn_id.clone(), + run, + } + }) + .collect() +} + +pub(crate) fn serialization_failure_hook_events_for_tool_use( + handlers: Vec, + turn_id: Option, + error_message: String, + tool_use_id: &str, +) -> Vec { + serialization_failure_hook_events(handlers, turn_id, error_message) + .into_iter() + .map(|event| hook_completed_for_tool_use(event, tool_use_id)) + .collect() +} + +pub(crate) fn hook_completed_for_tool_use( + mut event: HookCompletedEvent, + tool_use_id: &str, +) -> HookCompletedEvent { + event.run = hook_run_for_tool_use(event.run, tool_use_id); + event +} + +pub(crate) fn hook_run_for_tool_use(mut run: HookRunSummary, tool_use_id: &str) -> HookRunSummary { + run.id = format!("{}:{tool_use_id}", run.id); + run +} + +pub(crate) fn matcher_pattern_for_event( + event_name: HookEventName, + matcher: Option<&str>, +) -> Option<&str> { + match event_name { + HookEventName::PreToolUse + | HookEventName::PermissionRequest + | HookEventName::PostToolUse + | HookEventName::SessionStart + | HookEventName::SessionEnd + | HookEventName::SubagentStart + | HookEventName::SubagentStop + | HookEventName::PreCompact + | HookEventName::PostCompact => matcher, + HookEventName::UserPromptSubmit | HookEventName::Stop | HookEventName::Interrupt => None, + } +} + +pub(crate) fn validate_matcher_pattern(matcher: &str) -> Result<(), regex::Error> { + if is_match_all_matcher(matcher) || is_exact_matcher(matcher) { + return Ok(()); + } + regex::Regex::new(matcher).map(|_| ()) +} + +pub(crate) fn matches_matcher(matcher: Option<&str>, input: Option<&str>) -> bool { + match matcher { + None => true, + Some(matcher) if is_match_all_matcher(matcher) => true, + Some(matcher) if is_exact_matcher(matcher) => input + .map(|input| matcher.split('|').any(|candidate| candidate == input)) + .unwrap_or(false), + Some(matcher) => input + .and_then(|input| { + regex::Regex::new(matcher) + .ok() + .map(|regex| regex.is_match(input)) + }) + .unwrap_or(false), + } +} + +pub(crate) fn matcher_inputs<'a>( + tool_name: &'a str, + matcher_aliases: &'a [String], +) -> Vec<&'a str> { + // Keep the canonical name first so matcher previews and execution preserve + // the same primary identity that hook stdin will serialize. + std::iter::once(tool_name) + .chain(matcher_aliases.iter().map(String::as_str)) + .collect() +} + +fn is_match_all_matcher(matcher: &str) -> bool { + matcher.is_empty() || matcher == "*" +} + +fn is_exact_matcher(matcher: &str) -> bool { + matcher + .chars() + .all(|ch| ch.is_ascii_alphanumeric() || ch == '_' || ch == '|') +} + +#[cfg(test)] +mod tests { + use codex_protocol::protocol::HookEventName; + use pretty_assertions::assert_eq; + + use super::matcher_pattern_for_event; + use super::matches_matcher; + use super::validate_matcher_pattern; + + #[test] + fn matcher_omitted_matches_all_occurrences() { + assert!(matches_matcher(/*matcher*/ None, Some("Bash"))); + assert!(matches_matcher(/*matcher*/ None, Some("Write"))); + } + + #[test] + fn matcher_star_matches_all_occurrences() { + assert!(matches_matcher(Some("*"), Some("Bash"))); + assert!(matches_matcher(Some("*"), Some("Edit"))); + assert_eq!(validate_matcher_pattern("*"), Ok(())); + } + + #[test] + fn matcher_empty_string_matches_all_occurrences() { + assert!(matches_matcher(Some(""), Some("Bash"))); + assert!(matches_matcher(Some(""), Some("SessionStart"))); + assert_eq!(validate_matcher_pattern(""), Ok(())); + } + + #[test] + fn exact_matcher_supports_pipe_alternatives() { + assert!(matches_matcher(Some("Edit|Write"), Some("Edit"))); + assert!(matches_matcher(Some("Edit|Write"), Some("Write"))); + assert!(!matches_matcher(Some("Edit|Write"), Some("Bash"))); + assert_eq!(validate_matcher_pattern("Edit|Write"), Ok(())); + } + + #[test] + fn literal_matcher_uses_exact_matching() { + assert!(matches_matcher(Some("Bash"), Some("Bash"))); + assert!(!matches_matcher(Some("Bash"), Some("BashOutput"))); + assert!(matches_matcher( + Some("mcp__memory__create_entities"), + Some("mcp__memory__create_entities") + )); + assert!(!matches_matcher( + Some("mcp__memory"), + Some("mcp__memory__create_entities") + )); + assert_eq!(validate_matcher_pattern("mcp__memory"), Ok(())); + } + + #[test] + fn matcher_uses_regex_when_it_contains_regex_characters() { + assert!(matches_matcher(Some("^Bash"), Some("BashOutput"))); + assert_eq!(validate_matcher_pattern("^Bash"), Ok(())); + } + + #[test] + fn mcp_matchers_support_regex_wildcards() { + assert!(matches_matcher( + Some("mcp__memory__.*"), + Some("mcp__memory__create_entities") + )); + assert!(matches_matcher( + Some("mcp__.*__write.*"), + Some("mcp__filesystem__write_file") + )); + assert!(!matches_matcher( + Some("mcp__.*__write.*"), + Some("mcp__filesystem__read_file") + )); + assert_eq!(validate_matcher_pattern("mcp__memory__.*"), Ok(())); + } + + #[test] + fn matcher_supports_anchored_regexes() { + assert!(matches_matcher(Some("^Bash$"), Some("Bash"))); + assert!(!matches_matcher(Some("^Bash$"), Some("BashOutput"))); + assert_eq!(validate_matcher_pattern("^Bash$"), Ok(())); + } + + #[test] + fn invalid_regex_is_rejected() { + assert!(validate_matcher_pattern("[").is_err()); + assert!(!matches_matcher(Some("["), Some("Bash"))); + } + + #[test] + fn unsupported_events_ignore_matchers() { + assert_eq!( + matcher_pattern_for_event(HookEventName::UserPromptSubmit, Some("^hello")), + None + ); + assert_eq!( + matcher_pattern_for_event(HookEventName::Stop, Some("^done$")), + None + ); + assert_eq!( + matcher_pattern_for_event(HookEventName::Interrupt, Some("^interrupted$")), + None + ); + } + + #[test] + fn supported_events_keep_matchers() { + assert_eq!( + matcher_pattern_for_event(HookEventName::PreToolUse, Some("Bash")), + Some("Bash") + ); + assert_eq!( + matcher_pattern_for_event(HookEventName::PostToolUse, Some("Edit|Write")), + Some("Edit|Write") + ); + assert_eq!( + matcher_pattern_for_event(HookEventName::SessionStart, Some("startup|resume")), + Some("startup|resume") + ); + assert_eq!( + matcher_pattern_for_event(HookEventName::SessionEnd, Some("clear|other")), + Some("clear|other") + ); + assert_eq!( + matcher_pattern_for_event(HookEventName::PreCompact, Some("^auto$")), + Some("^auto$") + ); + assert_eq!( + matcher_pattern_for_event(HookEventName::PostCompact, Some("manual|auto")), + Some("manual|auto") + ); + } +} diff --git a/codex-rs/hooks/src/events/compact.rs b/codex-rs/hooks/src/events/compact.rs new file mode 100644 index 0000000000000000000000000000000000000000..82bd6c6cf55d1ce66755c7c9aeb3e4aba05de168 --- /dev/null +++ b/codex-rs/hooks/src/events/compact.rs @@ -0,0 +1,554 @@ +use std::path::PathBuf; + +use codex_protocol::ThreadId; +use codex_protocol::protocol::HookCompletedEvent; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookOutputEntry; +use codex_protocol::protocol::HookOutputEntryKind; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookRunSummary; +use codex_utils_absolute_path::AbsolutePathBuf; + +use super::common; +use crate::engine::ClaudeHooksEngine; +use crate::engine::ConfiguredHandler; +use crate::engine::HandlerRunResult; +use crate::engine::dispatcher; +use crate::engine::output_parser; +use crate::schema::PostCompactCommandInput; +use crate::schema::PreCompactCommandInput; +use crate::schema::SubagentCommandInputFields; + +#[derive(Debug, Clone)] +pub struct PreCompactRequest { + pub session_id: ThreadId, + pub turn_id: String, + pub subagent: Option, + pub cwd: AbsolutePathBuf, + pub transcript_path: Option, + pub model: String, + pub trigger: String, +} + +#[derive(Debug, Clone)] +pub struct PostCompactRequest { + pub session_id: ThreadId, + pub turn_id: String, + pub subagent: Option, + pub cwd: AbsolutePathBuf, + pub transcript_path: Option, + pub model: String, + pub trigger: String, +} + +#[derive(Debug)] +pub struct StatelessHookOutcome { + pub hook_events: Vec, + pub should_stop: bool, + pub stop_reason: Option, +} + +#[derive(Debug)] +pub struct PreCompactOutcome { + pub hook_events: Vec, + pub should_stop: bool, + pub stop_reason: Option, +} + +pub(crate) fn preview_pre( + handlers: &[ConfiguredHandler], + request: &PreCompactRequest, +) -> Vec { + dispatcher::select_handlers( + handlers, + HookEventName::PreCompact, + Some(request.trigger.as_str()), + ) + .into_iter() + .map(|handler| dispatcher::running_summary(&handler)) + .collect() +} + +pub(crate) async fn run_pre( + engine: &ClaudeHooksEngine, + request: PreCompactRequest, +) -> PreCompactOutcome { + let matched = dispatcher::select_handlers( + &engine.handlers, + HookEventName::PreCompact, + Some(request.trigger.as_str()), + ); + if matched.is_empty() { + return PreCompactOutcome { + hook_events: Vec::new(), + should_stop: false, + stop_reason: None, + }; + } + + let input_json = match pre_command_input_json(&request) { + Ok(input_json) => input_json, + Err(error) => { + return PreCompactOutcome { + hook_events: common::serialization_failure_hook_events( + matched, + Some(request.turn_id), + format!("failed to serialize pre compact hook input: {error}"), + ), + should_stop: false, + stop_reason: None, + }; + } + }; + + let results = dispatcher::execute_handlers( + engine, + matched, + input_json, + request.cwd.as_path(), + Some(request.turn_id), + parse_pre_completed, + ) + .await; + let should_stop = results.iter().any(|result| result.data.should_stop); + let stop_reason = results + .iter() + .find_map(|result| result.data.stop_reason.clone()); + PreCompactOutcome { + hook_events: results.into_iter().map(|result| result.completed).collect(), + should_stop, + stop_reason, + } +} + +fn pre_command_input_json(request: &PreCompactRequest) -> Result { + let subagent = SubagentCommandInputFields::from(request.subagent.as_ref()); + serde_json::to_string(&PreCompactCommandInput { + session_id: request.session_id.to_string(), + turn_id: request.turn_id.clone(), + agent_id: subagent.agent_id, + agent_type: subagent.agent_type, + transcript_path: crate::schema::NullableString::from_path(request.transcript_path.clone()), + cwd: request.cwd.display().to_string(), + hook_event_name: "PreCompact".to_string(), + model: request.model.clone(), + trigger: request.trigger.clone(), + }) +} + +pub(crate) fn preview_post( + handlers: &[ConfiguredHandler], + request: &PostCompactRequest, +) -> Vec { + dispatcher::select_handlers( + handlers, + HookEventName::PostCompact, + Some(request.trigger.as_str()), + ) + .into_iter() + .map(|handler| dispatcher::running_summary(&handler)) + .collect() +} + +pub(crate) async fn run_post( + engine: &ClaudeHooksEngine, + request: PostCompactRequest, +) -> StatelessHookOutcome { + let matched = dispatcher::select_handlers( + &engine.handlers, + HookEventName::PostCompact, + Some(request.trigger.as_str()), + ); + if matched.is_empty() { + return StatelessHookOutcome { + hook_events: Vec::new(), + should_stop: false, + stop_reason: None, + }; + } + + let input_json = match post_command_input_json(&request) { + Ok(input_json) => input_json, + Err(error) => { + return StatelessHookOutcome { + hook_events: common::serialization_failure_hook_events( + matched, + Some(request.turn_id), + format!("failed to serialize post compact hook input: {error}"), + ), + should_stop: false, + stop_reason: None, + }; + } + }; + + let results = dispatcher::execute_handlers( + engine, + matched, + input_json, + request.cwd.as_path(), + Some(request.turn_id), + parse_post_completed, + ) + .await; + let should_stop = results.iter().any(|result| result.data.should_stop); + let stop_reason = results + .iter() + .find_map(|result| result.data.stop_reason.clone()); + StatelessHookOutcome { + hook_events: results.into_iter().map(|result| result.completed).collect(), + should_stop, + stop_reason, + } +} + +fn post_command_input_json(request: &PostCompactRequest) -> Result { + let subagent = SubagentCommandInputFields::from(request.subagent.as_ref()); + serde_json::to_string(&PostCompactCommandInput { + session_id: request.session_id.to_string(), + turn_id: request.turn_id.clone(), + agent_id: subagent.agent_id, + agent_type: subagent.agent_type, + transcript_path: crate::schema::NullableString::from_path(request.transcript_path.clone()), + cwd: request.cwd.display().to_string(), + hook_event_name: "PostCompact".to_string(), + model: request.model.clone(), + trigger: request.trigger.clone(), + }) +} + +#[derive(Default)] +struct CompactHandlerData { + should_stop: bool, + stop_reason: Option, +} + +fn parse_pre_completed( + handler: &ConfiguredHandler, + run_result: HandlerRunResult, + turn_id: Option, +) -> dispatcher::ParsedHandler { + parse_completed( + handler, + run_result, + turn_id, + "PreCompact", + output_parser::parse_pre_compact, + ) +} + +fn parse_post_completed( + handler: &ConfiguredHandler, + run_result: HandlerRunResult, + turn_id: Option, +) -> dispatcher::ParsedHandler { + parse_completed( + handler, + run_result, + turn_id, + "PostCompact", + output_parser::parse_post_compact, + ) +} + +fn parse_completed( + handler: &ConfiguredHandler, + run_result: HandlerRunResult, + turn_id: Option, + event_label: &'static str, + parse_output: fn(&str) -> Option, +) -> dispatcher::ParsedHandler { + let mut entries = Vec::new(); + let mut status = HookRunStatus::Completed; + let mut should_stop = false; + let mut stop_reason = None; + + match run_result.error.as_deref() { + Some(error) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: error.to_string(), + }); + } + None => match run_result.exit_code { + Some(0) => { + let trimmed_stdout = run_result.stdout.trim(); + if trimmed_stdout.is_empty() { + } else if let Some(parsed) = parse_output(&run_result.stdout) { + if let Some(system_message) = parsed.universal.system_message { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Warning, + text: system_message, + }); + } + let _ = parsed.universal.suppress_output; + if handler.can_apply_control_effects() { + if !parsed.universal.continue_processing { + status = HookRunStatus::Stopped; + should_stop = true; + stop_reason = parsed.universal.stop_reason.clone(); + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Stop, + text: parsed.universal.stop_reason.unwrap_or_else(|| { + format!("{event_label} hook stopped execution") + }), + }); + } else if let Some(invalid_reason) = parsed.invalid_reason { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: invalid_reason, + }); + } + } + } else if output_parser::looks_like_json(&run_result.stdout) { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: format!("hook returned invalid {event_label} hook JSON output"), + }); + } + } + Some(code) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: common::trimmed_non_empty(&run_result.stderr) + .unwrap_or_else(|| format!("hook exited with code {code}")), + }); + } + None => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook process terminated without an exit code".to_string(), + }); + } + }, + } + + dispatcher::ParsedHandler { + completed: HookCompletedEvent { + turn_id, + run: dispatcher::completed_summary(handler, &run_result, status, entries), + }, + data: CompactHandlerData { + should_stop, + stop_reason, + }, + completion_order: 0, + } +} + +#[cfg(test)] +mod tests { + use codex_protocol::ThreadId; + use codex_protocol::protocol::HookEventName; + use codex_protocol::protocol::HookOutputEntry; + use codex_protocol::protocol::HookOutputEntryKind; + use codex_protocol::protocol::HookRunStatus; + use codex_utils_absolute_path::test_support::PathBufExt; + use codex_utils_absolute_path::test_support::test_path_buf; + use pretty_assertions::assert_eq; + use serde_json::json; + + use super::parse_post_completed; + use super::parse_pre_completed; + use super::post_command_input_json; + use super::pre_command_input_json; + use crate::engine::ConfiguredHandler; + use crate::engine::HandlerRunResult; + + #[test] + fn pre_compact_input_includes_lifecycle_metadata() { + let request = pre_request(); + let input_json = pre_command_input_json(&request).expect("serialize command input"); + let input: serde_json::Value = + serde_json::from_str(&input_json).expect("parse command input"); + + assert_eq!( + input, + json!({ + "session_id": pre_request().session_id.to_string(), + "turn_id": "turn-1", + "transcript_path": null, + "cwd": test_path_buf("/tmp").display().to_string(), + "hook_event_name": "PreCompact", + "model": "gpt-test", + "trigger": "manual", + }) + ); + } + + #[test] + fn post_compact_input_includes_lifecycle_metadata() { + let request = post_request(); + let input_json = post_command_input_json(&request).expect("serialize command input"); + let input: serde_json::Value = + serde_json::from_str(&input_json).expect("parse command input"); + + assert_eq!( + input, + json!({ + "session_id": post_request().session_id.to_string(), + "turn_id": "turn-1", + "transcript_path": null, + "cwd": test_path_buf("/tmp").display().to_string(), + "hook_event_name": "PostCompact", + "model": "gpt-test", + "trigger": "manual", + }) + ); + } + + #[test] + fn block_decision_is_not_supported_for_pre_compact() { + let parsed = parse_pre_completed( + &handler(HookEventName::PreCompact), + run_result( + Some(0), + r#"{"decision":"block","reason":"policy blocked compaction"}"#, + "", + ), + Some("turn-1".to_string()), + ); + + assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook returned invalid PreCompact hook JSON output".to_string(), + }] + ); + } + + #[test] + fn continue_false_stops_before_compaction() { + let parsed = parse_pre_completed( + &handler(HookEventName::PreCompact), + run_result(Some(0), r#"{"continue":false,"stopReason":"nope"}"#, ""), + Some("turn-1".to_string()), + ); + + assert_eq!(parsed.completed.run.status, HookRunStatus::Stopped); + assert_eq!(parsed.data.should_stop, true); + assert_eq!(parsed.data.stop_reason, Some("nope".to_string())); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Stop, + text: "nope".to_string(), + }] + ); + } + + #[test] + fn post_compact_continue_false_stops_after_compaction() { + let parsed = parse_post_completed( + &handler(HookEventName::PostCompact), + run_result( + Some(0), + r#"{"continue":false,"stopReason":"pause after compact"}"#, + "", + ), + Some("turn-1".to_string()), + ); + + assert_eq!(parsed.completed.run.status, HookRunStatus::Stopped); + assert_eq!(parsed.data.should_stop, true); + assert_eq!( + parsed.data.stop_reason, + Some("pause after compact".to_string()) + ); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Stop, + text: "pause after compact".to_string(), + }] + ); + } + + #[test] + fn pre_compact_ignores_plain_stdout() { + let parsed = parse_pre_completed( + &handler(HookEventName::PreCompact), + run_result(Some(0), "checking compact policy\n", ""), + Some("turn-1".to_string()), + ); + + assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); + assert_eq!(parsed.completed.run.entries, Vec::new()); + } + + #[test] + fn post_compact_ignores_plain_stdout() { + let parsed = parse_post_completed( + &handler(HookEventName::PostCompact), + run_result(Some(0), "logged compact summary\n", ""), + Some("turn-1".to_string()), + ); + + assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); + assert_eq!(parsed.completed.run.entries, Vec::new()); + } + + fn pre_request() -> super::PreCompactRequest { + super::PreCompactRequest { + session_id: ThreadId::from_string("00000000-0000-4000-8000-000000000001") + .expect("valid thread id"), + turn_id: "turn-1".to_string(), + subagent: None, + cwd: test_path_buf("/tmp").abs(), + transcript_path: None, + model: "gpt-test".to_string(), + trigger: "manual".to_string(), + } + } + + fn post_request() -> super::PostCompactRequest { + super::PostCompactRequest { + session_id: ThreadId::from_string("00000000-0000-4000-8000-000000000002") + .expect("valid thread id"), + turn_id: "turn-1".to_string(), + subagent: None, + cwd: test_path_buf("/tmp").abs(), + transcript_path: None, + model: "gpt-test".to_string(), + trigger: "manual".to_string(), + } + } + + fn handler(event_name: HookEventName) -> ConfiguredHandler { + ConfiguredHandler { + builtin: false, + event_name, + matcher: None, + timeout_sec: 5, + status_message: Some("running compact hook".to_string()), + additional_context_limit: Default::default(), + source_path: test_path_buf("/tmp/hooks.json").abs().into(), + source: codex_protocol::protocol::HookSource::User, + display_order: 0, + kind: crate::engine::ConfiguredHandlerKind::Command { + command: "python3 compact_hook.py".to_string(), + r#async: false, + env: std::collections::HashMap::new(), + }, + } + } + + fn run_result(exit_code: Option, stdout: &str, stderr: &str) -> HandlerRunResult { + HandlerRunResult { + started_at: 1_700_000_000, + completed_at: 1_700_000_001, + duration_ms: 12, + exit_code, + stdout: stdout.to_string(), + stderr: stderr.to_string(), + error: None, + } + } +} diff --git a/codex-rs/hooks/src/events/interrupt.rs b/codex-rs/hooks/src/events/interrupt.rs new file mode 100644 index 0000000000000000000000000000000000000000..743002eb9ca1b457132f466a9cddd964e8cd5942 --- /dev/null +++ b/codex-rs/hooks/src/events/interrupt.rs @@ -0,0 +1,183 @@ +use std::path::PathBuf; + +use codex_protocol::ThreadId; +use codex_protocol::protocol::HookCompletedEvent; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookOutputEntry; +use codex_protocol::protocol::HookOutputEntryKind; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookRunSummary; +use codex_utils_absolute_path::AbsolutePathBuf; +use serde_json::Map; +use serde_json::Value; + +use super::common; +use crate::engine::ClaudeHooksEngine; +use crate::engine::ConfiguredHandler; +use crate::engine::HandlerRunResult; +use crate::engine::HandlerSourcePath; +use crate::engine::dispatcher; +use crate::engine::output_parser; +use crate::schema::InterruptCommandInput; +use crate::schema::NullableString; + +#[derive(Debug, Clone)] +pub struct InterruptRequest { + pub session_id: ThreadId, + pub turn_id: String, + pub cwd: AbsolutePathBuf, + pub transcript_path: Option, + pub model: String, + pub permission_mode: String, + pub request_metadata: Option>, +} + +#[derive(Debug, Default)] +pub struct InterruptOutcome { + pub hook_events: Vec, +} + +#[derive(Debug, Default, PartialEq, Eq)] +struct InterruptHandlerData; + +pub(crate) fn preview(handlers: &[ConfiguredHandler]) -> Vec { + dispatcher::select_handlers( + handlers, + HookEventName::Interrupt, + /*matcher_input*/ None, + ) + .into_iter() + // Executor-scoped hooks run asynchronously and do not emit public hook lifecycle events. + .filter(|handler| matches!(handler.source_path, HandlerSourcePath::Local(_))) + .map(|handler| dispatcher::running_summary(&handler)) + .collect() +} + +pub(crate) async fn run(engine: &ClaudeHooksEngine, request: InterruptRequest) -> InterruptOutcome { + let matched = dispatcher::select_handlers( + &engine.handlers, + HookEventName::Interrupt, + /*matcher_input*/ None, + ); + if matched.is_empty() { + return InterruptOutcome::default(); + } + + let InterruptRequest { + session_id, + turn_id, + cwd, + transcript_path, + model, + permission_mode, + request_metadata, + } = request; + let input_json = match serde_json::to_string(&InterruptCommandInput { + session_id: session_id.to_string(), + turn_id: turn_id.clone(), + transcript_path: NullableString::from_path(transcript_path), + cwd: cwd.display().to_string(), + hook_event_name: "Interrupt".to_string(), + model, + permission_mode, + }) { + Ok(input_json) => input_json, + Err(error) => { + return InterruptOutcome { + hook_events: common::serialization_failure_hook_events( + matched, + Some(turn_id), + format!("failed to serialize interrupt hook input: {error}"), + ), + }; + } + }; + + let results = dispatcher::execute_handlers_with_metadata( + engine, + matched, + input_json, + cwd.as_path(), + Some(turn_id), + request_metadata.as_ref(), + parse_completed, + ) + .await; + + InterruptOutcome { + hook_events: results.into_iter().map(|result| result.completed).collect(), + } +} + +fn parse_completed( + handler: &ConfiguredHandler, + run_result: HandlerRunResult, + turn_id: Option, +) -> dispatcher::ParsedHandler { + let mut entries = Vec::new(); + let mut status = HookRunStatus::Completed; + + match run_result.error.as_deref() { + Some(error) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: error.to_string(), + }); + } + None => match run_result.exit_code { + Some(0) => { + let trimmed_stdout = run_result.stdout.trim(); + if trimmed_stdout.is_empty() { + } else if let Some(parsed) = output_parser::parse_interrupt(&run_result.stdout) { + if let Some(system_message) = parsed.system_message { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Warning, + text: system_message, + }); + } + } else { + status = HookRunStatus::Failed; + let text = if output_parser::looks_like_json(&run_result.stdout) { + "hook returned invalid interrupt hook JSON output" + } else { + "Interrupt hook returned non-JSON stdout" + }; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: text.to_string(), + }); + } + } + Some(exit_code) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: format!("hook exited with code {exit_code}"), + }); + } + None => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook exited without a status code".to_string(), + }); + } + }, + } + + let completed = HookCompletedEvent { + turn_id, + run: dispatcher::completed_summary(handler, &run_result, status, entries), + }; + + dispatcher::ParsedHandler { + completed, + data: InterruptHandlerData, + completion_order: 0, + } +} + +#[cfg(test)] +#[path = "interrupt_tests.rs"] +mod tests; diff --git a/codex-rs/hooks/src/events/interrupt_tests.rs b/codex-rs/hooks/src/events/interrupt_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..3d8642c8457bc86dfc4ffba9ffd37d0db3bddc83 --- /dev/null +++ b/codex-rs/hooks/src/events/interrupt_tests.rs @@ -0,0 +1,128 @@ +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookOutputEntry; +use codex_protocol::protocol::HookOutputEntryKind; +use codex_protocol::protocol::HookRunStatus; +use codex_utils_absolute_path::test_support::PathBufExt; +use codex_utils_absolute_path::test_support::test_path_buf; +use pretty_assertions::assert_eq; + +use super::InterruptHandlerData; +use super::parse_completed; +use crate::engine::ConfiguredHandler; +use crate::engine::ConfiguredHandlerKind; +use crate::engine::HandlerRunResult; + +#[test] +fn empty_stdout_succeeds() { + let parsed = parse_completed( + &handler(), + run_result(Some(0), "", ""), + /*turn_id*/ None, + ); + + assert_eq!(parsed.data, InterruptHandlerData); + assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); + assert!(parsed.completed.run.entries.is_empty()); +} + +#[test] +fn system_message_becomes_warning() { + let parsed = parse(r#"{"systemMessage":"watch the tide"}"#, Some(0)); + + assert_eq!(parsed.data, InterruptHandlerData); + assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); + assert_eq!( + parsed.completed.run.entries, + vec![warning("watch the tide")] + ); +} + +#[test] +fn invalid_json_outputs_fail() { + for stdout in [ + r#"{"continue":true}"#, + r#"{"stopReason":null}"#, + r#"{"suppressOutput":false}"#, + r#"{"decision":"block"}"#, + r#"{"systemMessage":"watch the tide""#, + ] { + assert_failed( + stdout, + Some(0), + "hook returned invalid interrupt hook JSON output", + ); + } +} + +#[test] +fn other_failures_use_standard_errors() { + for (exit_code, stdout, expected) in [ + (Some(0), "aloha", "Interrupt hook returned non-JSON stdout"), + (Some(2), "", "hook exited with code 2"), + ] { + assert_failed(stdout, exit_code, expected); + } +} + +fn assert_failed(stdout: &str, exit_code: Option, expected: &str) { + let parsed = parse(stdout, exit_code); + assert_eq!(parsed.data, InterruptHandlerData); + assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); + assert_eq!(parsed.completed.run.entries, vec![error(expected)]); +} + +fn parse( + stdout: &str, + exit_code: Option, +) -> crate::engine::dispatcher::ParsedHandler { + parse_completed( + &handler(), + run_result(exit_code, stdout, "ignored"), + Some("turn-1".to_string()), + ) +} + +fn warning(text: &str) -> HookOutputEntry { + HookOutputEntry { + kind: HookOutputEntryKind::Warning, + text: text.to_string(), + } +} + +fn error(text: &str) -> HookOutputEntry { + HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: text.to_string(), + } +} + +fn handler() -> ConfiguredHandler { + ConfiguredHandler { + builtin: false, + event_name: HookEventName::Interrupt, + matcher: None, + timeout_sec: 600, + status_message: None, + additional_context_limit: Default::default(), + source_path: test_path_buf("/tmp/hooks.json").abs().into(), + source: codex_protocol::protocol::HookSource::User, + display_order: 0, + kind: ConfiguredHandlerKind::Command { + command: "echo hook".to_string(), + r#async: false, + env: std::collections::HashMap::new(), + }, + } +} + +fn run_result(exit_code: Option, stdout: &str, stderr: &str) -> HandlerRunResult { + HandlerRunResult { + started_at: 1, + completed_at: 2, + duration_ms: 1, + exit_code, + stdout: stdout.to_string(), + stderr: stderr.to_string(), + error: None, + } +} diff --git a/codex-rs/hooks/src/events/mod.rs b/codex-rs/hooks/src/events/mod.rs new file mode 100644 index 0000000000000000000000000000000000000000..3d5ea1baad17c0b3814a59aa19c8b23feae49ed8 --- /dev/null +++ b/codex-rs/hooks/src/events/mod.rs @@ -0,0 +1,10 @@ +pub(crate) mod common; +pub mod compact; +pub mod interrupt; +pub mod permission_request; +pub mod post_tool_use; +pub mod pre_tool_use; +pub mod session_end; +pub mod session_start; +pub mod stop; +pub mod user_prompt_submit; diff --git a/codex-rs/hooks/src/events/permission_request.rs b/codex-rs/hooks/src/events/permission_request.rs new file mode 100644 index 0000000000000000000000000000000000000000..58da6c8e4d11082ff5ba402c2f2f705b75a9ed09 --- /dev/null +++ b/codex-rs/hooks/src/events/permission_request.rs @@ -0,0 +1,337 @@ +//! Permission-request hook execution. +//! +//! This event runs in the approval path, before guardian or user approval UI is +//! shown. Unlike `pre_tool_use`, handlers do not rewrite tool input or block by +//! stopping execution outright; instead they can return a concrete allow/deny +//! decision, or decline to decide and let the normal approval flow continue. +//! +//! The event also mirrors the rest of the hook system's lifecycle: +//! +//! 1. Preview matching handlers so the UI can render pending hook rows. +//! 2. Execute every matching handler in precedence order. +//! 3. Parse each handler into transcript-visible output plus an optional +//! decision. +//! 4. Fold the decisions conservatively: any deny wins, otherwise the last +//! allow wins, otherwise there is no hook verdict. +use std::path::PathBuf; + +use super::common; +use crate::engine::ClaudeHooksEngine; +use crate::engine::ConfiguredHandler; +use crate::engine::HandlerRunResult; +use crate::engine::dispatcher; +use crate::engine::output_parser; +use crate::schema::PermissionRequestCommandInput; +use crate::schema::SubagentCommandInputFields; +use codex_protocol::ThreadId; +use codex_protocol::protocol::HookCompletedEvent; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookOutputEntry; +use codex_protocol::protocol::HookOutputEntryKind; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookRunSummary; +use serde_json::Value; + +#[derive(Debug, Clone)] +pub struct PermissionRequestRequest { + pub session_id: ThreadId, + pub turn_id: String, + pub subagent: Option, + pub cwd: PathBuf, + pub transcript_path: Option, + pub model: String, + pub permission_mode: String, + pub tool_name: String, + pub matcher_aliases: Vec, + pub run_id_suffix: String, + pub tool_input: Value, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum PermissionRequestDecision { + Allow, + Deny { message: String }, +} + +#[derive(Debug)] +pub struct PermissionRequestOutcome { + pub hook_events: Vec, + pub decision: Option, +} + +#[derive(Debug, Default, PartialEq, Eq)] +struct PermissionRequestHandlerData { + decision: Option, +} + +pub(crate) fn preview( + handlers: &[ConfiguredHandler], + request: &PermissionRequestRequest, +) -> Vec { + let matcher_inputs = common::matcher_inputs(&request.tool_name, &request.matcher_aliases); + dispatcher::select_handlers_for_matcher_inputs( + handlers, + HookEventName::PermissionRequest, + &matcher_inputs, + ) + .into_iter() + .map(|handler| { + common::hook_run_for_tool_use( + dispatcher::running_summary(&handler), + &request.run_id_suffix, + ) + }) + .collect() +} + +pub(crate) async fn run( + engine: &ClaudeHooksEngine, + request: PermissionRequestRequest, +) -> PermissionRequestOutcome { + let matcher_inputs = common::matcher_inputs(&request.tool_name, &request.matcher_aliases); + let matched = dispatcher::select_handlers_for_matcher_inputs( + &engine.handlers, + HookEventName::PermissionRequest, + &matcher_inputs, + ); + if matched.is_empty() { + return PermissionRequestOutcome { + hook_events: Vec::new(), + decision: None, + }; + } + + let input_json = match serde_json::to_string(&build_command_input(&request)) { + Ok(input_json) => input_json, + Err(error) => { + let hook_events = common::serialization_failure_hook_events_for_tool_use( + matched, + Some(request.turn_id.clone()), + format!("failed to serialize permission request hook input: {error}"), + &request.run_id_suffix, + ); + return PermissionRequestOutcome { + hook_events, + decision: None, + }; + } + }; + + let results = dispatcher::execute_handlers( + engine, + matched, + input_json, + request.cwd.as_path(), + Some(request.turn_id.clone()), + parse_completed, + ) + .await; + + // Preserve the most specific matching allow, but treat any deny as final so + // broader policy layers cannot accidentally overrule a more specific block. + let decision = resolve_permission_request_decision( + results + .iter() + .filter_map(|result| result.data.decision.as_ref()), + ); + + PermissionRequestOutcome { + hook_events: results + .into_iter() + .map(|result| { + common::hook_completed_for_tool_use(result.completed, &request.run_id_suffix) + }) + .collect(), + decision, + } +} + +/// Resolve matching hook decisions conservatively: any deny wins immediately; +/// otherwise keep the highest-precedence allow so more specific handlers +/// override broader ones. +fn resolve_permission_request_decision<'a>( + decisions: impl IntoIterator, +) -> Option { + let mut resolved_allow = None; + for decision in decisions { + match decision { + PermissionRequestDecision::Allow => { + resolved_allow = Some(PermissionRequestDecision::Allow); + } + PermissionRequestDecision::Deny { message } => { + return Some(PermissionRequestDecision::Deny { + message: message.clone(), + }); + } + } + } + resolved_allow +} + +fn build_command_input(request: &PermissionRequestRequest) -> PermissionRequestCommandInput { + let subagent = SubagentCommandInputFields::from(request.subagent.as_ref()); + PermissionRequestCommandInput { + session_id: request.session_id.to_string(), + turn_id: request.turn_id.clone(), + agent_id: subagent.agent_id, + agent_type: subagent.agent_type, + transcript_path: crate::schema::NullableString::from_path(request.transcript_path.clone()), + cwd: request.cwd.display().to_string(), + hook_event_name: "PermissionRequest".to_string(), + model: request.model.clone(), + permission_mode: request.permission_mode.clone(), + tool_name: request.tool_name.clone(), + tool_input: request.tool_input.clone(), + } +} + +fn parse_completed( + handler: &ConfiguredHandler, + run_result: HandlerRunResult, + turn_id: Option, +) -> dispatcher::ParsedHandler { + let mut entries = Vec::new(); + let mut status = HookRunStatus::Completed; + let mut decision = None; + + match run_result.error.as_deref() { + Some(error) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: error.to_string(), + }); + } + None => match run_result.exit_code { + Some(0) => { + let trimmed_stdout = run_result.stdout.trim(); + if trimmed_stdout.is_empty() { + } else if let Some(parsed) = + output_parser::parse_permission_request(&run_result.stdout) + { + if let Some(system_message) = parsed.universal.system_message { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Warning, + text: system_message, + }); + } + if handler.can_apply_control_effects() { + if let Some(invalid_reason) = parsed.invalid_reason { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: invalid_reason, + }); + } else if let Some(parsed_decision) = parsed.decision { + match parsed_decision { + output_parser::PermissionRequestDecision::Allow => { + decision = Some(PermissionRequestDecision::Allow); + } + output_parser::PermissionRequestDecision::Deny { message } => { + status = HookRunStatus::Blocked; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: message.clone(), + }); + decision = Some(PermissionRequestDecision::Deny { message }); + } + } + } + } + } else if output_parser::looks_like_json(&run_result.stdout) { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook returned invalid permission-request JSON output".to_string(), + }); + } + } + Some(2) if handler.can_apply_control_effects() => { + if let Some(message) = common::trimmed_non_empty(&run_result.stderr) { + status = HookRunStatus::Blocked; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: message.clone(), + }); + decision = Some(PermissionRequestDecision::Deny { message }); + } else { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "PermissionRequest hook exited with code 2 but did not write a denial reason to stderr".to_string(), + }); + } + } + Some(exit_code) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: format!("hook exited with code {exit_code}"), + }); + } + None => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook exited without a status code".to_string(), + }); + } + }, + } + + let completed = HookCompletedEvent { + turn_id, + run: dispatcher::completed_summary(handler, &run_result, status, entries), + }; + + dispatcher::ParsedHandler { + completed, + data: PermissionRequestHandlerData { decision }, + completion_order: 0, + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::PermissionRequestDecision; + use super::resolve_permission_request_decision; + + #[test] + fn permission_request_deny_overrides_earlier_allow() { + let decisions = [ + PermissionRequestDecision::Allow, + PermissionRequestDecision::Deny { + message: "repo deny".to_string(), + }, + ]; + + assert_eq!( + resolve_permission_request_decision(decisions.iter()), + Some(PermissionRequestDecision::Deny { + message: "repo deny".to_string(), + }) + ); + } + + #[test] + fn permission_request_returns_allow_when_no_handler_denies() { + let decisions = [ + PermissionRequestDecision::Allow, + PermissionRequestDecision::Allow, + ]; + + assert_eq!( + resolve_permission_request_decision(decisions.iter()), + Some(PermissionRequestDecision::Allow) + ); + } + + #[test] + fn permission_request_returns_none_when_no_handler_decides() { + let decisions = Vec::::new(); + + assert_eq!(resolve_permission_request_decision(decisions.iter()), None); + } +} diff --git a/codex-rs/hooks/src/events/post_tool_use.rs b/codex-rs/hooks/src/events/post_tool_use.rs new file mode 100644 index 0000000000000000000000000000000000000000..aeb0ed20a7a12cb0ffd7ba979633c8a7557ccdae --- /dev/null +++ b/codex-rs/hooks/src/events/post_tool_use.rs @@ -0,0 +1,637 @@ +use std::path::PathBuf; + +use codex_protocol::ThreadId; +use codex_protocol::protocol::HookCompletedEvent; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookOutputEntry; +use codex_protocol::protocol::HookOutputEntryKind; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookRunSummary; +use codex_utils_absolute_path::AbsolutePathBuf; +use serde_json::Value; + +use super::common; +use crate::engine::ClaudeHooksEngine; +use crate::engine::ConfiguredHandler; +use crate::engine::HandlerRunResult; +use crate::engine::dispatcher; +use crate::engine::output_parser; +use crate::output_spill::AdditionalContext; +use crate::schema::PostToolUseCommandInput; +use crate::schema::SubagentCommandInputFields; + +#[derive(Debug, Clone)] +pub struct PostToolUseRequest { + pub session_id: ThreadId, + pub turn_id: String, + pub subagent: Option, + pub cwd: AbsolutePathBuf, + pub transcript_path: Option, + pub model: String, + pub permission_mode: String, + pub tool_name: String, + pub matcher_aliases: Vec, + pub tool_use_id: String, + pub tool_input: Value, + pub tool_response: Value, +} + +#[derive(Debug)] +pub struct PostToolUseOutcome { + pub hook_events: Vec, + pub should_block: bool, + pub additional_contexts: Vec, + pub feedback_message: Option, +} + +#[derive(Debug, Default, PartialEq, Eq)] +struct PostToolUseHandlerData { + should_block: bool, + additional_contexts_for_model: Vec, + feedback_messages_for_model: Vec, +} + +pub(crate) fn preview( + handlers: &[ConfiguredHandler], + request: &PostToolUseRequest, +) -> Vec { + let matcher_inputs = common::matcher_inputs(&request.tool_name, &request.matcher_aliases); + dispatcher::select_handlers_for_matcher_inputs( + handlers, + HookEventName::PostToolUse, + &matcher_inputs, + ) + .into_iter() + .map(|handler| { + common::hook_run_for_tool_use(dispatcher::running_summary(&handler), &request.tool_use_id) + }) + .collect() +} + +pub(crate) async fn run( + engine: &ClaudeHooksEngine, + request: PostToolUseRequest, +) -> PostToolUseOutcome { + let matcher_inputs = common::matcher_inputs(&request.tool_name, &request.matcher_aliases); + let matched = dispatcher::select_handlers_for_matcher_inputs( + &engine.handlers, + HookEventName::PostToolUse, + &matcher_inputs, + ); + if matched.is_empty() { + return PostToolUseOutcome { + hook_events: Vec::new(), + should_block: false, + additional_contexts: Vec::new(), + feedback_message: None, + }; + } + + let input_json = match command_input_json(&request) { + Ok(input_json) => input_json, + Err(error) => { + let hook_events = common::serialization_failure_hook_events_for_tool_use( + matched, + Some(request.turn_id.clone()), + format!("failed to serialize post tool use hook input: {error}"), + &request.tool_use_id, + ); + return serialization_failure_outcome(hook_events); + } + }; + + let results = dispatcher::execute_handlers( + engine, + matched, + input_json, + request.cwd.as_path(), + Some(request.turn_id.clone()), + parse_completed, + ) + .await; + + let additional_contexts = common::flatten_additional_contexts( + results + .iter() + .map(|result| result.data.additional_contexts_for_model.as_slice()), + ); + let additional_contexts = engine + .command_runtime + .output_spiller() + .maybe_spill_additional_contexts(additional_contexts) + .await; + let should_block = results.iter().any(|result| result.data.should_block); + let feedback_message = common::join_text_chunks( + results + .iter() + .flat_map(|result| result.data.feedback_messages_for_model.clone()) + .collect(), + ); + + PostToolUseOutcome { + hook_events: results + .into_iter() + .map(|result| { + common::hook_completed_for_tool_use(result.completed, &request.tool_use_id) + }) + .collect(), + should_block, + additional_contexts, + feedback_message, + } +} + +/// Serializes command stdin for a selected `PostToolUse` hook. +/// +/// Handler selection may include internal matcher aliases, but hook stdin keeps +/// the canonical `tool_name` for logs and for consumers that pair pre/post +/// events across processes. Shell-like tools pass `{ "command": ... }` as +/// `tool_input`; MCP tools pass their resolved JSON arguments. +fn command_input_json(request: &PostToolUseRequest) -> Result { + let subagent = SubagentCommandInputFields::from(request.subagent.as_ref()); + serde_json::to_string(&PostToolUseCommandInput { + session_id: request.session_id.to_string(), + turn_id: request.turn_id.clone(), + agent_id: subagent.agent_id, + agent_type: subagent.agent_type, + transcript_path: crate::schema::NullableString::from_path(request.transcript_path.clone()), + cwd: request.cwd.display().to_string(), + hook_event_name: "PostToolUse".to_string(), + model: request.model.clone(), + permission_mode: request.permission_mode.clone(), + tool_name: request.tool_name.clone(), + tool_input: request.tool_input.clone(), + tool_response: request.tool_response.clone(), + tool_use_id: request.tool_use_id.clone(), + }) +} + +fn parse_completed( + handler: &ConfiguredHandler, + run_result: HandlerRunResult, + turn_id: Option, +) -> dispatcher::ParsedHandler { + let mut entries = Vec::new(); + let mut status = HookRunStatus::Completed; + let mut should_block = false; + let mut additional_contexts_for_model = Vec::new(); + let mut feedback_messages_for_model = Vec::new(); + + match run_result.error.as_deref() { + Some(error) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: error.to_string(), + }); + } + None => match run_result.exit_code { + Some(0) => { + let trimmed_stdout = run_result.stdout.trim(); + if trimmed_stdout.is_empty() { + } else if let Some(parsed) = output_parser::parse_post_tool_use(&run_result.stdout) + { + if let Some(system_message) = parsed.universal.system_message { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Warning, + text: system_message, + }); + } + if (!handler.can_apply_control_effects() + || parsed.invalid_reason.is_none() && parsed.invalid_block_reason.is_none()) + && let Some(additional_context) = parsed.additional_context + { + common::append_additional_context( + &mut entries, + &mut additional_contexts_for_model, + handler, + additional_context, + ); + } + if handler.can_apply_control_effects() { + if !parsed.universal.continue_processing { + status = HookRunStatus::Stopped; + let stop_text = parsed.universal.stop_reason.unwrap_or_else(|| { + "PostToolUse hook stopped execution".to_string() + }); + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Stop, + text: stop_text.clone(), + }); + let model_feedback = parsed + .reason + .as_deref() + .and_then(common::trimmed_non_empty) + .unwrap_or(stop_text); + feedback_messages_for_model.push(model_feedback); + } else if let Some(invalid_reason) = parsed.invalid_reason { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: invalid_reason, + }); + } else if let Some(invalid_block_reason) = parsed.invalid_block_reason { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: invalid_block_reason, + }); + } else if parsed.should_block { + status = HookRunStatus::Blocked; + should_block = true; + if let Some(reason) = parsed.reason { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: reason.clone(), + }); + feedback_messages_for_model.push(reason); + } + } + } + } else if output_parser::looks_like_json(&run_result.stdout) { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook returned invalid post-tool-use JSON output".to_string(), + }); + } + } + Some(2) if handler.can_apply_control_effects() => { + if let Some(reason) = common::trimmed_non_empty(&run_result.stderr) { + status = HookRunStatus::Blocked; + should_block = true; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: reason.clone(), + }); + feedback_messages_for_model.push(reason); + } else { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "PostToolUse hook exited with code 2 but did not write feedback to stderr".to_string(), + }); + } + } + Some(exit_code) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: format!("hook exited with code {exit_code}"), + }); + } + None => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook exited without a status code".to_string(), + }); + } + }, + } + + let completed = HookCompletedEvent { + turn_id, + run: dispatcher::completed_summary(handler, &run_result, status, entries), + }; + + dispatcher::ParsedHandler { + completed, + data: PostToolUseHandlerData { + should_block, + additional_contexts_for_model, + feedback_messages_for_model, + }, + completion_order: 0, + } +} + +fn serialization_failure_outcome(hook_events: Vec) -> PostToolUseOutcome { + PostToolUseOutcome { + hook_events, + should_block: false, + additional_contexts: Vec::new(), + feedback_message: None, + } +} + +#[cfg(test)] +mod tests { + use codex_protocol::ThreadId; + use codex_protocol::protocol::HookEventName; + use codex_protocol::protocol::HookOutputEntry; + use codex_protocol::protocol::HookOutputEntryKind; + use codex_protocol::protocol::HookRunStatus; + use codex_utils_absolute_path::test_support::PathBufExt; + use codex_utils_absolute_path::test_support::test_path_buf; + use pretty_assertions::assert_eq; + use serde_json::json; + + use super::PostToolUseHandlerData; + use super::command_input_json; + use super::parse_completed; + use super::preview; + use crate::engine::ConfiguredHandler; + use crate::engine::HandlerRunResult; + use crate::events::common; + use crate::output_spill::AdditionalContext; + use crate::output_spill::AdditionalContextLimit; + + #[test] + fn command_input_uses_request_tool_name() { + let mut request = request_for_tool_use("call-apply-patch"); + request.tool_name = "apply_patch".to_string(); + + let input_json = command_input_json(&request).expect("serialize command input"); + let input: serde_json::Value = + serde_json::from_str(&input_json).expect("parse command input"); + + assert_eq!(input["tool_name"], "apply_patch"); + } + + #[test] + fn block_decision_stops_normal_processing() { + let parsed = parse_completed( + &handler(), + run_result( + Some(0), + r#"{"decision":"block","reason":"bash output looked sketchy"}"#, + "", + ), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PostToolUseHandlerData { + should_block: true, + additional_contexts_for_model: Vec::new(), + feedback_messages_for_model: vec!["bash output looked sketchy".to_string()], + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked); + } + + #[test] + fn additional_context_is_recorded() { + let mut handler = handler(); + handler.additional_context_limit = AdditionalContextLimit::from_config(Some(17)); + let parsed = parse_completed( + &handler, + run_result( + Some(0), + r#"{"hookSpecificOutput":{"hookEventName":"PostToolUse","additionalContext":"Remember the bash cleanup note."}}"#, + "", + ), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PostToolUseHandlerData { + should_block: false, + additional_contexts_for_model: vec![AdditionalContext { + text: "Remember the bash cleanup note.".to_string(), + limit: AdditionalContextLimit::from_config(Some(17)), + }], + feedback_messages_for_model: Vec::new(), + } + ); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Context, + text: "Remember the bash cleanup note.".to_string(), + }] + ); + } + + #[test] + fn unsupported_updated_mcp_tool_output_fails_open() { + let stdout = r#"{"hookSpecificOutput":{"hookEventName":"PostToolUse","updatedMCPToolOutput":{"ok":true},"additionalContext":"preserved"}}"#; + let parsed = parse_completed( + &handler(), + run_result(Some(0), stdout, ""), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PostToolUseHandlerData { + should_block: false, + additional_contexts_for_model: Vec::new(), + feedback_messages_for_model: Vec::new(), + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "PostToolUse hook returned unsupported updatedMCPToolOutput".to_string(), + }] + ); + + let async_handler = handler_with_async(/*async*/ true); + let parsed = parse_completed( + &async_handler, + run_result(Some(0), stdout, ""), + Some("turn-1".to_string()), + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Context, + text: "preserved".to_string(), + }], + ); + } + + #[test] + fn exit_two_blocks_with_feedback() { + let parsed = parse_completed( + &handler(), + run_result(Some(2), "", "post hook says pause"), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PostToolUseHandlerData { + should_block: true, + additional_contexts_for_model: Vec::new(), + feedback_messages_for_model: vec!["post hook says pause".to_string()], + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked); + } + + #[test] + fn continue_false_stops_with_reason() { + let parsed = parse_completed( + &handler(), + run_result( + Some(0), + r#"{"continue":false,"stopReason":"halt after bash output","reason":"post-tool hook says stop"}"#, + "", + ), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PostToolUseHandlerData { + should_block: false, + additional_contexts_for_model: Vec::new(), + feedback_messages_for_model: vec!["post-tool hook says stop".to_string()], + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Stopped); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Stop, + text: "halt after bash output".to_string(), + }] + ); + } + + #[test] + fn continue_false_without_reason_synthesizes_feedback() { + let parsed = parse_completed( + &handler(), + run_result(Some(0), r#"{"continue":false}"#, ""), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PostToolUseHandlerData { + should_block: false, + additional_contexts_for_model: Vec::new(), + feedback_messages_for_model: vec!["PostToolUse hook stopped execution".to_string()], + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Stopped); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Stop, + text: "PostToolUse hook stopped execution".to_string(), + }] + ); + } + + #[test] + fn plain_stdout_is_ignored_for_post_tool_use() { + let parsed = parse_completed( + &handler(), + run_result(Some(0), "plain text only", ""), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PostToolUseHandlerData { + should_block: false, + additional_contexts_for_model: Vec::new(), + feedback_messages_for_model: Vec::new(), + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); + assert_eq!(parsed.completed.run.entries, Vec::::new()); + } + + #[test] + fn preview_and_completed_run_ids_include_tool_use_id() { + let request = request_for_tool_use("tool-call-456"); + let runs = preview(&[handler()], &request); + + assert_eq!(runs.len(), 1); + assert_eq!( + runs[0].id, + format!( + "post-tool-use:0:{}:tool-call-456", + test_path_buf("/tmp/hooks.json").display() + ) + ); + + let parsed = parse_completed( + &handler(), + run_result(Some(0), "", ""), + Some("turn-1".to_string()), + ); + let completed = common::hook_completed_for_tool_use(parsed.completed, &request.tool_use_id); + + assert_eq!(completed.run.id, runs[0].id); + } + + #[test] + fn serialization_failure_run_ids_include_tool_use_id() { + let request = request_for_tool_use("tool-call-456"); + let runs = preview(&[handler()], &request); + + let completed = common::serialization_failure_hook_events_for_tool_use( + vec![handler()], + Some(request.turn_id.clone()), + "serialize failed".into(), + &request.tool_use_id, + ); + + assert_eq!(completed.len(), 1); + assert_eq!(completed[0].run.id, runs[0].id); + } + + fn handler() -> ConfiguredHandler { + handler_with_async(/*async*/ false) + } + + fn handler_with_async(r#async: bool) -> ConfiguredHandler { + ConfiguredHandler { + builtin: false, + event_name: HookEventName::PostToolUse, + matcher: Some("^Bash$".to_string()), + timeout_sec: 5, + status_message: Some("running post tool use hook".to_string()), + additional_context_limit: Default::default(), + source_path: test_path_buf("/tmp/hooks.json").abs().into(), + source: codex_protocol::protocol::HookSource::User, + display_order: 0, + kind: crate::engine::ConfiguredHandlerKind::Command { + command: "python3 post_tool_use_hook.py".to_string(), + r#async, + env: std::collections::HashMap::new(), + }, + } + } + + fn run_result(exit_code: Option, stdout: &str, stderr: &str) -> HandlerRunResult { + HandlerRunResult { + started_at: 1_700_000_000, + completed_at: 1_700_000_001, + duration_ms: 12, + exit_code, + stdout: stdout.to_string(), + stderr: stderr.to_string(), + error: None, + } + } + + fn request_for_tool_use(tool_use_id: &str) -> super::PostToolUseRequest { + super::PostToolUseRequest { + session_id: ThreadId::new(), + turn_id: "turn-1".to_string(), + subagent: None, + cwd: test_path_buf("/tmp").abs(), + transcript_path: None, + model: "gpt-test".to_string(), + permission_mode: "default".to_string(), + tool_name: "Bash".to_string(), + matcher_aliases: Vec::new(), + tool_use_id: tool_use_id.to_string(), + tool_input: json!({ "command": "echo hello" }), + tool_response: json!({"ok": true}), + } + } +} diff --git a/codex-rs/hooks/src/events/pre_tool_use.rs b/codex-rs/hooks/src/events/pre_tool_use.rs new file mode 100644 index 0000000000000000000000000000000000000000..c1baebd4f01fd1906f461fe84cb2bf0d73058bfc --- /dev/null +++ b/codex-rs/hooks/src/events/pre_tool_use.rs @@ -0,0 +1,820 @@ +use std::path::PathBuf; + +use codex_protocol::ThreadId; +use codex_protocol::protocol::HookCompletedEvent; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookOutputEntry; +use codex_protocol::protocol::HookOutputEntryKind; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookRunSummary; +use codex_utils_absolute_path::AbsolutePathBuf; +use serde_json::Value; + +use super::common; +use crate::engine::ClaudeHooksEngine; +use crate::engine::ConfiguredHandler; +use crate::engine::HandlerRunResult; +use crate::engine::dispatcher; +use crate::engine::output_parser; +use crate::output_spill::AdditionalContext; +use crate::schema::PreToolUseCommandInput; +use crate::schema::SubagentCommandInputFields; + +#[derive(Debug, Clone)] +pub struct PreToolUseRequest { + pub session_id: ThreadId, + pub turn_id: String, + pub subagent: Option, + pub cwd: AbsolutePathBuf, + pub transcript_path: Option, + pub model: String, + pub permission_mode: String, + pub tool_name: String, + pub matcher_aliases: Vec, + pub tool_use_id: String, + pub tool_input: Value, +} + +#[derive(Debug)] +pub struct PreToolUseOutcome { + pub hook_events: Vec, + pub should_block: bool, + pub block_reason: Option, + pub additional_contexts: Vec, + pub updated_input: Option, +} + +#[derive(Debug, Default, PartialEq, Eq)] +struct PreToolUseHandlerData { + should_block: bool, + block_reason: Option, + additional_contexts_for_model: Vec, + updated_input: Option, +} + +pub(crate) fn preview( + handlers: &[ConfiguredHandler], + request: &PreToolUseRequest, +) -> Vec { + let matcher_inputs = common::matcher_inputs(&request.tool_name, &request.matcher_aliases); + dispatcher::select_handlers_for_matcher_inputs( + handlers, + HookEventName::PreToolUse, + &matcher_inputs, + ) + .into_iter() + .map(|handler| { + common::hook_run_for_tool_use(dispatcher::running_summary(&handler), &request.tool_use_id) + }) + .collect() +} + +pub(crate) async fn run( + engine: &ClaudeHooksEngine, + request: PreToolUseRequest, +) -> PreToolUseOutcome { + let matcher_inputs = common::matcher_inputs(&request.tool_name, &request.matcher_aliases); + let matched = dispatcher::select_handlers_for_matcher_inputs( + &engine.handlers, + HookEventName::PreToolUse, + &matcher_inputs, + ); + if matched.is_empty() { + return PreToolUseOutcome { + hook_events: Vec::new(), + should_block: false, + block_reason: None, + additional_contexts: Vec::new(), + updated_input: None, + }; + } + + let input_json = match command_input_json(&request) { + Ok(input_json) => input_json, + Err(error) => { + let hook_events = common::serialization_failure_hook_events_for_tool_use( + matched, + Some(request.turn_id.clone()), + format!("failed to serialize pre tool use hook input: {error}"), + &request.tool_use_id, + ); + return serialization_failure_outcome(hook_events); + } + }; + + let results = dispatcher::execute_handlers( + engine, + matched, + input_json, + request.cwd.as_path(), + Some(request.turn_id.clone()), + parse_completed, + ) + .await; + + let should_block = results.iter().any(|result| result.data.should_block); + let block_reason = results + .iter() + .find_map(|result| result.data.block_reason.clone()); + let additional_contexts = common::flatten_additional_contexts( + results + .iter() + .map(|result| result.data.additional_contexts_for_model.as_slice()), + ); + let additional_contexts = engine + .command_runtime + .output_spiller() + .maybe_spill_additional_contexts(additional_contexts) + .await; + let updated_input = if should_block { + None + } else { + latest_updated_input(&results) + }; + + PreToolUseOutcome { + hook_events: results + .into_iter() + .map(|result| { + common::hook_completed_for_tool_use(result.completed, &request.tool_use_id) + }) + .collect(), + should_block, + block_reason, + additional_contexts, + updated_input, + } +} + +/// Chooses the rewrite from the hook that actually finished last. +/// +/// Hook results stay in configured order for stable reporting, but the +/// `PreToolUse` contract resolves competing rewrites by completion order. +fn latest_updated_input( + results: &[dispatcher::ParsedHandler], +) -> Option { + results + .iter() + .filter_map(|result| { + result + .data + .updated_input + .clone() + .map(|updated_input| (result.completion_order, updated_input)) + }) + .max_by_key(|(completion_order, _)| *completion_order) + .map(|(_, updated_input)| updated_input) +} + +/// Serializes command stdin for a selected `PreToolUse` hook. +/// +/// Handler selection may include internal matcher aliases, but hook stdin keeps +/// the canonical `tool_name` so audit logs and downstream policy decisions stay +/// stable. Shell-like tools pass `{ "command": ... }` as `tool_input`; MCP +/// tools pass their resolved JSON arguments. +fn command_input_json(request: &PreToolUseRequest) -> Result { + let subagent = SubagentCommandInputFields::from(request.subagent.as_ref()); + serde_json::to_string(&PreToolUseCommandInput { + session_id: request.session_id.to_string(), + turn_id: request.turn_id.clone(), + agent_id: subagent.agent_id, + agent_type: subagent.agent_type, + transcript_path: crate::schema::NullableString::from_path(request.transcript_path.clone()), + cwd: request.cwd.display().to_string(), + hook_event_name: "PreToolUse".to_string(), + model: request.model.clone(), + permission_mode: request.permission_mode.clone(), + tool_name: request.tool_name.clone(), + tool_input: request.tool_input.clone(), + tool_use_id: request.tool_use_id.clone(), + }) +} + +fn parse_completed( + handler: &ConfiguredHandler, + run_result: HandlerRunResult, + turn_id: Option, +) -> dispatcher::ParsedHandler { + let mut entries = Vec::new(); + let mut status = HookRunStatus::Completed; + let mut should_block = false; + let mut block_reason = None; + let mut additional_contexts_for_model = Vec::new(); + let mut updated_input = None; + + match run_result.error.as_deref() { + Some(error) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: error.to_string(), + }); + } + None => match run_result.exit_code { + Some(0) => { + let trimmed_stdout = run_result.stdout.trim(); + if trimmed_stdout.is_empty() { + } else if let Some(parsed) = output_parser::parse_pre_tool_use(&run_result.stdout) { + if let Some(system_message) = parsed.universal.system_message { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Warning, + text: system_message, + }); + } + if (!handler.can_apply_control_effects() || parsed.invalid_reason.is_none()) + && let Some(additional_context) = parsed.additional_context + { + common::append_additional_context( + &mut entries, + &mut additional_contexts_for_model, + handler, + additional_context, + ); + } + if handler.can_apply_control_effects() { + if let Some(invalid_reason) = parsed.invalid_reason { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: invalid_reason, + }); + } else if let Some(reason) = parsed.block_reason { + status = HookRunStatus::Blocked; + should_block = true; + block_reason = Some(reason.clone()); + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: reason, + }); + } else { + updated_input = parsed.updated_input; + } + } + } else if output_parser::looks_like_json(&run_result.stdout) { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook returned invalid pre-tool-use JSON output".to_string(), + }); + } + } + Some(2) if handler.can_apply_control_effects() => { + if let Some(reason) = common::trimmed_non_empty(&run_result.stderr) { + status = HookRunStatus::Blocked; + should_block = true; + block_reason = Some(reason.clone()); + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: reason, + }); + } else { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "PreToolUse hook exited with code 2 but did not write a blocking reason to stderr".to_string(), + }); + } + } + Some(exit_code) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: format!("hook exited with code {exit_code}"), + }); + } + None => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook exited without a status code".to_string(), + }); + } + }, + } + + let completed = HookCompletedEvent { + turn_id, + run: dispatcher::completed_summary(handler, &run_result, status, entries), + }; + + dispatcher::ParsedHandler { + completed, + data: PreToolUseHandlerData { + should_block, + block_reason, + additional_contexts_for_model, + updated_input, + }, + completion_order: 0, + } +} + +fn serialization_failure_outcome(hook_events: Vec) -> PreToolUseOutcome { + PreToolUseOutcome { + hook_events, + should_block: false, + block_reason: None, + additional_contexts: Vec::new(), + updated_input: None, + } +} + +#[cfg(test)] +mod tests { + use codex_protocol::ThreadId; + use codex_protocol::protocol::HookEventName; + use codex_protocol::protocol::HookOutputEntry; + use codex_protocol::protocol::HookOutputEntryKind; + use codex_protocol::protocol::HookRunStatus; + use codex_utils_absolute_path::test_support::PathBufExt; + use codex_utils_absolute_path::test_support::test_path_buf; + use pretty_assertions::assert_eq; + + use super::PreToolUseHandlerData; + use super::command_input_json; + use super::latest_updated_input; + use super::parse_completed; + use super::preview; + use crate::engine::ConfiguredHandler; + use crate::engine::HandlerRunResult; + use crate::events::common; + use crate::output_spill::AdditionalContext; + use crate::output_spill::AdditionalContextLimit; + + #[test] + fn command_input_uses_request_tool_name() { + let mut request = request_for_tool_use("call-apply-patch"); + request.tool_name = "apply_patch".to_string(); + + let input_json = command_input_json(&request).expect("serialize command input"); + let input: serde_json::Value = + serde_json::from_str(&input_json).expect("parse command input"); + + assert_eq!(input["tool_name"], "apply_patch"); + } + + #[test] + fn permission_decision_deny_blocks_processing() { + let parsed = parse_completed( + &handler(), + run_result( + Some(0), + r#"{"hookSpecificOutput":{"hookEventName":"PreToolUse","permissionDecision":"deny","permissionDecisionReason":"do not run that"}}"#, + "", + ), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PreToolUseHandlerData { + should_block: true, + block_reason: Some("do not run that".to_string()), + additional_contexts_for_model: Vec::new(), + updated_input: None, + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: "do not run that".to_string(), + }] + ); + } + + #[test] + fn permission_decision_allow_can_update_input() { + let parsed = parse_completed( + &handler(), + run_result( + Some(0), + r#"{"hookSpecificOutput":{"hookEventName":"PreToolUse","permissionDecision":"allow","updatedInput":{"command":"echo rewritten"}}}"#, + "", + ), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PreToolUseHandlerData { + should_block: false, + block_reason: None, + additional_contexts_for_model: Vec::new(), + updated_input: Some(serde_json::json!({ "command": "echo rewritten" })), + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); + assert_eq!(parsed.completed.run.entries, vec![]); + } + + #[test] + fn last_completed_updated_input_wins() { + let mut later_configured = parse_completed( + &handler(), + run_result( + Some(0), + r#"{"hookSpecificOutput":{"hookEventName":"PreToolUse","permissionDecision":"allow","updatedInput":{"command":"echo configured later"}}}"#, + "", + ), + Some("turn-1".to_string()), + ); + later_configured.completion_order = 0; + let mut earlier_configured = parse_completed( + &handler(), + run_result( + Some(0), + r#"{"hookSpecificOutput":{"hookEventName":"PreToolUse","permissionDecision":"allow","updatedInput":{"command":"echo finished later"}}}"#, + "", + ), + Some("turn-1".to_string()), + ); + earlier_configured.completion_order = 1; + + assert_eq!( + latest_updated_input(&[later_configured, earlier_configured]), + Some(serde_json::json!({ "command": "echo finished later" })) + ); + } + + #[test] + fn permission_decision_allow_without_updated_input_fails_open() { + let stdout = r#"{"hookSpecificOutput":{"hookEventName":"PreToolUse","permissionDecision":"allow","additionalContext":"preserved"}}"#; + let parsed = parse_completed( + &handler(), + run_result(Some(0), stdout, ""), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PreToolUseHandlerData { + should_block: false, + block_reason: None, + additional_contexts_for_model: Vec::new(), + updated_input: None, + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "PreToolUse hook returned unsupported permissionDecision:allow".to_string(), + }] + ); + + let async_handler = handler_with_async(/*async*/ true); + let parsed = parse_completed( + &async_handler, + run_result(Some(0), stdout, ""), + Some("turn-1".to_string()), + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Context, + text: "preserved".to_string(), + }], + ); + } + + #[test] + fn deprecated_block_decision_blocks_processing() { + let parsed = parse_completed( + &handler(), + run_result( + Some(0), + r#"{"decision":"block","reason":"do not run that"}"#, + "", + ), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PreToolUseHandlerData { + should_block: true, + block_reason: Some("do not run that".to_string()), + additional_contexts_for_model: Vec::new(), + updated_input: None, + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: "do not run that".to_string(), + }] + ); + } + + #[test] + fn deprecated_block_decision_with_additional_context_blocks_processing() { + let parsed = parse_completed( + &handler(), + run_result( + Some(0), + r#"{"decision":"block","reason":"do not run that","hookSpecificOutput":{"hookEventName":"PreToolUse","additionalContext":"remember this"}}"#, + "", + ), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PreToolUseHandlerData { + should_block: true, + block_reason: Some("do not run that".to_string()), + additional_contexts_for_model: vec![AdditionalContext { + text: "remember this".to_string(), + limit: Default::default(), + }], + updated_input: None, + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked); + assert_eq!( + parsed.completed.run.entries, + vec![ + HookOutputEntry { + kind: HookOutputEntryKind::Context, + text: "remember this".to_string(), + }, + HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: "do not run that".to_string(), + }, + ] + ); + } + + #[test] + fn unsupported_permission_decision_fails_open() { + let parsed = parse_completed( + &handler(), + run_result( + Some(0), + r#"{"hookSpecificOutput":{"hookEventName":"PreToolUse","permissionDecision":"ask","permissionDecisionReason":"please confirm"}}"#, + "", + ), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PreToolUseHandlerData { + should_block: false, + block_reason: None, + additional_contexts_for_model: Vec::new(), + updated_input: None, + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "PreToolUse hook returned unsupported permissionDecision:ask".to_string(), + }] + ); + } + + #[test] + fn deprecated_approve_decision_fails_open() { + let parsed = parse_completed( + &handler(), + run_result(Some(0), r#"{"decision":"approve"}"#, ""), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PreToolUseHandlerData { + should_block: false, + block_reason: None, + additional_contexts_for_model: Vec::new(), + updated_input: None, + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "PreToolUse hook returned unsupported decision:approve".to_string(), + }] + ); + } + + #[test] + fn additional_context_is_recorded() { + let mut handler = handler(); + handler.additional_context_limit = AdditionalContextLimit::from_config(Some(13)); + let parsed = parse_completed( + &handler, + run_result( + Some(0), + r#"{"hookSpecificOutput":{"hookEventName":"PreToolUse","permissionDecision":"deny","permissionDecisionReason":"do not run that","additionalContext":"nope"}}"#, + "", + ), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PreToolUseHandlerData { + should_block: true, + block_reason: Some("do not run that".to_string()), + additional_contexts_for_model: vec![AdditionalContext { + text: "nope".to_string(), + limit: AdditionalContextLimit::from_config(Some(13)), + }], + updated_input: None, + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked); + assert_eq!( + parsed.completed.run.entries, + vec![ + HookOutputEntry { + kind: HookOutputEntryKind::Context, + text: "nope".to_string(), + }, + HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: "do not run that".to_string(), + }, + ] + ); + } + + #[test] + fn plain_stdout_is_ignored() { + let parsed = parse_completed( + &handler(), + run_result(Some(0), "hook ran successfully\n", ""), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PreToolUseHandlerData { + should_block: false, + block_reason: None, + additional_contexts_for_model: Vec::new(), + updated_input: None, + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); + assert_eq!(parsed.completed.run.entries, vec![]); + } + + #[test] + fn invalid_json_like_stdout_fails_instead_of_becoming_noop() { + let parsed = parse_completed( + &handler(), + run_result(Some(0), "{\"decision\":\n", ""), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PreToolUseHandlerData { + should_block: false, + block_reason: None, + additional_contexts_for_model: Vec::new(), + updated_input: None, + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook returned invalid pre-tool-use JSON output".to_string(), + }] + ); + } + + #[test] + fn exit_code_two_blocks_processing() { + let parsed = parse_completed( + &handler(), + run_result(Some(2), "", "blocked by policy\n"), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + PreToolUseHandlerData { + should_block: true, + block_reason: Some("blocked by policy".to_string()), + additional_contexts_for_model: Vec::new(), + updated_input: None, + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: "blocked by policy".to_string(), + }] + ); + } + + #[test] + fn preview_and_completed_run_ids_include_tool_use_id() { + let request = request_for_tool_use("tool-call-123"); + let runs = preview(&[handler()], &request); + + assert_eq!(runs.len(), 1); + assert_eq!( + runs[0].id, + format!( + "pre-tool-use:0:{}:tool-call-123", + test_path_buf("/tmp/hooks.json").display() + ) + ); + + let parsed = parse_completed( + &handler(), + run_result(Some(0), "", ""), + Some("turn-1".to_string()), + ); + let completed = common::hook_completed_for_tool_use(parsed.completed, &request.tool_use_id); + + assert_eq!(completed.run.id, runs[0].id); + } + + #[test] + fn serialization_failure_run_ids_include_tool_use_id() { + let request = request_for_tool_use("tool-call-123"); + let runs = preview(&[handler()], &request); + + let completed = common::serialization_failure_hook_events_for_tool_use( + vec![handler()], + Some(request.turn_id.clone()), + "serialize failed".into(), + &request.tool_use_id, + ); + + assert_eq!(completed.len(), 1); + assert_eq!(completed[0].run.id, runs[0].id); + } + + fn handler() -> ConfiguredHandler { + handler_with_async(/*async*/ false) + } + + fn handler_with_async(r#async: bool) -> ConfiguredHandler { + ConfiguredHandler { + builtin: false, + event_name: HookEventName::PreToolUse, + matcher: Some("^Bash$".to_string()), + timeout_sec: 5, + status_message: None, + additional_context_limit: Default::default(), + source_path: test_path_buf("/tmp/hooks.json").abs().into(), + source: codex_protocol::protocol::HookSource::User, + display_order: 0, + kind: crate::engine::ConfiguredHandlerKind::Command { + command: "echo hook".to_string(), + r#async, + env: std::collections::HashMap::new(), + }, + } + } + + fn run_result(exit_code: Option, stdout: &str, stderr: &str) -> HandlerRunResult { + HandlerRunResult { + started_at: 1, + completed_at: 2, + duration_ms: 1, + exit_code, + stdout: stdout.to_string(), + stderr: stderr.to_string(), + error: None, + } + } + + fn request_for_tool_use(tool_use_id: &str) -> super::PreToolUseRequest { + super::PreToolUseRequest { + session_id: ThreadId::new(), + turn_id: "turn-1".to_string(), + subagent: None, + cwd: test_path_buf("/tmp").abs(), + transcript_path: None, + model: "gpt-test".to_string(), + permission_mode: "default".to_string(), + tool_name: "Bash".to_string(), + matcher_aliases: Vec::new(), + tool_use_id: tool_use_id.to_string(), + tool_input: serde_json::json!({ "command": "echo hello" }), + } + } +} diff --git a/codex-rs/hooks/src/events/session_end.rs b/codex-rs/hooks/src/events/session_end.rs new file mode 100644 index 0000000000000000000000000000000000000000..708d0e53d98a5743b76d284a4de261ed320de3e3 --- /dev/null +++ b/codex-rs/hooks/src/events/session_end.rs @@ -0,0 +1,139 @@ +use std::path::PathBuf; + +use codex_protocol::ThreadId; +use codex_protocol::protocol::HookCompletedEvent; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookOutputEntry; +use codex_protocol::protocol::HookOutputEntryKind; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookRunSummary; +use codex_utils_absolute_path::AbsolutePathBuf; + +use super::common; +use crate::engine::ClaudeHooksEngine; +use crate::engine::ConfiguredHandler; +use crate::engine::HandlerRunResult; +use crate::engine::dispatcher; +use crate::schema::NullableString; +use crate::schema::SessionEndCommandInput; + +pub(crate) const SESSION_END_DEFAULT_TIMEOUT_SEC: u64 = 1; +/// Keep below app-server's in-process `SHUTDOWN_TIMEOUT`: SessionEnd runs during +/// teardown and must leave headroom within the existing five-second bound. +pub(crate) const SESSION_END_MAX_TIMEOUT_SEC: u64 = 3; +const SESSION_END_REASON: &str = "other"; + +#[derive(Debug, Clone)] +pub struct SessionEndRequest { + pub session_id: ThreadId, + pub turn_id: String, + pub cwd: AbsolutePathBuf, + pub transcript_path: Option, +} + +#[derive(Debug, Default)] +pub struct SessionEndOutcome { + pub hook_events: Vec, +} + +pub(crate) fn preview(handlers: &[ConfiguredHandler]) -> Vec { + dispatcher::select_handlers( + handlers, + HookEventName::SessionEnd, + Some(SESSION_END_REASON), + ) + .into_iter() + .map(|handler| dispatcher::running_summary(&handler)) + .collect() +} + +pub(crate) async fn run( + engine: &ClaudeHooksEngine, + request: SessionEndRequest, +) -> SessionEndOutcome { + let matched = dispatcher::select_handlers( + &engine.handlers, + HookEventName::SessionEnd, + Some(SESSION_END_REASON), + ); + if matched.is_empty() { + return SessionEndOutcome::default(); + } + + let input_json = match serde_json::to_string(&SessionEndCommandInput { + session_id: request.session_id.to_string(), + transcript_path: NullableString::from_path(request.transcript_path.clone()), + cwd: request.cwd.display().to_string(), + hook_event_name: "SessionEnd".to_string(), + reason: SESSION_END_REASON.to_string(), + }) { + Ok(input_json) => input_json, + Err(error) => { + return SessionEndOutcome { + hook_events: common::serialization_failure_hook_events( + matched, + Some(request.turn_id.clone()), + format!("failed to serialize session end hook input: {error}"), + ), + }; + } + }; + + let results = dispatcher::execute_handlers( + engine, + matched, + input_json, + request.cwd.as_path(), + Some(request.turn_id), + parse_completed, + ) + .await; + SessionEndOutcome { + hook_events: results.into_iter().map(|result| result.completed).collect(), + } +} + +fn parse_completed( + handler: &ConfiguredHandler, + run_result: HandlerRunResult, + turn_id: Option, +) -> dispatcher::ParsedHandler<()> { + let (status, entries) = match (run_result.error.as_deref(), run_result.exit_code) { + (Some(error), _) => ( + HookRunStatus::Failed, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: error.to_string(), + }], + ), + (None, Some(0)) => (HookRunStatus::Completed, Vec::new()), + (None, Some(code)) => ( + HookRunStatus::Failed, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: common::trimmed_non_empty(&run_result.stderr) + .unwrap_or_else(|| format!("hook exited with code {code}")), + }], + ), + (None, None) => ( + HookRunStatus::Failed, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook process terminated without an exit code".to_string(), + }], + ), + }; + + dispatcher::ParsedHandler { + completed: HookCompletedEvent { + turn_id, + run: dispatcher::completed_summary(handler, &run_result, status, entries), + }, + data: (), + completion_order: 0, + } +} + +#[cfg(test)] +#[path = "session_end_tests.rs"] +mod tests; diff --git a/codex-rs/hooks/src/events/session_end_tests.rs b/codex-rs/hooks/src/events/session_end_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..0c100ee4a4c6fcec5f5ca158a13f7fa66247f69e --- /dev/null +++ b/codex-rs/hooks/src/events/session_end_tests.rs @@ -0,0 +1,78 @@ +use std::collections::HashMap; + +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookSource; +use codex_utils_absolute_path::test_support::PathBufExt; +use codex_utils_absolute_path::test_support::test_path_buf; +use pretty_assertions::assert_eq; + +use super::parse_completed; +use super::preview; +use crate::engine::ConfiguredHandler; +use crate::engine::HandlerRunResult; + +#[test] +fn session_end_matches_other_reason() { + let selected = preview(&[ + ConfiguredHandler { + display_order: 0, + ..handler(Some("clear")) + }, + ConfiguredHandler { + display_order: 1, + ..handler(Some("other")) + }, + ConfiguredHandler { + display_order: 2, + ..handler(/*matcher*/ None) + }, + ]); + + assert_eq!( + selected + .iter() + .map(|run| run.display_order) + .collect::>(), + vec![1, 2] + ); +} + +#[test] +fn session_end_ignores_successful_output() { + let completed = parse_completed( + &handler(/*matcher*/ None), + HandlerRunResult { + started_at: 1, + completed_at: 2, + duration_ms: 1, + exit_code: Some(0), + stdout: r#"{"continue":false,"decision":"block","reason":"ignored"}"#.to_string(), + stderr: String::new(), + error: None, + }, + /*turn_id*/ None, + ); + + assert_eq!(completed.completed.run.status, HookRunStatus::Completed); + assert_eq!(completed.completed.run.entries, Vec::new()); +} + +fn handler(matcher: Option<&str>) -> ConfiguredHandler { + ConfiguredHandler { + builtin: false, + event_name: HookEventName::SessionEnd, + matcher: matcher.map(str::to_string), + timeout_sec: 2, + status_message: None, + additional_context_limit: Default::default(), + source_path: test_path_buf("/tmp/hooks.json").abs().into(), + source: HookSource::User, + display_order: 0, + kind: crate::engine::ConfiguredHandlerKind::Command { + command: "echo hook".to_string(), + r#async: false, + env: HashMap::new(), + }, + } +} diff --git a/codex-rs/hooks/src/events/session_start.rs b/codex-rs/hooks/src/events/session_start.rs new file mode 100644 index 0000000000000000000000000000000000000000..3aff0fe946eb64a0fc042522c5f60e465fef0d7c --- /dev/null +++ b/codex-rs/hooks/src/events/session_start.rs @@ -0,0 +1,573 @@ +use std::path::PathBuf; + +use codex_protocol::ThreadId; +use codex_protocol::protocol::HookCompletedEvent; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookOutputEntry; +use codex_protocol::protocol::HookOutputEntryKind; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookRunSummary; +use codex_utils_absolute_path::AbsolutePathBuf; + +use super::common; +use crate::engine::ClaudeHooksEngine; +use crate::engine::ConfiguredHandler; +use crate::engine::HandlerRunResult; +use crate::engine::dispatcher; +use crate::engine::output_parser; +use crate::output_spill::AdditionalContext; +use crate::schema::NullableString; +use crate::schema::SessionStartCommandInput; +use crate::schema::SubagentStartCommandInput; + +#[derive(Debug, Clone, Copy)] +pub enum SessionStartSource { + Startup, + Resume, + Clear, + Compact, + Fork, +} + +impl SessionStartSource { + pub fn as_str(self) -> &'static str { + match self { + Self::Startup => "startup", + Self::Resume => "resume", + Self::Clear => "clear", + Self::Compact => "compact", + Self::Fork => "fork", + } + } +} + +#[derive(Debug, Clone)] +pub struct SessionStartRequest { + pub session_id: ThreadId, + pub cwd: AbsolutePathBuf, + pub transcript_path: Option, + pub model: String, + pub permission_mode: String, + pub target: StartHookTarget, +} + +#[derive(Debug, Clone)] +pub enum StartHookTarget { + SessionStart { + source: SessionStartSource, + }, + SubagentStart { + turn_id: String, + agent_id: String, + agent_type: String, + }, +} + +impl StartHookTarget { + fn event_name(&self) -> HookEventName { + match self { + Self::SessionStart { .. } => HookEventName::SessionStart, + Self::SubagentStart { .. } => HookEventName::SubagentStart, + } + } + + fn matcher_input(&self) -> &str { + match self { + Self::SessionStart { source } => source.as_str(), + Self::SubagentStart { agent_type, .. } => agent_type.as_str(), + } + } +} + +#[derive(Debug)] +pub struct SessionStartOutcome { + pub hook_events: Vec, + pub should_stop: bool, + pub stop_reason: Option, + pub additional_contexts: Vec, +} + +#[derive(Debug, PartialEq, Eq)] +struct SessionStartHandlerData { + should_stop: bool, + stop_reason: Option, + additional_contexts_for_model: Vec, +} + +pub(crate) fn preview( + handlers: &[ConfiguredHandler], + request: &SessionStartRequest, +) -> Vec { + dispatcher::select_handlers( + handlers, + request.target.event_name(), + Some(request.target.matcher_input()), + ) + .into_iter() + .map(|handler| dispatcher::running_summary(&handler)) + .collect() +} + +pub(crate) async fn run( + engine: &ClaudeHooksEngine, + request: SessionStartRequest, + turn_id: Option, +) -> SessionStartOutcome { + let matched = dispatcher::select_handlers( + &engine.handlers, + request.target.event_name(), + Some(request.target.matcher_input()), + ); + if matched.is_empty() { + return SessionStartOutcome { + hook_events: Vec::new(), + should_stop: false, + stop_reason: None, + additional_contexts: Vec::new(), + }; + } + + let (input_json, turn_id) = match request.target { + StartHookTarget::SessionStart { source } => { + let input_json = match serde_json::to_string(&SessionStartCommandInput::new( + request.session_id.to_string(), + request.transcript_path.clone(), + request.cwd.display().to_string(), + request.model.clone(), + request.permission_mode.clone(), + source.as_str().to_string(), + )) { + Ok(input_json) => input_json, + Err(error) => { + return serialization_failure_outcome( + common::serialization_failure_hook_events( + matched, + turn_id, + format!("failed to serialize session start hook input: {error}"), + ), + ); + } + }; + (input_json, turn_id) + } + StartHookTarget::SubagentStart { + turn_id: subagent_turn_id, + agent_id, + agent_type, + } => { + let input = SubagentStartCommandInput { + session_id: request.session_id.to_string(), + turn_id: subagent_turn_id.clone(), + transcript_path: NullableString::from_path(request.transcript_path.clone()), + cwd: request.cwd.display().to_string(), + hook_event_name: "SubagentStart".to_string(), + model: request.model.clone(), + permission_mode: request.permission_mode.clone(), + agent_id, + agent_type, + }; + let input_json = match serde_json::to_string(&input) { + Ok(input_json) => input_json, + Err(error) => { + return serialization_failure_outcome( + common::serialization_failure_hook_events( + matched, + Some(subagent_turn_id), + format!("failed to serialize subagent start hook input: {error}"), + ), + ); + } + }; + (input_json, Some(subagent_turn_id)) + } + }; + + let results = dispatcher::execute_handlers( + engine, + matched, + input_json, + request.cwd.as_path(), + turn_id, + parse_completed, + ) + .await; + + let should_stop = results.iter().any(|result| result.data.should_stop); + let stop_reason = results + .iter() + .find_map(|result| result.data.stop_reason.clone()); + let additional_contexts = common::flatten_additional_contexts( + results + .iter() + .map(|result| result.data.additional_contexts_for_model.as_slice()), + ); + let additional_contexts = engine + .command_runtime + .output_spiller() + .maybe_spill_additional_contexts(additional_contexts) + .await; + + SessionStartOutcome { + hook_events: results.into_iter().map(|result| result.completed).collect(), + should_stop, + stop_reason, + additional_contexts, + } +} + +/// Interprets completed `SessionStart` and `SubagentStart` hook runs. +/// +/// The two events have different input payloads but share most output +/// handling: hook JSON can emit warnings/context, invalid JSON-looking stdout +/// fails, and plain stdout becomes model context. Only `SessionStart` honors +/// `continue:false`; `SubagentStart` stays context-injection-only. +fn parse_completed( + handler: &ConfiguredHandler, + run_result: HandlerRunResult, + turn_id: Option, +) -> dispatcher::ParsedHandler { + let mut entries = Vec::new(); + let mut status = HookRunStatus::Completed; + let mut should_stop = false; + let mut stop_reason = None; + let mut additional_contexts_for_model = Vec::new(); + + match run_result.error.as_deref() { + Some(error) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: error.to_string(), + }); + } + None => match run_result.exit_code { + Some(0) => { + let trimmed_stdout = run_result.stdout.trim(); + if trimmed_stdout.is_empty() { + } else if let Some(parsed) = match handler.event_name { + HookEventName::SessionStart => { + output_parser::parse_session_start(&run_result.stdout) + } + HookEventName::SubagentStart => { + output_parser::parse_subagent_start(&run_result.stdout) + } + event_name => { + panic!("expected start hook event, got {event_name:?}") + } + } { + if let Some(system_message) = parsed.universal.system_message { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Warning, + text: system_message, + }); + } + if let Some(additional_context) = parsed.additional_context { + common::append_additional_context( + &mut entries, + &mut additional_contexts_for_model, + handler, + additional_context, + ); + } + let _ = parsed.universal.suppress_output; + if handler.can_apply_control_effects() + && handler.event_name == HookEventName::SessionStart + && !parsed.universal.continue_processing + { + status = HookRunStatus::Stopped; + should_stop = true; + stop_reason = parsed.universal.stop_reason.clone(); + if let Some(stop_reason_text) = parsed.universal.stop_reason { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Stop, + text: stop_reason_text, + }); + } + } + } else if output_parser::looks_like_json(&run_result.stdout) { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: match handler.event_name { + HookEventName::SessionStart => { + "hook returned invalid session start JSON output" + } + HookEventName::SubagentStart => { + "hook returned invalid subagent start JSON output" + } + event_name => { + panic!("expected start hook event, got {event_name:?}") + } + } + .to_string(), + }); + } else { + let additional_context = trimmed_stdout.to_string(); + common::append_additional_context( + &mut entries, + &mut additional_contexts_for_model, + handler, + additional_context, + ); + } + } + Some(exit_code) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: format!("hook exited with code {exit_code}"), + }); + } + None => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook exited without a status code".to_string(), + }); + } + }, + } + + let completed = HookCompletedEvent { + turn_id, + run: dispatcher::completed_summary(handler, &run_result, status, entries), + }; + + dispatcher::ParsedHandler { + completed, + data: SessionStartHandlerData { + should_stop, + stop_reason, + additional_contexts_for_model, + }, + completion_order: 0, + } +} + +fn serialization_failure_outcome(hook_events: Vec) -> SessionStartOutcome { + SessionStartOutcome { + hook_events, + should_stop: false, + stop_reason: None, + additional_contexts: Vec::new(), + } +} + +#[cfg(test)] +mod tests { + use codex_protocol::protocol::HookEventName; + use codex_protocol::protocol::HookOutputEntry; + use codex_protocol::protocol::HookOutputEntryKind; + use codex_protocol::protocol::HookRunStatus; + use codex_utils_absolute_path::test_support::PathBufExt; + use codex_utils_absolute_path::test_support::test_path_buf; + use pretty_assertions::assert_eq; + + use super::SessionStartHandlerData; + use super::parse_completed; + use crate::engine::ConfiguredHandler; + use crate::engine::HandlerRunResult; + use crate::output_spill::AdditionalContext; + use crate::output_spill::AdditionalContextLimit; + + #[test] + fn plain_stdout_becomes_model_context() { + let mut handler = handler(); + handler.additional_context_limit = AdditionalContextLimit::from_config(Some(7)); + let parsed = parse_completed( + &handler, + run_result(Some(0), "hello from hook\n", ""), + /*turn_id*/ None, + ); + + assert_eq!( + parsed.data, + SessionStartHandlerData { + should_stop: false, + stop_reason: None, + additional_contexts_for_model: vec![AdditionalContext { + text: "hello from hook".to_string(), + limit: AdditionalContextLimit::from_config(Some(7)), + }], + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Context, + text: "hello from hook".to_string(), + }] + ); + } + + #[test] + fn continue_false_preserves_context_for_later_turns() { + let parsed = parse_completed( + &handler(), + run_result( + Some(0), + r#"{"continue":false,"stopReason":"pause","hookSpecificOutput":{"hookEventName":"SessionStart","additionalContext":"do not inject"}}"#, + "", + ), + /*turn_id*/ None, + ); + + assert_eq!( + parsed.data, + SessionStartHandlerData { + should_stop: true, + stop_reason: Some("pause".to_string()), + additional_contexts_for_model: vec![AdditionalContext { + text: "do not inject".to_string(), + limit: Default::default(), + }], + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Stopped); + assert_eq!( + parsed.completed.run.entries, + vec![ + HookOutputEntry { + kind: HookOutputEntryKind::Context, + text: "do not inject".to_string(), + }, + HookOutputEntry { + kind: HookOutputEntryKind::Stop, + text: "pause".to_string(), + }, + ] + ); + } + + #[test] + fn invalid_json_like_stdout_fails_instead_of_becoming_model_context() { + let parsed = parse_completed( + &handler(), + run_result( + Some(0), + r#"{"hookSpecificOutput":{"hookEventName":"SessionStart""#, + "", + ), + /*turn_id*/ None, + ); + + assert_eq!( + parsed.data, + SessionStartHandlerData { + should_stop: false, + stop_reason: None, + additional_contexts_for_model: Vec::new(), + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook returned invalid session start JSON output".to_string(), + }] + ); + } + + #[test] + fn subagent_start_plain_stdout_becomes_model_context() { + let mut handler = handler_for(HookEventName::SubagentStart); + handler.additional_context_limit = AdditionalContextLimit::from_config(Some(4_096)); + let parsed = parse_completed( + &handler, + run_result(Some(0), "hello from subagent hook\n", ""), + /*turn_id*/ Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + SessionStartHandlerData { + should_stop: false, + stop_reason: None, + additional_contexts_for_model: vec![AdditionalContext { + text: "hello from subagent hook".to_string(), + limit: AdditionalContextLimit::from_config(Some(4_096)), + }], + } + ); + assert_eq!(parsed.completed.turn_id.as_deref(), Some("turn-1")); + assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Context, + text: "hello from subagent hook".to_string(), + }] + ); + } + + #[test] + fn subagent_start_continue_false_is_ignored() { + let parsed = parse_completed( + &handler_for(HookEventName::SubagentStart), + run_result( + Some(0), + r#"{"continue":false,"stopReason":"skip child","hookSpecificOutput":{"hookEventName":"SubagentStart","additionalContext":"child context"}}"#, + "", + ), + /*turn_id*/ Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + SessionStartHandlerData { + should_stop: false, + stop_reason: None, + additional_contexts_for_model: vec![AdditionalContext { + text: "child context".to_string(), + limit: Default::default(), + }], + } + ); + assert_eq!(parsed.completed.turn_id.as_deref(), Some("turn-1")); + assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Context, + text: "child context".to_string(), + }] + ); + } + + fn handler() -> ConfiguredHandler { + handler_for(HookEventName::SessionStart) + } + + fn handler_for(event_name: HookEventName) -> ConfiguredHandler { + ConfiguredHandler { + builtin: false, + event_name, + matcher: None, + timeout_sec: 600, + status_message: None, + additional_context_limit: Default::default(), + source_path: test_path_buf("/tmp/hooks.json").abs().into(), + source: codex_protocol::protocol::HookSource::User, + display_order: 0, + kind: crate::engine::ConfiguredHandlerKind::Command { + command: "echo hook".to_string(), + r#async: false, + env: std::collections::HashMap::new(), + }, + } + } + + fn run_result(exit_code: Option, stdout: &str, stderr: &str) -> HandlerRunResult { + HandlerRunResult { + started_at: 1, + completed_at: 2, + duration_ms: 1, + exit_code, + stdout: stdout.to_string(), + stderr: stderr.to_string(), + error: None, + } + } +} diff --git a/codex-rs/hooks/src/events/stop.rs b/codex-rs/hooks/src/events/stop.rs new file mode 100644 index 0000000000000000000000000000000000000000..be2be9100e356427b298f24872113852eab3862a --- /dev/null +++ b/codex-rs/hooks/src/events/stop.rs @@ -0,0 +1,723 @@ +use std::path::PathBuf; + +use codex_protocol::ThreadId; +use codex_protocol::items::HookPromptFragment; +use codex_protocol::protocol::HookCompletedEvent; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookOutputEntry; +use codex_protocol::protocol::HookOutputEntryKind; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookRunSummary; +use codex_protocol::protocol::HookSource; +use codex_utils_absolute_path::AbsolutePathBuf; +use serde_json::Map; +use serde_json::Value; + +use super::common; +use crate::engine::ClaudeHooksEngine; +use crate::engine::ConfiguredHandler; +use crate::engine::HandlerRunResult; +use crate::engine::HandlerSourcePath; +use crate::engine::dispatcher; +use crate::engine::output_parser; +use crate::schema::NullableString; +use crate::schema::StopCommandInput; +use crate::schema::SubagentStopCommandInput; + +#[derive(Debug, Clone)] +pub struct StopRequest { + pub session_id: ThreadId, + pub turn_id: String, + pub cwd: AbsolutePathBuf, + pub transcript_path: Option, + pub model: String, + pub permission_mode: String, + pub request_metadata: Option>, + pub stop_hook_active: bool, + pub last_assistant_message: Option, + pub target: StopHookTarget, +} + +#[derive(Debug, Clone)] +pub enum StopHookTarget { + Stop, + /// Internal memory work runs policy and executor hooks, not project completion checks. + MemoryConsolidation, + SubagentStop { + agent_id: String, + agent_type: String, + agent_transcript_path: Option, + }, +} + +impl StopHookTarget { + fn event_name(&self) -> HookEventName { + match self { + Self::Stop | Self::MemoryConsolidation => HookEventName::Stop, + Self::SubagentStop { .. } => HookEventName::SubagentStop, + } + } + + fn matcher_input(&self) -> Option<&str> { + match self { + Self::Stop | Self::MemoryConsolidation => None, + Self::SubagentStop { agent_type, .. } => Some(agent_type.as_str()), + } + } + + fn select_handlers(&self, handlers: &[ConfiguredHandler]) -> Vec { + dispatcher::select_handlers(handlers, self.event_name(), self.matcher_input()) + .into_iter() + .filter(|handler| { + !matches!(self, Self::MemoryConsolidation) + || matches!( + handler.source_path, + HandlerSourcePath::ExecutorScoped { .. } + ) + || match handler.source { + HookSource::User + | HookSource::Project + | HookSource::SessionFlags + | HookSource::Plugin => false, + HookSource::System + | HookSource::Mdm + | HookSource::CloudRequirements + | HookSource::CloudManagedConfig + | HookSource::LegacyManagedConfigFile + | HookSource::LegacyManagedConfigMdm + // Required hooks can have unknown attribution; retain them fail-closed. + | HookSource::Unknown => true, + } + }) + .collect() + } +} + +#[derive(Debug, Default)] +pub struct StopOutcome { + pub hook_events: Vec, + pub should_stop: bool, + pub stop_reason: Option, + pub should_block: bool, + pub block_reason: Option, + pub continuation_fragments: Vec, +} + +#[derive(Debug, Default, PartialEq, Eq)] +struct StopHandlerData { + should_stop: bool, + stop_reason: Option, + should_block: bool, + block_reason: Option, + continuation_fragments: Vec, +} + +pub(crate) fn preview( + handlers: &[ConfiguredHandler], + request: &StopRequest, +) -> Vec { + request + .target + .select_handlers(handlers) + .into_iter() + .filter(|handler| matches!(handler.source_path, HandlerSourcePath::Local(_))) + .map(|handler| dispatcher::running_summary(&handler)) + .collect() +} + +pub(crate) async fn run(engine: &ClaudeHooksEngine, request: StopRequest) -> StopOutcome { + let matched = request.target.select_handlers(&engine.handlers); + if matched.is_empty() { + return StopOutcome { + hook_events: Vec::new(), + should_stop: false, + stop_reason: None, + should_block: false, + block_reason: None, + continuation_fragments: Vec::new(), + }; + } + + // Memory workers terminate on managed rejection rather than continuing the turn, + // so their executor cleanup must also run when a managed hook blocks completion. + let (executor_cleanup, matched): (Vec<_>, Vec<_>) = matched.into_iter().partition(|handler| { + matches!(request.target, StopHookTarget::MemoryConsolidation) + && matches!( + handler.source_path, + HandlerSourcePath::ExecutorScoped { .. } + ) + }); + let input_json = match request.target { + StopHookTarget::Stop | StopHookTarget::MemoryConsolidation => { + let input = StopCommandInput { + session_id: request.session_id.to_string(), + turn_id: request.turn_id.clone(), + transcript_path: NullableString::from_path(request.transcript_path.clone()), + cwd: request.cwd.display().to_string(), + hook_event_name: "Stop".to_string(), + model: request.model.clone(), + permission_mode: request.permission_mode.clone(), + stop_hook_active: request.stop_hook_active, + last_assistant_message: NullableString::from_string( + request.last_assistant_message.clone(), + ), + }; + match serde_json::to_string(&input) { + Ok(input_json) => input_json, + Err(error) => { + return serialization_failure_outcome( + common::serialization_failure_hook_events( + matched, + Some(request.turn_id), + format!("failed to serialize stop hook input: {error}"), + ), + ); + } + } + } + StopHookTarget::SubagentStop { + agent_id, + agent_type, + agent_transcript_path, + } => { + let input = SubagentStopCommandInput { + session_id: request.session_id.to_string(), + turn_id: request.turn_id.clone(), + transcript_path: NullableString::from_path(request.transcript_path.clone()), + agent_transcript_path: NullableString::from_path(agent_transcript_path), + cwd: request.cwd.display().to_string(), + hook_event_name: "SubagentStop".to_string(), + model: request.model.clone(), + permission_mode: request.permission_mode.clone(), + stop_hook_active: request.stop_hook_active, + agent_id, + agent_type, + last_assistant_message: NullableString::from_string( + request.last_assistant_message.clone(), + ), + }; + match serde_json::to_string(&input) { + Ok(input_json) => input_json, + Err(error) => { + return serialization_failure_outcome( + common::serialization_failure_hook_events( + matched, + Some(request.turn_id), + format!("failed to serialize subagent stop hook input: {error}"), + ), + ); + } + } + } + }; + + let results = dispatcher::execute_handlers_with_metadata( + engine, + matched, + input_json.clone(), + request.cwd.as_path(), + Some(request.turn_id.clone()), + request.request_metadata.as_ref(), + parse_completed, + ) + .await; + + if !executor_cleanup.is_empty() { + dispatcher::execute_handlers_with_metadata( + engine, + executor_cleanup, + input_json, + request.cwd.as_path(), + Some(request.turn_id), + request.request_metadata.as_ref(), + parse_completed, + ) + .await; + } + + let aggregate = aggregate_results(results.iter().map(|result| &result.data)); + + StopOutcome { + hook_events: results.into_iter().map(|result| result.completed).collect(), + should_stop: aggregate.should_stop, + stop_reason: aggregate.stop_reason, + should_block: aggregate.should_block, + block_reason: aggregate.block_reason, + continuation_fragments: aggregate.continuation_fragments, + } +} + +fn parse_completed( + handler: &ConfiguredHandler, + run_result: HandlerRunResult, + turn_id: Option, +) -> dispatcher::ParsedHandler { + let mut entries = Vec::new(); + let mut status = HookRunStatus::Completed; + let mut should_stop = false; + let mut stop_reason = None; + let mut should_block = false; + let mut block_reason = None; + let mut continuation_prompt = None; + let hook_event_name = match handler.event_name { + HookEventName::Stop | HookEventName::SubagentStop => handler.event_name, + event_name => { + panic!("expected stop hook event, got {event_name:?}"); + } + }; + + match run_result.error.as_deref() { + Some(error) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: error.to_string(), + }); + } + None => match run_result.exit_code { + Some(0) => { + let trimmed_stdout = run_result.stdout.trim(); + if trimmed_stdout.is_empty() { + } else if let Some(parsed) = match hook_event_name { + HookEventName::Stop => output_parser::parse_stop(&run_result.stdout), + HookEventName::SubagentStop => { + output_parser::parse_subagent_stop(&run_result.stdout) + } + _ => unreachable!("validated stop hook event"), + } { + if let Some(system_message) = parsed.universal.system_message { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Warning, + text: system_message, + }); + } + let _ = parsed.universal.suppress_output; + if handler.can_apply_control_effects() { + if !parsed.universal.continue_processing { + status = HookRunStatus::Stopped; + should_stop = true; + stop_reason = parsed.universal.stop_reason.clone(); + if let Some(stop_reason_text) = parsed.universal.stop_reason { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Stop, + text: stop_reason_text, + }); + } + } else if let Some(invalid_block_reason) = parsed.invalid_block_reason { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: invalid_block_reason, + }); + } else if parsed.should_block + && let Some(reason) = + parsed.reason.as_deref().and_then(common::trimmed_non_empty) + { + status = HookRunStatus::Blocked; + should_block = true; + block_reason = Some(reason.clone()); + continuation_prompt = Some(reason.clone()); + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: reason, + }); + } + } + } else if handler.can_apply_control_effects() + || output_parser::looks_like_json(&run_result.stdout) + { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: match hook_event_name { + HookEventName::Stop => "hook returned invalid stop hook JSON output", + HookEventName::SubagentStop => { + "hook returned invalid subagent stop hook JSON output" + } + _ => unreachable!("validated stop hook event"), + } + .to_string(), + }); + } + } + Some(2) if handler.can_apply_control_effects() => { + if let Some(reason) = common::trimmed_non_empty(&run_result.stderr) { + status = HookRunStatus::Blocked; + should_block = true; + block_reason = Some(reason.clone()); + continuation_prompt = Some(reason.clone()); + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: reason, + }); + } else { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: match hook_event_name { + HookEventName::Stop => { + "Stop hook exited with code 2 but did not write a continuation prompt to stderr" + } + HookEventName::SubagentStop => { + "SubagentStop hook exited with code 2 but did not write a continuation prompt to stderr" + } + _ => unreachable!("validated stop hook event"), + } + .to_string(), + }); + } + } + Some(exit_code) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: format!("hook exited with code {exit_code}"), + }); + } + None => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook exited without a status code".to_string(), + }); + } + }, + } + + let completed = HookCompletedEvent { + turn_id, + run: dispatcher::completed_summary(handler, &run_result, status, entries), + }; + let continuation_fragments = continuation_prompt + .map(|prompt| { + vec![HookPromptFragment::from_single_hook( + prompt, + completed.run.id.clone(), + )] + }) + .unwrap_or_default(); + + dispatcher::ParsedHandler { + completed, + data: StopHandlerData { + should_stop, + stop_reason, + should_block, + block_reason, + continuation_fragments, + }, + completion_order: 0, + } +} + +fn aggregate_results<'a>( + results: impl IntoIterator, +) -> StopHandlerData { + let results = results.into_iter().collect::>(); + let should_stop = results.iter().any(|result| result.should_stop); + let stop_reason = results.iter().find_map(|result| result.stop_reason.clone()); + let should_block = !should_stop && results.iter().any(|result| result.should_block); + let block_reason = if should_block { + common::join_text_chunks( + results + .iter() + .filter_map(|result| result.block_reason.clone()) + .collect(), + ) + } else { + None + }; + let continuation_fragments = if should_block { + results + .iter() + .filter(|result| result.should_block) + .flat_map(|result| result.continuation_fragments.clone()) + .collect() + } else { + Vec::new() + }; + + StopHandlerData { + should_stop, + stop_reason, + should_block, + block_reason, + continuation_fragments, + } +} + +fn serialization_failure_outcome(hook_events: Vec) -> StopOutcome { + StopOutcome { + hook_events, + should_stop: false, + stop_reason: None, + should_block: false, + block_reason: None, + continuation_fragments: Vec::new(), + } +} + +#[cfg(test)] +mod tests { + use codex_protocol::protocol::HookEventName; + use codex_protocol::protocol::HookOutputEntry; + use codex_protocol::protocol::HookOutputEntryKind; + use codex_protocol::protocol::HookRunStatus; + use codex_utils_absolute_path::test_support::PathBufExt; + use codex_utils_absolute_path::test_support::test_path_buf; + use pretty_assertions::assert_eq; + + use codex_protocol::items::HookPromptFragment; + + use super::StopHandlerData; + use super::aggregate_results; + use super::parse_completed; + use crate::engine::ConfiguredHandler; + use crate::engine::HandlerRunResult; + + #[test] + fn block_decision_with_reason_sets_continuation_prompt() { + let parsed = parse_completed( + &handler(), + run_result( + Some(0), + r#"{"decision":"block","reason":"retry with tests"}"#, + "", + ), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + StopHandlerData { + should_stop: false, + stop_reason: None, + should_block: true, + block_reason: Some("retry with tests".to_string()), + continuation_fragments: vec![HookPromptFragment { + text: "retry with tests".to_string(), + hook_run_id: parsed.completed.run.id.clone(), + }], + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked); + } + + #[test] + fn block_decision_without_reason_is_invalid() { + let parsed = parse_completed( + &handler(), + run_result(Some(0), r#"{"decision":"block"}"#, ""), + Some("turn-1".to_string()), + ); + + assert_eq!(parsed.data, StopHandlerData::default()); + assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "Stop hook returned decision:block without a non-empty reason".to_string(), + }] + ); + + let async_handler = handler_with_async(/*async*/ true); + let parsed = parse_completed( + &async_handler, + run_result(Some(0), r#"{"decision":"block"}"#, ""), + Some("turn-1".to_string()), + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); + assert_eq!(parsed.completed.run.entries, Vec::new()); + } + + #[test] + fn continue_false_overrides_block_decision() { + let parsed = parse_completed( + &handler(), + run_result( + Some(0), + r#"{"continue":false,"stopReason":"done","decision":"block","reason":"keep going"}"#, + "", + ), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + StopHandlerData { + should_stop: true, + stop_reason: Some("done".to_string()), + should_block: false, + block_reason: None, + continuation_fragments: Vec::new(), + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Stopped); + } + + #[test] + fn exit_code_two_uses_stderr_feedback_only() { + let parsed = parse_completed( + &handler(), + run_result(Some(2), "ignored stdout", "retry with tests"), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + StopHandlerData { + should_stop: false, + stop_reason: None, + should_block: true, + block_reason: Some("retry with tests".to_string()), + continuation_fragments: vec![HookPromptFragment { + text: "retry with tests".to_string(), + hook_run_id: parsed.completed.run.id.clone(), + }], + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked); + } + + #[test] + fn exit_code_two_without_stderr_does_not_block() { + let parsed = parse_completed( + &handler(), + run_result(Some(2), "", " "), + /*turn_id*/ None, + ); + + assert_eq!(parsed.data, StopHandlerData::default()); + assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: + "Stop hook exited with code 2 but did not write a continuation prompt to stderr" + .to_string(), + }] + ); + } + + #[test] + fn block_decision_with_blank_reason_fails_instead_of_blocking() { + let parsed = parse_completed( + &handler(), + run_result(Some(0), "{\"decision\":\"block\",\"reason\":\" \"}", ""), + Some("turn-1".to_string()), + ); + + assert_eq!(parsed.data, StopHandlerData::default()); + assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "Stop hook returned decision:block without a non-empty reason".to_string(), + }] + ); + } + + #[test] + fn invalid_stdout_fails_instead_of_silently_nooping() { + let parsed = parse_completed( + &handler(), + run_result(Some(0), "not json", ""), + Some("turn-1".to_string()), + ); + + assert_eq!(parsed.data, StopHandlerData::default()); + assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook returned invalid stop hook JSON output".to_string(), + }] + ); + + let async_handler = handler_with_async(/*async*/ true); + let parsed = parse_completed( + &async_handler, + run_result(Some(0), "not json", ""), + Some("turn-1".to_string()), + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); + assert_eq!(parsed.completed.run.entries, Vec::new()); + } + + #[test] + fn aggregate_results_concatenates_blocking_reasons_in_declaration_order() { + let aggregate = aggregate_results([ + &StopHandlerData { + should_stop: false, + stop_reason: None, + should_block: true, + block_reason: Some("first".to_string()), + continuation_fragments: vec![HookPromptFragment::from_single_hook( + "first", "run-1", + )], + }, + &StopHandlerData { + should_stop: false, + stop_reason: None, + should_block: true, + block_reason: Some("second".to_string()), + continuation_fragments: vec![HookPromptFragment::from_single_hook( + "second", "run-2", + )], + }, + ]); + + assert_eq!( + aggregate, + StopHandlerData { + should_stop: false, + stop_reason: None, + should_block: true, + block_reason: Some("first\n\nsecond".to_string()), + continuation_fragments: vec![ + HookPromptFragment::from_single_hook("first", "run-1"), + HookPromptFragment::from_single_hook("second", "run-2"), + ], + } + ); + } + + fn handler() -> ConfiguredHandler { + handler_with_async(/*async*/ false) + } + + fn handler_with_async(r#async: bool) -> ConfiguredHandler { + ConfiguredHandler { + builtin: false, + event_name: HookEventName::Stop, + matcher: None, + timeout_sec: 600, + status_message: None, + additional_context_limit: Default::default(), + source_path: test_path_buf("/tmp/hooks.json").abs().into(), + source: codex_protocol::protocol::HookSource::User, + display_order: 0, + kind: crate::engine::ConfiguredHandlerKind::Command { + command: "echo hook".to_string(), + r#async, + env: std::collections::HashMap::new(), + }, + } + } + + fn run_result(exit_code: Option, stdout: &str, stderr: &str) -> HandlerRunResult { + HandlerRunResult { + started_at: 1, + completed_at: 2, + duration_ms: 1, + exit_code, + stdout: stdout.to_string(), + stderr: stderr.to_string(), + error: None, + } + } +} diff --git a/codex-rs/hooks/src/events/user_prompt_submit.rs b/codex-rs/hooks/src/events/user_prompt_submit.rs new file mode 100644 index 0000000000000000000000000000000000000000..9124431bc194284ef4a009c76008851a984ff319 --- /dev/null +++ b/codex-rs/hooks/src/events/user_prompt_submit.rs @@ -0,0 +1,492 @@ +use std::path::PathBuf; + +use codex_protocol::ThreadId; +use codex_protocol::protocol::HookCompletedEvent; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookOutputEntry; +use codex_protocol::protocol::HookOutputEntryKind; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookRunSummary; +use codex_utils_absolute_path::AbsolutePathBuf; + +use super::common; +use crate::engine::ClaudeHooksEngine; +use crate::engine::ConfiguredHandler; +use crate::engine::HandlerRunResult; +use crate::engine::dispatcher; +use crate::engine::output_parser; +use crate::output_spill::AdditionalContext; +use crate::schema::NullableString; +use crate::schema::SubagentCommandInputFields; +use crate::schema::UserPromptSubmitCommandInput; + +#[derive(Debug, Clone)] +pub struct UserPromptSubmitRequest { + pub session_id: ThreadId, + pub turn_id: String, + pub subagent: Option, + pub cwd: AbsolutePathBuf, + pub transcript_path: Option, + pub model: String, + pub permission_mode: String, + pub prompt: String, +} + +#[derive(Debug)] +pub struct UserPromptSubmitOutcome { + pub hook_events: Vec, + pub should_stop: bool, + pub stop_reason: Option, + pub additional_contexts: Vec, +} + +#[derive(Debug, PartialEq, Eq)] +struct UserPromptSubmitHandlerData { + should_stop: bool, + stop_reason: Option, + additional_contexts_for_model: Vec, +} + +pub(crate) fn preview( + handlers: &[ConfiguredHandler], + _request: &UserPromptSubmitRequest, +) -> Vec { + dispatcher::select_handlers( + handlers, + HookEventName::UserPromptSubmit, + /*matcher_input*/ None, + ) + .into_iter() + .map(|handler| dispatcher::running_summary(&handler)) + .collect() +} + +pub(crate) async fn run( + engine: &ClaudeHooksEngine, + request: UserPromptSubmitRequest, +) -> UserPromptSubmitOutcome { + let matched = dispatcher::select_handlers( + &engine.handlers, + HookEventName::UserPromptSubmit, + /*matcher_input*/ None, + ); + if matched.is_empty() { + return UserPromptSubmitOutcome { + hook_events: Vec::new(), + should_stop: false, + stop_reason: None, + additional_contexts: Vec::new(), + }; + } + + let subagent = SubagentCommandInputFields::from(request.subagent.as_ref()); + let input_json = match serde_json::to_string(&UserPromptSubmitCommandInput { + session_id: request.session_id.to_string(), + turn_id: request.turn_id.clone(), + agent_id: subagent.agent_id, + agent_type: subagent.agent_type, + transcript_path: NullableString::from_path(request.transcript_path.clone()), + cwd: request.cwd.display().to_string(), + hook_event_name: "UserPromptSubmit".to_string(), + model: request.model.clone(), + permission_mode: request.permission_mode.clone(), + prompt: request.prompt.clone(), + }) { + Ok(input_json) => input_json, + Err(error) => { + return serialization_failure_outcome(common::serialization_failure_hook_events( + matched, + Some(request.turn_id), + format!("failed to serialize user prompt submit hook input: {error}"), + )); + } + }; + + let results = dispatcher::execute_handlers( + engine, + matched, + input_json, + request.cwd.as_path(), + Some(request.turn_id), + parse_completed, + ) + .await; + + let should_stop = results.iter().any(|result| result.data.should_stop); + let stop_reason = results + .iter() + .find_map(|result| result.data.stop_reason.clone()); + let additional_contexts = common::flatten_additional_contexts( + results + .iter() + .map(|result| result.data.additional_contexts_for_model.as_slice()), + ); + let additional_contexts = engine + .command_runtime + .output_spiller() + .maybe_spill_additional_contexts(additional_contexts) + .await; + + UserPromptSubmitOutcome { + hook_events: results.into_iter().map(|result| result.completed).collect(), + should_stop, + stop_reason, + additional_contexts, + } +} + +fn parse_completed( + handler: &ConfiguredHandler, + run_result: HandlerRunResult, + turn_id: Option, +) -> dispatcher::ParsedHandler { + let mut entries = Vec::new(); + let mut status = HookRunStatus::Completed; + let mut should_stop = false; + let mut stop_reason = None; + let mut additional_contexts_for_model = Vec::new(); + + match run_result.error.as_deref() { + Some(error) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: error.to_string(), + }); + } + None => match run_result.exit_code { + Some(0) => { + let trimmed_stdout = run_result.stdout.trim(); + if trimmed_stdout.is_empty() { + } else if let Some(parsed) = + output_parser::parse_user_prompt_submit(&run_result.stdout) + { + if let Some(system_message) = parsed.universal.system_message { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Warning, + text: system_message, + }); + } + if (!handler.can_apply_control_effects() + || parsed.invalid_block_reason.is_none()) + && let Some(additional_context) = parsed.additional_context + { + common::append_additional_context( + &mut entries, + &mut additional_contexts_for_model, + handler, + additional_context, + ); + } + let _ = parsed.universal.suppress_output; + if handler.can_apply_control_effects() { + if !parsed.universal.continue_processing { + status = HookRunStatus::Stopped; + should_stop = true; + stop_reason = parsed.universal.stop_reason.clone(); + if let Some(stop_reason_text) = parsed.universal.stop_reason { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Stop, + text: stop_reason_text, + }); + } + } else if let Some(invalid_block_reason) = parsed.invalid_block_reason { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: invalid_block_reason, + }); + } else if parsed.should_block { + status = HookRunStatus::Blocked; + should_stop = true; + stop_reason = parsed.reason.clone(); + if let Some(reason) = parsed.reason { + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: reason, + }); + } + } + } + } else if output_parser::looks_like_json(&run_result.stdout) { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook returned invalid user prompt submit JSON output".to_string(), + }); + } else { + let additional_context = trimmed_stdout.to_string(); + common::append_additional_context( + &mut entries, + &mut additional_contexts_for_model, + handler, + additional_context, + ); + } + } + Some(2) if handler.can_apply_control_effects() => { + if let Some(reason) = common::trimmed_non_empty(&run_result.stderr) { + status = HookRunStatus::Blocked; + should_stop = true; + stop_reason = Some(reason.clone()); + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: reason, + }); + } else { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "UserPromptSubmit hook exited with code 2 but did not write a blocking reason to stderr".to_string(), + }); + } + } + Some(exit_code) => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: format!("hook exited with code {exit_code}"), + }); + } + None => { + status = HookRunStatus::Failed; + entries.push(HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "hook exited without a status code".to_string(), + }); + } + }, + } + + let completed = HookCompletedEvent { + turn_id, + run: dispatcher::completed_summary(handler, &run_result, status, entries), + }; + + dispatcher::ParsedHandler { + completed, + data: UserPromptSubmitHandlerData { + should_stop, + stop_reason, + additional_contexts_for_model, + }, + completion_order: 0, + } +} + +fn serialization_failure_outcome(hook_events: Vec) -> UserPromptSubmitOutcome { + UserPromptSubmitOutcome { + hook_events, + should_stop: false, + stop_reason: None, + additional_contexts: Vec::new(), + } +} + +#[cfg(test)] +mod tests { + use codex_protocol::protocol::HookEventName; + use codex_protocol::protocol::HookOutputEntry; + use codex_protocol::protocol::HookOutputEntryKind; + use codex_protocol::protocol::HookRunStatus; + use codex_utils_absolute_path::test_support::PathBufExt; + use codex_utils_absolute_path::test_support::test_path_buf; + use pretty_assertions::assert_eq; + + use super::UserPromptSubmitHandlerData; + use super::parse_completed; + use crate::engine::ConfiguredHandler; + use crate::engine::HandlerRunResult; + use crate::output_spill::AdditionalContext; + + #[test] + fn continue_false_preserves_context_for_later_turns() { + let parsed = parse_completed( + &handler(), + run_result( + Some(0), + r#"{"continue":false,"stopReason":"pause","hookSpecificOutput":{"hookEventName":"UserPromptSubmit","additionalContext":"do not inject"}}"#, + "", + ), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + UserPromptSubmitHandlerData { + should_stop: true, + stop_reason: Some("pause".to_string()), + additional_contexts_for_model: vec![AdditionalContext { + text: "do not inject".to_string(), + limit: Default::default(), + }], + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Stopped); + assert_eq!( + parsed.completed.run.entries, + vec![ + HookOutputEntry { + kind: HookOutputEntryKind::Context, + text: "do not inject".to_string(), + }, + HookOutputEntry { + kind: HookOutputEntryKind::Stop, + text: "pause".to_string(), + }, + ] + ); + } + + #[test] + fn claude_block_decision_blocks_processing() { + let parsed = parse_completed( + &handler(), + run_result( + Some(0), + r#"{"decision":"block","reason":"slow down","hookSpecificOutput":{"hookEventName":"UserPromptSubmit","additionalContext":"do not inject"}}"#, + "", + ), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + UserPromptSubmitHandlerData { + should_stop: true, + stop_reason: Some("slow down".to_string()), + additional_contexts_for_model: vec![AdditionalContext { + text: "do not inject".to_string(), + limit: Default::default(), + }], + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked); + assert_eq!( + parsed.completed.run.entries, + vec![ + HookOutputEntry { + kind: HookOutputEntryKind::Context, + text: "do not inject".to_string(), + }, + HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: "slow down".to_string(), + }, + ] + ); + } + + #[test] + fn claude_block_decision_requires_reason() { + let stdout = r#"{"decision":"block","hookSpecificOutput":{"hookEventName":"UserPromptSubmit","additionalContext":"do not inject"}}"#; + let parsed = parse_completed( + &handler(), + run_result(Some(0), stdout, ""), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + UserPromptSubmitHandlerData { + should_stop: false, + stop_reason: None, + additional_contexts_for_model: Vec::new(), + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Error, + text: "UserPromptSubmit hook returned decision:block without a non-empty reason" + .to_string(), + }] + ); + + let async_handler = handler_with_async(/*async*/ true); + let parsed = parse_completed( + &async_handler, + run_result(Some(0), stdout, ""), + Some("turn-1".to_string()), + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Completed); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Context, + text: "do not inject".to_string(), + }], + ); + } + + #[test] + fn exit_code_two_blocks_processing() { + let parsed = parse_completed( + &handler(), + run_result(Some(2), "", "blocked by policy\n"), + Some("turn-1".to_string()), + ); + + assert_eq!( + parsed.data, + UserPromptSubmitHandlerData { + should_stop: true, + stop_reason: Some("blocked by policy".to_string()), + additional_contexts_for_model: Vec::new(), + } + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Blocked); + assert_eq!( + parsed.completed.run.entries, + vec![HookOutputEntry { + kind: HookOutputEntryKind::Feedback, + text: "blocked by policy".to_string(), + }] + ); + + let async_handler = handler_with_async(/*async*/ true); + let parsed = parse_completed( + &async_handler, + run_result(Some(2), "", "blocked by policy\n"), + Some("turn-1".to_string()), + ); + assert_eq!(parsed.completed.run.status, HookRunStatus::Failed); + assert!(!parsed.data.should_stop); + } + + fn handler() -> ConfiguredHandler { + handler_with_async(/*async*/ false) + } + + fn handler_with_async(r#async: bool) -> ConfiguredHandler { + ConfiguredHandler { + builtin: false, + event_name: HookEventName::UserPromptSubmit, + matcher: None, + timeout_sec: 5, + status_message: None, + additional_context_limit: Default::default(), + source_path: test_path_buf("/tmp/hooks.json").abs().into(), + source: codex_protocol::protocol::HookSource::User, + display_order: 0, + kind: crate::engine::ConfiguredHandlerKind::Command { + command: "echo hook".to_string(), + r#async, + env: std::collections::HashMap::new(), + }, + } + } + + fn run_result(exit_code: Option, stdout: &str, stderr: &str) -> HandlerRunResult { + HandlerRunResult { + started_at: 1, + completed_at: 2, + duration_ms: 1, + exit_code, + stdout: stdout.to_string(), + stderr: stderr.to_string(), + error: None, + } + } +} diff --git a/codex-rs/hooks/src/legacy_notify.rs b/codex-rs/hooks/src/legacy_notify.rs new file mode 100644 index 0000000000000000000000000000000000000000..5be7d7b32ee13572cd9c0f9af276bb5b083f7159 --- /dev/null +++ b/codex-rs/hooks/src/legacy_notify.rs @@ -0,0 +1,183 @@ +use std::ffi::OsString; +use std::process::Stdio; +use std::sync::Arc; + +use serde::Serialize; + +use crate::Hook; +use crate::HookEvent; +use crate::HookPayload; +use crate::HookResult; +use crate::registry::command_from_argv; + +/// Legacy notify payload appended as the final argv argument for backward compatibility. +#[derive(Debug, Clone, PartialEq, Serialize)] +#[serde(tag = "type", rename_all = "kebab-case")] +enum UserNotification { + #[serde(rename_all = "kebab-case")] + AgentTurnComplete { + thread_id: String, + turn_id: String, + cwd: String, + #[serde(skip_serializing_if = "Option::is_none")] + client: Option, + input_messages: Vec, + last_assistant_message: Option, + }, +} + +pub fn legacy_notify_json(payload: &HookPayload) -> Result { + match &payload.hook_event { + HookEvent::AfterAgent { event } => { + serde_json::to_string(&UserNotification::AgentTurnComplete { + thread_id: event.thread_id.to_string(), + turn_id: event.turn_id.clone(), + cwd: payload.cwd.display().to_string(), + client: payload.client.clone(), + input_messages: event.input_messages.clone(), + last_assistant_message: event.last_assistant_message.clone(), + }) + } + } +} + +// TODO: Remove this hook and its environment plumbing when legacy `notify` support is removed. +pub(crate) fn notify_hook(argv: Vec, environment: Arc>) -> Hook { + let argv = Arc::new(argv); + Hook { + name: "legacy_notify".to_string(), + func: Arc::new(move |payload: &HookPayload| { + let argv = Arc::clone(&argv); + let environment = Arc::clone(&environment); + Box::pin(async move { + let mut command = match command_from_argv(&argv, environment.iter().cloned()) { + Some(command) => command, + None => return HookResult::Success, + }; + if let Ok(notify_payload) = legacy_notify_json(payload) { + command.arg(notify_payload); + } + + command + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()); + + match command.spawn() { + Ok(_) => HookResult::Success, + Err(err) => HookResult::FailedContinue(err.into()), + } + }) + }), + } +} + +#[cfg(test)] +mod tests { + use anyhow::Result; + use codex_protocol::ThreadId; + use codex_protocol::shell_environment::CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR; + use codex_utils_absolute_path::test_support::PathBufExt; + use codex_utils_absolute_path::test_support::test_path_buf; + use pretty_assertions::assert_eq; + use serde_json::Value; + use serde_json::json; + + use super::*; + use crate::HookEventAfterAgent; + + fn expected_notification_json() -> Value { + let cwd = test_path_buf("/Users/example/project"); + json!({ + "type": "agent-turn-complete", + "thread-id": "b5f6c1c2-1111-2222-3333-444455556666", + "turn-id": "12345", + "cwd": cwd.display().to_string(), + "client": "codex-tui", + "input-messages": ["Rename `foo` to `bar` and update the callsites."], + "last-assistant-message": "Rename complete and verified `cargo build` succeeds.", + }) + } + + #[test] + fn test_user_notification() -> Result<()> { + let notification = UserNotification::AgentTurnComplete { + thread_id: "b5f6c1c2-1111-2222-3333-444455556666".to_string(), + turn_id: "12345".to_string(), + cwd: test_path_buf("/Users/example/project") + .display() + .to_string(), + client: Some("codex-tui".to_string()), + input_messages: vec!["Rename `foo` to `bar` and update the callsites.".to_string()], + last_assistant_message: Some( + "Rename complete and verified `cargo build` succeeds.".to_string(), + ), + }; + let serialized = serde_json::to_string(¬ification)?; + let actual: Value = serde_json::from_str(&serialized)?; + assert_eq!(actual, expected_notification_json()); + Ok(()) + } + + #[test] + fn legacy_notify_json_matches_historical_wire_shape() -> Result<()> { + let payload = HookPayload { + session_id: ThreadId::new(), + cwd: test_path_buf("/Users/example/project").abs(), + client: Some("codex-tui".to_string()), + triggered_at: chrono::Utc::now(), + hook_event: HookEvent::AfterAgent { + event: HookEventAfterAgent { + thread_id: ThreadId::from_string("b5f6c1c2-1111-2222-3333-444455556666") + .expect("valid thread id"), + turn_id: "12345".to_string(), + input_messages: vec![ + "Rename `foo` to `bar` and update the callsites.".to_string(), + ], + last_assistant_message: Some( + "Rename complete and verified `cargo build` succeeds.".to_string(), + ), + }, + }, + }; + + let serialized = legacy_notify_json(&payload)?; + let actual: Value = serde_json::from_str(&serialized)?; + assert_eq!(actual, expected_notification_json()); + + Ok(()) + } + + #[test] + fn legacy_notify_command_replays_session_snapshot_and_scrubs_credentials() { + let argv = vec!["notify-command".to_string()]; + let environment = vec![ + ( + OsString::from("CODEX_LEGACY_NOTIFY_SNAPSHOT"), + OsString::from("captured"), + ), + ( + OsString::from(CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR), + OsString::from("restricted-token"), + ), + ]; + + let command = command_from_argv(&argv, environment) + .expect("legacy notification command should be configured"); + let configured_environment = command + .as_std() + .get_envs() + .filter_map(|(name, value)| { + value.map(|value| (name.to_os_string(), value.to_os_string())) + }) + .collect::>(); + + assert_eq!( + configured_environment, + vec![( + OsString::from("CODEX_LEGACY_NOTIFY_SNAPSHOT"), + OsString::from("captured"), + )] + ); + } +} diff --git a/codex-rs/hooks/src/lib.rs b/codex-rs/hooks/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..6abc173535b23b5ee6bd057fd3f49d9b82ea7fe0 --- /dev/null +++ b/codex-rs/hooks/src/lib.rs @@ -0,0 +1,123 @@ +mod config_rules; +mod declarations; +mod engine; +pub(crate) mod events; +mod legacy_notify; +mod mcp; +mod output_spill; +mod registry; +mod schema; +mod types; + +use codex_protocol::protocol::HookEventName; + +pub use config_rules::hook_states_from_stack; +pub use declarations::PluginHookDeclaration; +pub use declarations::plugin_hook_declarations; +pub use engine::HookListEntry; +pub use engine::HookListEntryHandler; +pub use engine::dispatcher::hook_execution_mode_label; +pub use engine::dispatcher::hook_handler_type_label; +pub use events::common::SubagentHookContext; +/// Hook event names as they appear in hooks JSON and config files. +pub const HOOK_EVENT_NAMES: [&str; 12] = [ + "PreToolUse", + "PermissionRequest", + "PostToolUse", + "PreCompact", + "PostCompact", + "SessionStart", + "SessionEnd", + "UserPromptSubmit", + "SubagentStart", + "SubagentStop", + "Stop", + "Interrupt", +]; + +/// Hook event names whose matcher fields are meaningful during dispatch. +/// +/// Other events can appear in hooks JSON, but Codex ignores their matcher +/// fields because those events do not dispatch against a tool, compaction +/// trigger, session-start source, or session-end reason. +pub const HOOK_EVENT_NAMES_WITH_MATCHERS: [&str; 9] = [ + "PreToolUse", + "PermissionRequest", + "PostToolUse", + "PreCompact", + "PostCompact", + "SessionStart", + "SessionEnd", + "SubagentStart", + "SubagentStop", +]; + +pub use events::compact::PostCompactRequest; +pub use events::compact::PreCompactOutcome; +pub use events::compact::PreCompactRequest; +pub use events::compact::StatelessHookOutcome; +pub use events::interrupt::InterruptOutcome; +pub use events::interrupt::InterruptRequest; +pub use events::permission_request::PermissionRequestDecision; +pub use events::permission_request::PermissionRequestOutcome; +pub use events::permission_request::PermissionRequestRequest; +pub use events::post_tool_use::PostToolUseOutcome; +pub use events::post_tool_use::PostToolUseRequest; +pub use events::pre_tool_use::PreToolUseOutcome; +pub use events::pre_tool_use::PreToolUseRequest; +pub use events::session_end::SessionEndOutcome; +pub use events::session_end::SessionEndRequest; +pub use events::session_start::SessionStartOutcome; +pub use events::session_start::SessionStartRequest; +pub use events::session_start::SessionStartSource; +pub use events::session_start::StartHookTarget; +pub use events::stop::StopHookTarget; +pub use events::stop::StopOutcome; +pub use events::stop::StopRequest; +pub use events::user_prompt_submit::UserPromptSubmitOutcome; +pub use events::user_prompt_submit::UserPromptSubmitRequest; +pub use legacy_notify::legacy_notify_json; +pub use mcp::HookMcpCall; +pub use mcp::HookMcpExecutor; +pub use registry::HookListOutcome; +pub use registry::Hooks; +pub use registry::HooksConfig; +pub use registry::list_hooks; +pub use schema::write_schema_fixtures; +pub use types::Hook; +pub use types::HookEvent; +pub use types::HookEventAfterAgent; +pub use types::HookPayload; +pub use types::HookResponse; +pub use types::HookResult; + +/// Returns the hook event label used in persisted hook-state keys. +pub fn hook_event_key_label(event_name: HookEventName) -> &'static str { + match event_name { + HookEventName::PreToolUse => "pre_tool_use", + HookEventName::PermissionRequest => "permission_request", + HookEventName::PostToolUse => "post_tool_use", + HookEventName::PreCompact => "pre_compact", + HookEventName::PostCompact => "post_compact", + HookEventName::SessionStart => "session_start", + HookEventName::SessionEnd => "session_end", + HookEventName::UserPromptSubmit => "user_prompt_submit", + HookEventName::SubagentStart => "subagent_start", + HookEventName::SubagentStop => "subagent_stop", + HookEventName::Stop => "stop", + HookEventName::Interrupt => "interrupt", + } +} + +/// Builds the persisted config-state key for one discovered hook handler. +pub fn hook_key( + key_source: &str, + event_name: HookEventName, + group_index: usize, + handler_index: usize, +) -> String { + format!( + "{key_source}:{}:{group_index}:{handler_index}", + hook_event_key_label(event_name) + ) +} diff --git a/codex-rs/hooks/src/mcp.rs b/codex-rs/hooks/src/mcp.rs new file mode 100644 index 0000000000000000000000000000000000000000..49d907df80e0241d17cf2e1b60bfe18f138ec8c5 --- /dev/null +++ b/codex-rs/hooks/src/mcp.rs @@ -0,0 +1,24 @@ +use std::time::Duration; + +use futures::future::BoxFuture; +use serde_json::Map; +use serde_json::Value; + +/// One MCP tool call requested by a configured hook handler. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct HookMcpCall { + pub server: String, + pub tool: String, + pub environment_id: Option, + pub metadata: Option>, + pub input: Map, + pub timeout: Duration, +} + +/// Executes already-connected MCP tools on behalf of hooks without coupling this crate to core. +/// +/// Implementations own server readiness, policy enforcement, timeout handling, and elicitation. +pub trait HookMcpExecutor: Send + Sync { + /// Returns text that is interpreted using ordinary command-hook output semantics. + fn execute(&self, call: HookMcpCall) -> BoxFuture<'_, anyhow::Result>; +} diff --git a/codex-rs/hooks/src/output_spill.rs b/codex-rs/hooks/src/output_spill.rs new file mode 100644 index 0000000000000000000000000000000000000000..b60ac973995a15b39c8499584acf4448b7eaffab --- /dev/null +++ b/codex-rs/hooks/src/output_spill.rs @@ -0,0 +1,135 @@ +use codex_protocol::ThreadId; +use codex_protocol::items::HookPromptFragment; +use codex_utils_absolute_path::AbsolutePathBuf; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::approx_token_count; +use codex_utils_output_truncation::formatted_truncate_text; +use tokio::fs; +use tracing::warn; +use uuid::Uuid; + +const HOOK_OUTPUTS_DIR: &str = "hook_outputs"; +pub(crate) const DEFAULT_HOOK_OUTPUT_TOKEN_LIMIT: usize = 2_500; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct AdditionalContextLimit { + token_limit: usize, +} + +impl AdditionalContextLimit { + pub(crate) fn from_config(value: Option) -> Self { + Self { + token_limit: value.unwrap_or(DEFAULT_HOOK_OUTPUT_TOKEN_LIMIT), + } + } +} + +impl Default for AdditionalContextLimit { + fn default() -> Self { + Self::from_config(/*value*/ None) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct AdditionalContext { + pub text: String, + pub limit: AdditionalContextLimit, +} + +#[derive(Clone)] +pub(crate) struct HookOutputSpiller { + output_dir: AbsolutePathBuf, +} + +impl HookOutputSpiller { + pub(crate) fn new(thread_id: ThreadId) -> Self { + Self { + output_dir: AbsolutePathBuf::resolve_path_against_base(std::env::temp_dir(), "/") + .join(HOOK_OUTPUTS_DIR) + .join(thread_id.to_string()), + } + } + + /// Keeps hook text within the model-visible hook-output budget. + /// + /// Oversized text is written in full under the OS temp directory at + /// `/hook_outputs//` + /// and replaced with the same head/tail preview style used for other truncated + /// output, plus a path back to the preserved full text. + pub(crate) async fn maybe_spill_text(&self, text: String) -> String { + self.maybe_spill_text_with_limit(text, AdditionalContextLimit::default()) + .await + } + + async fn maybe_spill_text_with_limit( + &self, + text: String, + limit: AdditionalContextLimit, + ) -> String { + let token_limit = limit.token_limit; + if token_limit == 0 || approx_token_count(&text) <= token_limit { + return text; + } + + let path = self.output_dir.join(format!("{}.txt", Uuid::new_v4())); + if let Some(parent) = path.parent() + && let Err(err) = fs::create_dir_all(parent.as_ref()).await + { + warn!( + "failed to create hook output directory {}: {err}", + parent.display() + ); + return formatted_truncate_text(&text, TruncationPolicy::Tokens(token_limit)); + } + + if let Err(err) = fs::write(path.as_ref(), &text).await { + warn!("failed to write hook output {}: {err}", path.display()); + return formatted_truncate_text(&text, TruncationPolicy::Tokens(token_limit)); + } + + spilled_hook_output_preview(&text, &path, token_limit) + } + + pub(crate) async fn maybe_spill_additional_contexts( + &self, + contexts: Vec, + ) -> Vec { + let mut spilled = Vec::with_capacity(contexts.len()); + for context in contexts { + spilled.push( + self.maybe_spill_text_with_limit(context.text, context.limit) + .await, + ); + } + spilled + } + + pub(crate) async fn maybe_spill_prompt_fragments( + &self, + fragments: Vec, + ) -> Vec { + let mut spilled = Vec::with_capacity(fragments.len()); + for fragment in fragments { + spilled.push(HookPromptFragment { + text: self.maybe_spill_text(fragment.text).await, + hook_run_id: fragment.hook_run_id, + }); + } + spilled + } +} + +/// Builds the model-visible replacement for a spilled hook output. +/// +/// The path footer is budgeted before truncation so adding the recovery path +/// does not consume the configured preview budget. +fn spilled_hook_output_preview(text: &str, path: &AbsolutePathBuf, token_limit: usize) -> String { + let footer = format!("\n\nFull hook output saved to: {}", path.display()); + let preview_policy = + TruncationPolicy::Tokens(token_limit.saturating_sub(approx_token_count(&footer))); + format!("{}{footer}", formatted_truncate_text(text, preview_policy)) +} + +#[cfg(test)] +#[path = "output_spill_tests.rs"] +mod tests; diff --git a/codex-rs/hooks/src/output_spill_tests.rs b/codex-rs/hooks/src/output_spill_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..6b72c956ce16793d12cea015ee3d27997be28d94 --- /dev/null +++ b/codex-rs/hooks/src/output_spill_tests.rs @@ -0,0 +1,83 @@ +use super::*; +use anyhow::Context; +use anyhow::Result; +use tempfile::tempdir; + +#[test] +fn hook_output_spiller_is_scoped_to_its_thread() { + let thread_id = ThreadId::new(); + let spiller = HookOutputSpiller::new(thread_id); + + assert_eq!( + spiller.output_dir, + AbsolutePathBuf::resolve_path_against_base(std::env::temp_dir(), "/") + .join(HOOK_OUTPUTS_DIR) + .join(thread_id.to_string()) + ); +} + +#[tokio::test] +async fn small_hook_output_remains_inline() -> Result<()> { + let dir = tempdir()?; + let output_dir = AbsolutePathBuf::from_absolute_path(dir.path())?.join(HOOK_OUTPUTS_DIR); + let spiller = HookOutputSpiller { + output_dir: output_dir.clone(), + }; + + let output = spiller.maybe_spill_text("short".to_string()).await; + + assert_eq!(output, "short"); + assert!(!output_dir.exists()); + Ok(()) +} + +#[tokio::test] +async fn large_hook_output_spills_to_file() -> Result<()> { + let dir = tempdir()?; + let text = "hook output ".repeat(1_000); + let output_dir = AbsolutePathBuf::from_absolute_path(dir.path())?.join(HOOK_OUTPUTS_DIR); + let spiller = HookOutputSpiller { output_dir }; + + let output = spiller.maybe_spill_text(text.clone()).await; + + assert!(output.contains("tokens truncated")); + let path = output + .lines() + .find_map(|line| line.strip_prefix("Full hook output saved to: ")) + .context("spill path")?; + assert_eq!(fs::read_to_string(path).await?, text); + Ok(()) +} + +#[tokio::test] +async fn additional_contexts_apply_limits_individually() -> Result<()> { + let dir = tempdir()?; + let limited_text = "limited hook output ".repeat(1_000); + let unlimited_text = "unlimited hook output ".repeat(5_000); + assert!(approx_token_count(&unlimited_text) > 10_000); + let output_dir = AbsolutePathBuf::from_absolute_path(dir.path())?.join(HOOK_OUTPUTS_DIR); + let spiller = HookOutputSpiller { output_dir }; + let output = spiller + .maybe_spill_additional_contexts(vec![ + AdditionalContext { + text: limited_text.clone(), + limit: AdditionalContextLimit::from_config(Some(1)), + }, + AdditionalContext { + text: unlimited_text.clone(), + limit: AdditionalContextLimit::from_config(Some(0)), + }, + AdditionalContext { + text: unlimited_text.clone(), + limit: AdditionalContextLimit::from_config(Some(usize::MAX)), + }, + ]) + .await; + let [limited_output, zero_limit_output, high_limit_output] = output.as_slice() else { + panic!("expected one output for each additional context"); + }; + assert!(limited_output.contains("Full hook output saved to:")); + assert_eq!(zero_limit_output, &unlimited_text); + assert_eq!(high_limit_output, &unlimited_text); + Ok(()) +} diff --git a/codex-rs/hooks/src/registry.rs b/codex-rs/hooks/src/registry.rs new file mode 100644 index 0000000000000000000000000000000000000000..5a033f9c64dacaf031719e60c3c42d52b2d9eb07 --- /dev/null +++ b/codex-rs/hooks/src/registry.rs @@ -0,0 +1,336 @@ +use crate::engine::ClaudeHooksEngine; +use crate::engine::CommandShell; +use crate::engine::HookListEntry; +use crate::engine::command_runner::CommandHookRuntime; +use crate::events::compact::PostCompactRequest; +use crate::events::compact::PreCompactOutcome; +use crate::events::compact::PreCompactRequest; +use crate::events::compact::StatelessHookOutcome; +use crate::events::interrupt::InterruptOutcome; +use crate::events::interrupt::InterruptRequest; +use crate::events::permission_request::PermissionRequestOutcome; +use crate::events::permission_request::PermissionRequestRequest; +use crate::events::post_tool_use::PostToolUseOutcome; +use crate::events::post_tool_use::PostToolUseRequest; +use crate::events::pre_tool_use::PreToolUseOutcome; +use crate::events::pre_tool_use::PreToolUseRequest; +use crate::events::session_end::SessionEndOutcome; +use crate::events::session_end::SessionEndRequest; +use crate::events::session_start::SessionStartOutcome; +use crate::events::session_start::SessionStartRequest; +use crate::events::stop::StopOutcome; +use crate::events::stop::StopRequest; +use crate::events::user_prompt_submit::UserPromptSubmitOutcome; +use crate::events::user_prompt_submit::UserPromptSubmitRequest; +use crate::mcp::HookMcpExecutor; +use crate::types::Hook; +use crate::types::HookEvent; +use crate::types::HookPayload; +use crate::types::HookResponse; +use async_channel::Receiver; +use codex_config::ConfigLayerStack; +use codex_plugin::ExecutorPluginHookSource; +use codex_plugin::PluginHookSource; +use codex_protocol::ThreadId; +use codex_protocol::shell_environment::scrub_non_inheritable_env_vars; +use std::ffi::OsString; +use std::sync::Arc; +use std::time::Duration; +use tokio::process::Command; + +#[derive(Default, Clone)] +pub struct HooksConfig { + pub legacy_notify_argv: Option>, + pub feature_enabled: bool, + pub bypass_hook_trust: bool, + pub config_layer_stack: Option, + pub plugin_hook_sources: Vec, + pub plugin_hook_load_warnings: Vec, + pub shell_program: Option, + pub shell_args: Vec, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct HookListOutcome { + pub hooks: Vec, + pub warnings: Vec, +} + +#[derive(Clone)] +pub struct Hooks { + // TODO: Once legacy `notify` is removed, capture this snapshot in `CommandHookRuntime::new` + // and remove the environment plumbing from `Hooks` and `from_config`. + environment: Arc>, + after_agent: Vec, + engine: ClaudeHooksEngine, + plugin_hook_sources: Vec, + plugin_hook_load_warnings: Vec, +} + +impl Hooks { + /// Bind this session's hook runtime and output files to its thread, rejecting unloadable + /// required managed hooks. + pub fn new( + config: HooksConfig, + thread_id: ThreadId, + mcp_executor: Arc, + ) -> anyhow::Result<(Self, Receiver)> { + let (result_sender, result_receiver) = async_channel::unbounded(); + let environment = Arc::new(std::env::vars_os().collect()); + let hooks = Self::from_config(config, mcp_executor, Arc::clone(&environment), |shell| { + CommandHookRuntime::new(shell, environment, thread_id, result_sender) + }); + let required_load_errors = hooks.engine.required_load_errors(); + if !required_load_errors.is_empty() { + anyhow::bail!( + "failed to load required managed hooks: {}", + required_load_errors.join("; ") + ); + } + Ok((hooks, result_receiver)) + } + + /// Preserve in-flight background hooks while applying a refreshed configuration. + pub fn reconfigured(&self, config: HooksConfig) -> Self { + Self::from_config( + config, + Arc::clone(&self.engine.mcp_executor), + Arc::clone(&self.environment), + |shell| self.engine.command_runtime.reconfigured(shell), + ) + } + + pub fn matches_plugin_hooks<'a>( + &self, + sources: impl IntoIterator, + warnings: impl IntoIterator, + ) -> bool { + self.plugin_hook_sources.iter().eq(sources) + && self.plugin_hook_load_warnings.iter().eq(warnings) + } + + pub fn with_executor_hooks(&self, executor_hooks: Vec) -> Self { + let mut hooks = self.clone(); + hooks.engine.set_executor_hooks(executor_hooks); + hooks + } + + fn from_config( + config: HooksConfig, + mcp_executor: Arc, + environment: Arc>, + build_runtime: impl FnOnce(CommandShell) -> CommandHookRuntime, + ) -> Self { + let after_agent = config + .legacy_notify_argv + .filter(|argv| !argv.is_empty() && !argv[0].is_empty()) + .map(|argv| crate::legacy_notify::notify_hook(argv, Arc::clone(&environment))) + .into_iter() + .collect(); + let command_runtime = build_runtime(CommandShell { + program: config.shell_program.unwrap_or_default(), + args: config.shell_args, + }); + let engine = ClaudeHooksEngine::new( + config.feature_enabled, + config.bypass_hook_trust, + config.config_layer_stack.as_ref(), + config.plugin_hook_sources.clone(), + config.plugin_hook_load_warnings.clone(), + command_runtime, + mcp_executor, + ); + Self { + environment, + after_agent, + engine, + plugin_hook_sources: config.plugin_hook_sources, + plugin_hook_load_warnings: config.plugin_hook_load_warnings, + } + } + + /// Abort and join outstanding async hooks during session shutdown. + pub async fn shutdown(&self) { + self.engine.command_runtime.shutdown().await; + } + + pub fn startup_warnings(&self) -> &[String] { + self.engine.warnings() + } + + fn hooks_for_event(&self, hook_event: &HookEvent) -> &[Hook] { + match hook_event { + HookEvent::AfterAgent { .. } => &self.after_agent, + } + } + + pub async fn dispatch(&self, hook_payload: HookPayload) -> Vec { + let hooks = self.hooks_for_event(&hook_payload.hook_event); + let mut outcomes = Vec::with_capacity(hooks.len()); + for hook in hooks { + let outcome = hook.execute(&hook_payload).await; + let should_abort_operation = outcome.result.should_abort_operation(); + outcomes.push(outcome); + if should_abort_operation { + break; + } + } + + outcomes + } + + pub fn preview_session_start( + &self, + request: &SessionStartRequest, + ) -> Vec { + self.engine.preview_session_start(request) + } + + pub fn preview_pre_tool_use( + &self, + request: &PreToolUseRequest, + ) -> Vec { + self.engine.preview_pre_tool_use(request) + } + + pub fn preview_permission_request( + &self, + request: &PermissionRequestRequest, + ) -> Vec { + self.engine.preview_permission_request(request) + } + + /// Maximum configured timeout among PermissionRequest hooks. + /// + /// Matching handlers run concurrently, so their aggregate timeout is bounded by this maximum. + pub fn max_permission_request_timeout(&self) -> Duration { + self.engine.max_permission_request_timeout() + } + + pub fn preview_post_tool_use( + &self, + request: &PostToolUseRequest, + ) -> Vec { + self.engine.preview_post_tool_use(request) + } + + pub async fn run_session_start( + &self, + request: SessionStartRequest, + turn_id: Option, + ) -> SessionStartOutcome { + self.engine.run_session_start(request, turn_id).await + } + + pub async fn run_pre_tool_use(&self, request: PreToolUseRequest) -> PreToolUseOutcome { + self.engine.run_pre_tool_use(request).await + } + + pub async fn run_permission_request( + &self, + request: PermissionRequestRequest, + ) -> PermissionRequestOutcome { + self.engine.run_permission_request(request).await + } + + pub async fn run_post_tool_use(&self, request: PostToolUseRequest) -> PostToolUseOutcome { + self.engine.run_post_tool_use(request).await + } + + pub fn preview_pre_compact( + &self, + request: &PreCompactRequest, + ) -> Vec { + self.engine.preview_pre_compact(request) + } + + pub async fn run_pre_compact(&self, request: PreCompactRequest) -> PreCompactOutcome { + self.engine.run_pre_compact(request).await + } + + pub fn preview_post_compact( + &self, + request: &PostCompactRequest, + ) -> Vec { + self.engine.preview_post_compact(request) + } + + pub async fn run_post_compact(&self, request: PostCompactRequest) -> StatelessHookOutcome { + self.engine.run_post_compact(request).await + } + + pub fn preview_user_prompt_submit( + &self, + request: &UserPromptSubmitRequest, + ) -> Vec { + self.engine.preview_user_prompt_submit(request) + } + + pub async fn run_user_prompt_submit( + &self, + request: UserPromptSubmitRequest, + ) -> UserPromptSubmitOutcome { + self.engine.run_user_prompt_submit(request).await + } + + pub fn preview_stop( + &self, + request: &StopRequest, + ) -> Vec { + self.engine.preview_stop(request) + } + + pub async fn run_stop(&self, request: StopRequest) -> StopOutcome { + self.engine.run_stop(request).await + } + + pub fn preview_session_end(&self) -> Vec { + self.engine.preview_session_end() + } + + pub async fn run_session_end(&self, request: SessionEndRequest) -> SessionEndOutcome { + self.engine.run_session_end(request).await + } + + pub fn preview_interrupt(&self) -> Vec { + self.engine.preview_interrupt() + } + + pub async fn run_interrupt(&self, request: InterruptRequest) -> InterruptOutcome { + self.engine.run_interrupt(request).await + } +} + +pub fn list_hooks(config: HooksConfig) -> HookListOutcome { + if !config.feature_enabled { + return HookListOutcome::default(); + } + + let discovered = crate::engine::discovery::discover_handlers( + config.config_layer_stack.as_ref(), + config.plugin_hook_sources, + config.plugin_hook_load_warnings, + config.bypass_hook_trust, + ); + HookListOutcome { + hooks: discovered.hook_entries, + warnings: discovered.warnings, + } +} + +// TODO: Remove this legacy-notify-only command builder when `notify` support is removed. +pub(crate) fn command_from_argv( + argv: &[String], + environment: impl IntoIterator, +) -> Option { + let (program, args) = argv.split_first()?; + if program.is_empty() { + return None; + } + let mut command = Command::new(program); + command.args(args); + command.env_clear(); + command.envs(environment); + scrub_non_inheritable_env_vars(command.as_std_mut()); + Some(command) +} diff --git a/codex-rs/hooks/src/schema.rs b/codex-rs/hooks/src/schema.rs new file mode 100644 index 0000000000000000000000000000000000000000..bb50807af5fc242f4fb65a20357920a9fc315a63 --- /dev/null +++ b/codex-rs/hooks/src/schema.rs @@ -0,0 +1,1254 @@ +use schemars::JsonSchema; +use schemars::r#gen::SchemaGenerator; +use schemars::r#gen::SchemaSettings; +use schemars::schema::InstanceType; +use schemars::schema::RootSchema; +use schemars::schema::Schema; +use schemars::schema::SchemaObject; +use serde::Deserialize; +use serde::Serialize; +use serde_json::Map; +use serde_json::Value; +use std::path::Path; +use std::path::PathBuf; + +use crate::events::common::SubagentHookContext; + +const GENERATED_DIR: &str = "generated"; +const POST_TOOL_USE_INPUT_FIXTURE: &str = "post-tool-use.command.input.schema.json"; +const POST_TOOL_USE_OUTPUT_FIXTURE: &str = "post-tool-use.command.output.schema.json"; +const PERMISSION_REQUEST_INPUT_FIXTURE: &str = "permission-request.command.input.schema.json"; +const PERMISSION_REQUEST_OUTPUT_FIXTURE: &str = "permission-request.command.output.schema.json"; +const POST_COMPACT_INPUT_FIXTURE: &str = "post-compact.command.input.schema.json"; +const POST_COMPACT_OUTPUT_FIXTURE: &str = "post-compact.command.output.schema.json"; +const PRE_TOOL_USE_INPUT_FIXTURE: &str = "pre-tool-use.command.input.schema.json"; +const PRE_TOOL_USE_OUTPUT_FIXTURE: &str = "pre-tool-use.command.output.schema.json"; +const PRE_COMPACT_INPUT_FIXTURE: &str = "pre-compact.command.input.schema.json"; +const PRE_COMPACT_OUTPUT_FIXTURE: &str = "pre-compact.command.output.schema.json"; +const SESSION_START_INPUT_FIXTURE: &str = "session-start.command.input.schema.json"; +const SESSION_START_OUTPUT_FIXTURE: &str = "session-start.command.output.schema.json"; +const SESSION_END_INPUT_FIXTURE: &str = "session-end.command.input.schema.json"; +const USER_PROMPT_SUBMIT_INPUT_FIXTURE: &str = "user-prompt-submit.command.input.schema.json"; +const USER_PROMPT_SUBMIT_OUTPUT_FIXTURE: &str = "user-prompt-submit.command.output.schema.json"; +const SUBAGENT_START_INPUT_FIXTURE: &str = "subagent-start.command.input.schema.json"; +const SUBAGENT_START_OUTPUT_FIXTURE: &str = "subagent-start.command.output.schema.json"; +const SUBAGENT_STOP_INPUT_FIXTURE: &str = "subagent-stop.command.input.schema.json"; +const SUBAGENT_STOP_OUTPUT_FIXTURE: &str = "subagent-stop.command.output.schema.json"; +const STOP_INPUT_FIXTURE: &str = "stop.command.input.schema.json"; +const STOP_OUTPUT_FIXTURE: &str = "stop.command.output.schema.json"; +const INTERRUPT_INPUT_FIXTURE: &str = "interrupt.command.input.schema.json"; +const INTERRUPT_OUTPUT_FIXTURE: &str = "interrupt.command.output.schema.json"; + +#[derive(Debug, Clone, Serialize)] +#[serde(transparent)] +pub(crate) struct NullableString(Option); + +impl NullableString { + pub(crate) fn from_path(path: Option) -> Self { + Self(path.map(|path| path.display().to_string())) + } + + pub(crate) fn from_string(value: Option) -> Self { + Self(value) + } +} + +impl JsonSchema for NullableString { + fn schema_name() -> String { + "NullableString".to_string() + } + + fn json_schema(_gen: &mut SchemaGenerator) -> Schema { + Schema::Object(SchemaObject { + instance_type: Some(vec![InstanceType::String, InstanceType::Null].into()), + ..Default::default() + }) + } +} + +#[derive(Debug, Clone, Default)] +pub(crate) struct SubagentCommandInputFields { + pub agent_id: Option, + pub agent_type: Option, +} + +impl From> for SubagentCommandInputFields { + fn from(value: Option<&SubagentHookContext>) -> Self { + match value { + Some(context) => Self { + agent_id: Some(context.agent_id.clone()), + agent_type: Some(context.agent_type.clone()), + }, + None => Self::default(), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +pub(crate) struct HookUniversalOutputWire { + #[serde(default = "default_continue")] + pub r#continue: bool, + #[serde(default)] + pub stop_reason: Option, + #[serde(default)] + pub suppress_output: bool, + #[serde(default)] + pub system_message: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)] +pub(crate) enum HookEventNameWire { + #[serde(rename = "PreToolUse")] + PreToolUse, + #[serde(rename = "PermissionRequest")] + PermissionRequest, + #[serde(rename = "PostToolUse")] + PostToolUse, + #[serde(rename = "PreCompact")] + PreCompact, + #[serde(rename = "PostCompact")] + PostCompact, + #[serde(rename = "SessionStart")] + SessionStart, + #[serde(rename = "UserPromptSubmit")] + UserPromptSubmit, + #[serde(rename = "SubagentStart")] + SubagentStart, + #[serde(rename = "SubagentStop")] + SubagentStop, + #[serde(rename = "Stop")] + Stop, + #[serde(rename = "Interrupt")] + Interrupt, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +#[schemars(rename = "pre-tool-use.command.output")] +pub(crate) struct PreToolUseCommandOutputWire { + #[serde(flatten)] + pub universal: HookUniversalOutputWire, + #[serde(default)] + pub decision: Option, + #[serde(default)] + pub reason: Option, + #[serde(default)] + pub hook_specific_output: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +#[schemars(rename = "post-tool-use.command.output")] +pub(crate) struct PostToolUseCommandOutputWire { + #[serde(flatten)] + pub universal: HookUniversalOutputWire, + #[serde(default)] + pub decision: Option, + #[serde(default)] + pub reason: Option, + #[serde(default)] + pub hook_specific_output: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +#[schemars(rename = "permission-request.command.output")] +pub(crate) struct PermissionRequestCommandOutputWire { + #[serde(flatten)] + pub universal: HookUniversalOutputWire, + #[serde(default)] + pub hook_specific_output: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +#[schemars(rename = "pre-compact.command.output")] +pub(crate) struct PreCompactCommandOutputWire { + #[serde(flatten)] + pub universal: HookUniversalOutputWire, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +#[schemars(rename = "post-compact.command.output")] +pub(crate) struct PostCompactCommandOutputWire { + #[serde(flatten)] + pub universal: HookUniversalOutputWire, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +pub(crate) struct PermissionRequestHookSpecificOutputWire { + #[schemars(schema_with = "permission_request_hook_event_name_schema")] + pub hook_event_name: HookEventNameWire, + #[serde(default)] + pub decision: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +pub(crate) struct PermissionRequestDecisionWire { + pub behavior: PermissionRequestBehaviorWire, + /// Reserved for a future input-rewrite capability. + /// + /// PermissionRequest hooks currently fail closed if this field is present. + #[serde(default)] + pub updated_input: Option, + /// Reserved for a future permission-rewrite capability. + /// + /// PermissionRequest hooks currently fail closed if this field is present. + #[serde(default)] + pub updated_permissions: Option, + #[serde(default)] + pub message: Option, + /// Reserved for future short-circuiting semantics. + /// + /// PermissionRequest hooks currently fail closed if this field is `true`. + #[serde(default)] + pub interrupt: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)] +pub(crate) enum PermissionRequestBehaviorWire { + #[serde(rename = "allow")] + Allow, + #[serde(rename = "deny")] + Deny, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +pub(crate) struct PostToolUseHookSpecificOutputWire { + #[schemars(schema_with = "post_tool_use_hook_event_name_schema")] + pub hook_event_name: HookEventNameWire, + #[serde(default)] + pub additional_context: Option, + #[serde(default)] + #[serde(rename = "updatedMCPToolOutput")] + pub updated_mcp_tool_output: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +pub(crate) struct PreToolUseHookSpecificOutputWire { + #[schemars(schema_with = "pre_tool_use_hook_event_name_schema")] + pub hook_event_name: HookEventNameWire, + #[serde(default)] + pub permission_decision: Option, + #[serde(default)] + pub permission_decision_reason: Option, + #[serde(default)] + pub updated_input: Option, + #[serde(default)] + pub additional_context: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)] +pub(crate) enum PreToolUsePermissionDecisionWire { + #[serde(rename = "allow")] + Allow, + #[serde(rename = "deny")] + Deny, + #[serde(rename = "ask")] + Ask, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)] +pub(crate) enum PreToolUseDecisionWire { + #[serde(rename = "approve")] + Approve, + #[serde(rename = "block")] + Block, +} + +#[derive(Debug, Clone, Serialize, JsonSchema)] +#[serde(deny_unknown_fields)] +#[schemars(rename = "pre-tool-use.command.input")] +pub(crate) struct PreToolUseCommandInput { + pub session_id: String, + /// Codex extension: expose the active turn id to internal turn-scoped hooks. + pub turn_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_type: Option, + pub transcript_path: NullableString, + pub cwd: String, + #[schemars(schema_with = "pre_tool_use_hook_event_name_schema")] + pub hook_event_name: String, + pub model: String, + #[schemars(schema_with = "permission_mode_schema")] + pub permission_mode: String, + pub tool_name: String, + pub tool_input: Value, + pub tool_use_id: String, +} + +#[derive(Debug, Clone, Serialize, JsonSchema)] +#[serde(deny_unknown_fields)] +#[schemars(rename = "permission-request.command.input")] +pub(crate) struct PermissionRequestCommandInput { + pub session_id: String, + /// Codex extension: expose the active turn id to internal turn-scoped hooks. + pub turn_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_type: Option, + pub transcript_path: NullableString, + pub cwd: String, + #[schemars(schema_with = "permission_request_hook_event_name_schema")] + pub hook_event_name: String, + pub model: String, + #[schemars(schema_with = "permission_mode_schema")] + pub permission_mode: String, + pub tool_name: String, + pub tool_input: Value, +} + +#[derive(Debug, Clone, Serialize, JsonSchema)] +#[serde(deny_unknown_fields)] +#[schemars(rename = "post-tool-use.command.input")] +pub(crate) struct PostToolUseCommandInput { + pub session_id: String, + /// Codex extension: expose the active turn id to internal turn-scoped hooks. + pub turn_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_type: Option, + pub transcript_path: NullableString, + pub cwd: String, + #[schemars(schema_with = "post_tool_use_hook_event_name_schema")] + pub hook_event_name: String, + pub model: String, + #[schemars(schema_with = "permission_mode_schema")] + pub permission_mode: String, + pub tool_name: String, + pub tool_input: Value, + pub tool_response: Value, + pub tool_use_id: String, +} + +#[derive(Debug, Clone, Serialize, JsonSchema)] +#[serde(deny_unknown_fields)] +#[schemars(rename = "pre-compact.command.input")] +pub(crate) struct PreCompactCommandInput { + pub session_id: String, + /// Codex extension: expose the active turn id to internal turn-scoped hooks. + pub turn_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_type: Option, + pub transcript_path: NullableString, + pub cwd: String, + #[schemars(schema_with = "pre_compact_hook_event_name_schema")] + pub hook_event_name: String, + pub model: String, + #[schemars(schema_with = "compaction_trigger_schema")] + pub trigger: String, +} + +#[derive(Debug, Clone, Serialize, JsonSchema)] +#[serde(deny_unknown_fields)] +#[schemars(rename = "post-compact.command.input")] +pub(crate) struct PostCompactCommandInput { + pub session_id: String, + /// Codex extension: expose the active turn id to internal turn-scoped hooks. + pub turn_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_type: Option, + pub transcript_path: NullableString, + pub cwd: String, + #[schemars(schema_with = "post_compact_hook_event_name_schema")] + pub hook_event_name: String, + pub model: String, + #[schemars(schema_with = "compaction_trigger_schema")] + pub trigger: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +#[schemars(rename = "session-start.command.output")] +pub(crate) struct SessionStartCommandOutputWire { + #[serde(flatten)] + pub universal: HookUniversalOutputWire, + #[serde(default)] + pub hook_specific_output: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +pub(crate) struct SessionStartHookSpecificOutputWire { + #[schemars(schema_with = "session_start_hook_event_name_schema")] + pub hook_event_name: HookEventNameWire, + #[serde(default)] + pub additional_context: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +#[schemars(rename = "subagent-start.command.output")] +pub(crate) struct SubagentStartCommandOutputWire { + #[serde(flatten)] + pub universal: HookUniversalOutputWire, + #[serde(default)] + pub hook_specific_output: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +pub(crate) struct SubagentStartHookSpecificOutputWire { + #[schemars(schema_with = "subagent_start_hook_event_name_schema")] + pub hook_event_name: HookEventNameWire, + #[serde(default)] + pub additional_context: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +#[schemars(rename = "user-prompt-submit.command.output")] +pub(crate) struct UserPromptSubmitCommandOutputWire { + #[serde(flatten)] + pub universal: HookUniversalOutputWire, + #[serde(default)] + pub decision: Option, + #[serde(default)] + pub reason: Option, + #[serde(default)] + pub hook_specific_output: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +pub(crate) struct UserPromptSubmitHookSpecificOutputWire { + #[schemars(schema_with = "user_prompt_submit_hook_event_name_schema")] + pub hook_event_name: HookEventNameWire, + #[serde(default)] + pub additional_context: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +#[schemars(rename = "stop.command.output")] +pub(crate) struct StopCommandOutputWire { + #[serde(flatten)] + pub universal: HookUniversalOutputWire, + #[serde(default)] + pub decision: Option, + /// Claude requires `reason` when `decision` is `block`; we enforce that + /// semantic rule during output parsing rather than in the JSON schema. + #[serde(default)] + pub reason: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +#[schemars(rename = "subagent-stop.command.output")] +pub(crate) struct SubagentStopCommandOutputWire { + #[serde(flatten)] + pub universal: HookUniversalOutputWire, + #[serde(default)] + pub decision: Option, + /// Claude requires `reason` when `decision` is `block`; we enforce that + /// semantic rule during output parsing rather than in the JSON schema. + #[serde(default)] + pub reason: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "camelCase")] +#[serde(deny_unknown_fields)] +#[schemars(rename = "interrupt.command.output")] +pub(crate) struct InterruptCommandOutputWire { + #[serde(default)] + pub system_message: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)] +pub(crate) enum BlockDecisionWire { + #[serde(rename = "block")] + Block, +} + +#[derive(Debug, Clone, Serialize, JsonSchema)] +#[serde(deny_unknown_fields)] +#[schemars(rename = "session-start.command.input")] +pub(crate) struct SessionStartCommandInput { + pub session_id: String, + pub transcript_path: NullableString, + pub cwd: String, + #[schemars(schema_with = "session_start_hook_event_name_schema")] + pub hook_event_name: String, + pub model: String, + #[schemars(schema_with = "permission_mode_schema")] + pub permission_mode: String, + #[schemars(schema_with = "session_start_source_schema")] + pub source: String, +} + +#[derive(Debug, Clone, Serialize, JsonSchema)] +#[serde(deny_unknown_fields)] +#[schemars(rename = "session-end.command.input")] +pub(crate) struct SessionEndCommandInput { + pub session_id: String, + pub transcript_path: NullableString, + pub cwd: String, + #[schemars(schema_with = "session_end_hook_event_name_schema")] + pub hook_event_name: String, + #[schemars(schema_with = "session_end_reason_schema")] + pub reason: String, +} + +impl SessionStartCommandInput { + pub(crate) fn new( + session_id: impl Into, + transcript_path: Option, + cwd: impl Into, + model: impl Into, + permission_mode: impl Into, + source: impl Into, + ) -> Self { + Self { + session_id: session_id.into(), + transcript_path: NullableString::from_path(transcript_path), + cwd: cwd.into(), + hook_event_name: "SessionStart".to_string(), + model: model.into(), + permission_mode: permission_mode.into(), + source: source.into(), + } + } +} + +#[derive(Debug, Clone, Serialize, JsonSchema)] +#[serde(deny_unknown_fields)] +#[schemars(rename = "subagent-start.command.input")] +pub(crate) struct SubagentStartCommandInput { + pub session_id: String, + /// Codex extension: expose the active turn id to internal turn-scoped hooks. + pub turn_id: String, + pub transcript_path: NullableString, + pub cwd: String, + #[schemars(schema_with = "subagent_start_hook_event_name_schema")] + pub hook_event_name: String, + pub model: String, + #[schemars(schema_with = "permission_mode_schema")] + pub permission_mode: String, + pub agent_id: String, + pub agent_type: String, +} + +#[derive(Debug, Clone, Serialize, JsonSchema)] +#[serde(deny_unknown_fields)] +#[schemars(rename = "user-prompt-submit.command.input")] +pub(crate) struct UserPromptSubmitCommandInput { + pub session_id: String, + /// Codex extension: expose the active turn id to internal turn-scoped hooks. + pub turn_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_type: Option, + pub transcript_path: NullableString, + pub cwd: String, + #[schemars(schema_with = "user_prompt_submit_hook_event_name_schema")] + pub hook_event_name: String, + pub model: String, + #[schemars(schema_with = "permission_mode_schema")] + pub permission_mode: String, + pub prompt: String, +} + +#[derive(Debug, Clone, Serialize, JsonSchema)] +#[serde(deny_unknown_fields)] +#[schemars(rename = "stop.command.input")] +pub(crate) struct StopCommandInput { + pub session_id: String, + /// Codex extension: expose the active turn id to internal turn-scoped hooks. + pub turn_id: String, + pub transcript_path: NullableString, + pub cwd: String, + #[schemars(schema_with = "stop_hook_event_name_schema")] + pub hook_event_name: String, + pub model: String, + #[schemars(schema_with = "permission_mode_schema")] + pub permission_mode: String, + pub stop_hook_active: bool, + pub last_assistant_message: NullableString, +} + +#[derive(Debug, Clone, Serialize, JsonSchema)] +#[serde(deny_unknown_fields)] +#[schemars(rename = "subagent-stop.command.input")] +pub(crate) struct SubagentStopCommandInput { + pub session_id: String, + /// Codex extension: expose the active turn id to internal turn-scoped hooks. + pub turn_id: String, + pub transcript_path: NullableString, + pub agent_transcript_path: NullableString, + pub cwd: String, + #[schemars(schema_with = "subagent_stop_hook_event_name_schema")] + pub hook_event_name: String, + pub model: String, + #[schemars(schema_with = "permission_mode_schema")] + pub permission_mode: String, + pub stop_hook_active: bool, + pub agent_id: String, + pub agent_type: String, + pub last_assistant_message: NullableString, +} + +#[derive(Debug, Clone, Serialize, JsonSchema)] +#[serde(deny_unknown_fields)] +#[schemars(rename = "interrupt.command.input")] +pub(crate) struct InterruptCommandInput { + pub session_id: String, + /// Codex extension: expose the active turn id to internal turn-scoped hooks. + pub turn_id: String, + pub transcript_path: NullableString, + pub cwd: String, + #[schemars(schema_with = "interrupt_hook_event_name_schema")] + pub hook_event_name: String, + pub model: String, + #[schemars(schema_with = "permission_mode_schema")] + pub permission_mode: String, +} + +pub fn write_schema_fixtures(schema_root: &Path) -> anyhow::Result<()> { + let generated_dir = schema_root.join(GENERATED_DIR); + ensure_empty_dir(&generated_dir)?; + + write_schema( + &generated_dir.join(POST_TOOL_USE_INPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(POST_TOOL_USE_OUTPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(PERMISSION_REQUEST_INPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(PERMISSION_REQUEST_OUTPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(POST_COMPACT_INPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(POST_COMPACT_OUTPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(PRE_COMPACT_INPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(PRE_COMPACT_OUTPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(PRE_TOOL_USE_INPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(PRE_TOOL_USE_OUTPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(SESSION_START_INPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(SESSION_START_OUTPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(SESSION_END_INPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(USER_PROMPT_SUBMIT_INPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(USER_PROMPT_SUBMIT_OUTPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(SUBAGENT_START_INPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(SUBAGENT_START_OUTPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(SUBAGENT_STOP_INPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(SUBAGENT_STOP_OUTPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(STOP_INPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(STOP_OUTPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(INTERRUPT_INPUT_FIXTURE), + schema_json::()?, + )?; + write_schema( + &generated_dir.join(INTERRUPT_OUTPUT_FIXTURE), + schema_json::()?, + )?; + + Ok(()) +} + +fn write_schema(path: &Path, json: Vec) -> anyhow::Result<()> { + std::fs::write(path, json)?; + Ok(()) +} + +fn ensure_empty_dir(dir: &Path) -> anyhow::Result<()> { + if dir.exists() { + std::fs::remove_dir_all(dir)?; + } + std::fs::create_dir_all(dir)?; + Ok(()) +} + +fn schema_json() -> anyhow::Result> +where + T: JsonSchema, +{ + let schema = schema_for_type::(); + let value = serde_json::to_value(schema)?; + let value = canonicalize_json(&value); + Ok(serde_json::to_vec_pretty(&value)?) +} + +fn schema_for_type() -> RootSchema +where + T: JsonSchema, +{ + SchemaSettings::draft07() + .with(|settings| { + settings.option_add_null_type = false; + }) + .into_generator() + .into_root_schema_for::() +} + +fn canonicalize_json(value: &Value) -> Value { + match value { + Value::Array(items) => Value::Array(items.iter().map(canonicalize_json).collect()), + Value::Object(map) => { + let mut entries: Vec<_> = map.iter().collect(); + entries.sort_by_key(|(key, _)| *key); + let mut sorted = Map::with_capacity(map.len()); + for (key, child) in entries { + sorted.insert(key.clone(), canonicalize_json(child)); + } + Value::Object(sorted) + } + _ => value.clone(), + } +} + +fn session_start_hook_event_name_schema(_gen: &mut SchemaGenerator) -> Schema { + string_const_schema("SessionStart") +} + +fn session_end_hook_event_name_schema(_gen: &mut SchemaGenerator) -> Schema { + string_const_schema("SessionEnd") +} + +fn session_end_reason_schema(_gen: &mut SchemaGenerator) -> Schema { + string_const_schema("other") +} + +fn post_tool_use_hook_event_name_schema(_gen: &mut SchemaGenerator) -> Schema { + string_const_schema("PostToolUse") +} + +fn pre_compact_hook_event_name_schema(_gen: &mut SchemaGenerator) -> Schema { + string_const_schema("PreCompact") +} + +fn post_compact_hook_event_name_schema(_gen: &mut SchemaGenerator) -> Schema { + string_const_schema("PostCompact") +} + +fn pre_tool_use_hook_event_name_schema(_gen: &mut SchemaGenerator) -> Schema { + string_const_schema("PreToolUse") +} + +fn permission_request_hook_event_name_schema(_gen: &mut SchemaGenerator) -> Schema { + string_const_schema("PermissionRequest") +} + +fn user_prompt_submit_hook_event_name_schema(_gen: &mut SchemaGenerator) -> Schema { + string_const_schema("UserPromptSubmit") +} + +fn subagent_start_hook_event_name_schema(_gen: &mut SchemaGenerator) -> Schema { + string_const_schema("SubagentStart") +} + +fn subagent_stop_hook_event_name_schema(_gen: &mut SchemaGenerator) -> Schema { + string_const_schema("SubagentStop") +} + +fn stop_hook_event_name_schema(_gen: &mut SchemaGenerator) -> Schema { + string_const_schema("Stop") +} + +fn interrupt_hook_event_name_schema(_gen: &mut SchemaGenerator) -> Schema { + string_const_schema("Interrupt") +} + +fn permission_mode_schema(_gen: &mut SchemaGenerator) -> Schema { + string_enum_schema(&[ + "default", + "acceptEdits", + "plan", + "dontAsk", + "bypassPermissions", + ]) +} + +fn session_start_source_schema(_gen: &mut SchemaGenerator) -> Schema { + string_enum_schema(&["startup", "resume", "clear", "compact", "fork"]) +} + +fn compaction_trigger_schema(_gen: &mut SchemaGenerator) -> Schema { + string_enum_schema(&["manual", "auto"]) +} + +fn string_const_schema(value: &str) -> Schema { + let mut schema = SchemaObject { + instance_type: Some(InstanceType::String.into()), + ..Default::default() + }; + schema.const_value = Some(Value::String(value.to_string())); + Schema::Object(schema) +} + +fn string_enum_schema(values: &[&str]) -> Schema { + let mut schema = SchemaObject { + instance_type: Some(InstanceType::String.into()), + ..Default::default() + }; + schema.enum_values = Some( + values + .iter() + .map(|value| Value::String((*value).to_string())) + .collect(), + ); + Schema::Object(schema) +} + +fn default_continue() -> bool { + true +} + +#[cfg(test)] +mod tests { + use super::INTERRUPT_INPUT_FIXTURE; + use super::INTERRUPT_OUTPUT_FIXTURE; + use super::InterruptCommandInput; + use super::NullableString; + use super::PERMISSION_REQUEST_INPUT_FIXTURE; + use super::PERMISSION_REQUEST_OUTPUT_FIXTURE; + use super::POST_COMPACT_INPUT_FIXTURE; + use super::POST_COMPACT_OUTPUT_FIXTURE; + use super::POST_TOOL_USE_INPUT_FIXTURE; + use super::POST_TOOL_USE_OUTPUT_FIXTURE; + use super::PRE_COMPACT_INPUT_FIXTURE; + use super::PRE_COMPACT_OUTPUT_FIXTURE; + use super::PRE_TOOL_USE_INPUT_FIXTURE; + use super::PRE_TOOL_USE_OUTPUT_FIXTURE; + use super::PermissionRequestCommandInput; + use super::PermissionRequestCommandOutputWire; + use super::PostCompactCommandInput; + use super::PostToolUseCommandInput; + use super::PostToolUseCommandOutputWire; + use super::PreCompactCommandInput; + use super::PreToolUseCommandInput; + use super::PreToolUseCommandOutputWire; + use super::SESSION_END_INPUT_FIXTURE; + use super::SESSION_START_INPUT_FIXTURE; + use super::SESSION_START_OUTPUT_FIXTURE; + use super::STOP_INPUT_FIXTURE; + use super::STOP_OUTPUT_FIXTURE; + use super::SUBAGENT_START_INPUT_FIXTURE; + use super::SUBAGENT_START_OUTPUT_FIXTURE; + use super::SUBAGENT_STOP_INPUT_FIXTURE; + use super::SUBAGENT_STOP_OUTPUT_FIXTURE; + use super::SessionStartCommandOutputWire; + use super::StopCommandInput; + use super::SubagentCommandInputFields; + use super::SubagentStartCommandInput; + use super::SubagentStartCommandOutputWire; + use super::SubagentStopCommandInput; + use super::USER_PROMPT_SUBMIT_INPUT_FIXTURE; + use super::USER_PROMPT_SUBMIT_OUTPUT_FIXTURE; + use super::UserPromptSubmitCommandInput; + use super::UserPromptSubmitCommandOutputWire; + use super::schema_json; + use super::write_schema_fixtures; + use crate::events::common::SubagentHookContext; + use pretty_assertions::assert_eq; + use schemars::JsonSchema; + use serde_json::Value; + use serde_json::json; + use tempfile::TempDir; + + fn expected_fixture(name: &str) -> &'static str { + match name { + POST_TOOL_USE_INPUT_FIXTURE => { + include_str!("../schema/generated/post-tool-use.command.input.schema.json") + } + POST_TOOL_USE_OUTPUT_FIXTURE => { + include_str!("../schema/generated/post-tool-use.command.output.schema.json") + } + PERMISSION_REQUEST_INPUT_FIXTURE => { + include_str!("../schema/generated/permission-request.command.input.schema.json") + } + PERMISSION_REQUEST_OUTPUT_FIXTURE => { + include_str!("../schema/generated/permission-request.command.output.schema.json") + } + POST_COMPACT_INPUT_FIXTURE => { + include_str!("../schema/generated/post-compact.command.input.schema.json") + } + POST_COMPACT_OUTPUT_FIXTURE => { + include_str!("../schema/generated/post-compact.command.output.schema.json") + } + PRE_COMPACT_INPUT_FIXTURE => { + include_str!("../schema/generated/pre-compact.command.input.schema.json") + } + PRE_COMPACT_OUTPUT_FIXTURE => { + include_str!("../schema/generated/pre-compact.command.output.schema.json") + } + PRE_TOOL_USE_INPUT_FIXTURE => { + include_str!("../schema/generated/pre-tool-use.command.input.schema.json") + } + PRE_TOOL_USE_OUTPUT_FIXTURE => { + include_str!("../schema/generated/pre-tool-use.command.output.schema.json") + } + SESSION_START_INPUT_FIXTURE => { + include_str!("../schema/generated/session-start.command.input.schema.json") + } + SESSION_START_OUTPUT_FIXTURE => { + include_str!("../schema/generated/session-start.command.output.schema.json") + } + SESSION_END_INPUT_FIXTURE => { + include_str!("../schema/generated/session-end.command.input.schema.json") + } + USER_PROMPT_SUBMIT_INPUT_FIXTURE => { + include_str!("../schema/generated/user-prompt-submit.command.input.schema.json") + } + USER_PROMPT_SUBMIT_OUTPUT_FIXTURE => { + include_str!("../schema/generated/user-prompt-submit.command.output.schema.json") + } + SUBAGENT_START_INPUT_FIXTURE => { + include_str!("../schema/generated/subagent-start.command.input.schema.json") + } + SUBAGENT_START_OUTPUT_FIXTURE => { + include_str!("../schema/generated/subagent-start.command.output.schema.json") + } + SUBAGENT_STOP_INPUT_FIXTURE => { + include_str!("../schema/generated/subagent-stop.command.input.schema.json") + } + SUBAGENT_STOP_OUTPUT_FIXTURE => { + include_str!("../schema/generated/subagent-stop.command.output.schema.json") + } + STOP_INPUT_FIXTURE => { + include_str!("../schema/generated/stop.command.input.schema.json") + } + STOP_OUTPUT_FIXTURE => { + include_str!("../schema/generated/stop.command.output.schema.json") + } + INTERRUPT_INPUT_FIXTURE => { + include_str!("../schema/generated/interrupt.command.input.schema.json") + } + INTERRUPT_OUTPUT_FIXTURE => { + include_str!("../schema/generated/interrupt.command.output.schema.json") + } + _ => panic!("unexpected fixture name: {name}"), + } + } + + fn normalize_newlines(value: &str) -> String { + value.replace("\r\n", "\n") + } + + fn assert_output_hook_event_name_const(definition: &str, expected: &str) { + let schema: Value = + serde_json::from_slice(&schema_json::().expect("serialize hook output schema")) + .expect("parse hook output schema"); + + assert_eq!( + schema["definitions"][definition]["properties"]["hookEventName"], + json!({ + "const": expected, + "type": "string", + }) + ); + } + + #[test] + fn generated_hook_schemas_match_fixtures() { + let temp_dir = TempDir::new().expect("create temp dir"); + let schema_root = temp_dir.path().join("schema"); + write_schema_fixtures(&schema_root).expect("write generated hook schemas"); + + for fixture in [ + POST_TOOL_USE_INPUT_FIXTURE, + POST_TOOL_USE_OUTPUT_FIXTURE, + PERMISSION_REQUEST_INPUT_FIXTURE, + PERMISSION_REQUEST_OUTPUT_FIXTURE, + POST_COMPACT_INPUT_FIXTURE, + POST_COMPACT_OUTPUT_FIXTURE, + PRE_COMPACT_INPUT_FIXTURE, + PRE_COMPACT_OUTPUT_FIXTURE, + PRE_TOOL_USE_INPUT_FIXTURE, + PRE_TOOL_USE_OUTPUT_FIXTURE, + SESSION_START_INPUT_FIXTURE, + SESSION_START_OUTPUT_FIXTURE, + SESSION_END_INPUT_FIXTURE, + USER_PROMPT_SUBMIT_INPUT_FIXTURE, + USER_PROMPT_SUBMIT_OUTPUT_FIXTURE, + SUBAGENT_START_INPUT_FIXTURE, + SUBAGENT_START_OUTPUT_FIXTURE, + SUBAGENT_STOP_INPUT_FIXTURE, + SUBAGENT_STOP_OUTPUT_FIXTURE, + STOP_INPUT_FIXTURE, + STOP_OUTPUT_FIXTURE, + INTERRUPT_INPUT_FIXTURE, + INTERRUPT_OUTPUT_FIXTURE, + ] { + let expected = normalize_newlines(expected_fixture(fixture)); + let actual = std::fs::read_to_string(schema_root.join("generated").join(fixture)) + .unwrap_or_else(|err| panic!("read generated schema {fixture}: {err}")); + let actual = normalize_newlines(&actual); + assert_eq!(expected, actual, "fixture should match generated schema"); + } + } + + #[test] + fn hook_specific_output_event_names_are_event_specific_in_output_schemas() { + assert_output_hook_event_name_const::( + "PermissionRequestHookSpecificOutputWire", + "PermissionRequest", + ); + assert_output_hook_event_name_const::( + "PostToolUseHookSpecificOutputWire", + "PostToolUse", + ); + assert_output_hook_event_name_const::( + "PreToolUseHookSpecificOutputWire", + "PreToolUse", + ); + assert_output_hook_event_name_const::( + "SessionStartHookSpecificOutputWire", + "SessionStart", + ); + assert_output_hook_event_name_const::( + "SubagentStartHookSpecificOutputWire", + "SubagentStart", + ); + assert_output_hook_event_name_const::( + "UserPromptSubmitHookSpecificOutputWire", + "UserPromptSubmit", + ); + } + + #[test] + fn turn_scoped_hook_inputs_include_codex_turn_id_extension() { + // Codex intentionally diverges from Claude's public hook docs here so + // internal hook consumers can key off the active turn. + let pre_tool_use: Value = serde_json::from_slice( + &schema_json::().expect("serialize pre tool use input schema"), + ) + .expect("parse pre tool use input schema"); + let post_tool_use: Value = serde_json::from_slice( + &schema_json::() + .expect("serialize post tool use input schema"), + ) + .expect("parse post tool use input schema"); + let pre_compact: Value = serde_json::from_slice( + &schema_json::().expect("serialize pre compact input schema"), + ) + .expect("parse pre compact input schema"); + let post_compact: Value = serde_json::from_slice( + &schema_json::().expect("serialize post compact input schema"), + ) + .expect("parse post compact input schema"); + let permission_request: Value = serde_json::from_slice( + &schema_json::() + .expect("serialize permission request input schema"), + ) + .expect("parse permission request input schema"); + let user_prompt_submit: Value = serde_json::from_slice( + &schema_json::() + .expect("serialize user prompt submit input schema"), + ) + .expect("parse user prompt submit input schema"); + let subagent_start: Value = serde_json::from_slice( + &schema_json::() + .expect("serialize subagent start input schema"), + ) + .expect("parse subagent start input schema"); + let subagent_stop: Value = serde_json::from_slice( + &schema_json::() + .expect("serialize subagent stop input schema"), + ) + .expect("parse subagent stop input schema"); + let stop: Value = serde_json::from_slice( + &schema_json::().expect("serialize stop input schema"), + ) + .expect("parse stop input schema"); + let interrupt: Value = serde_json::from_slice( + &schema_json::().expect("serialize interrupt input schema"), + ) + .expect("parse interrupt input schema"); + + for schema in [ + &pre_tool_use, + &permission_request, + &post_tool_use, + &pre_compact, + &post_compact, + &user_prompt_submit, + &subagent_start, + &subagent_stop, + &stop, + &interrupt, + ] { + assert_eq!(schema["properties"]["turn_id"]["type"], "string"); + assert!( + schema["required"] + .as_array() + .expect("schema required fields") + .contains(&Value::String("turn_id".to_string())) + ); + } + } + + #[test] + fn subagent_context_fields_are_optional_for_hooks_that_run_inside_subagents() { + let schemas = [ + schema_json::().expect("serialize pre tool use input schema"), + schema_json::() + .expect("serialize permission request input schema"), + schema_json::().expect("serialize post tool use input schema"), + schema_json::().expect("serialize pre compact input schema"), + schema_json::().expect("serialize post compact input schema"), + schema_json::() + .expect("serialize user prompt submit input schema"), + ]; + + for schema in schemas { + let schema: Value = serde_json::from_slice(&schema).expect("parse hook input schema"); + assert_eq!(schema["properties"]["agent_id"]["type"], "string"); + assert_eq!(schema["properties"]["agent_type"]["type"], "string"); + let required = schema["required"] + .as_array() + .expect("schema required fields"); + assert!(!required.contains(&Value::String("agent_id".to_string()))); + assert!(!required.contains(&Value::String("agent_type".to_string()))); + } + } + + #[test] + fn subagent_context_fields_serialize_flat_and_omit_when_absent() { + let subagent = SubagentCommandInputFields::from(Some(&SubagentHookContext { + agent_id: "agent-1".to_string(), + agent_type: "worker".to_string(), + })); + let input = PreToolUseCommandInput { + session_id: "session-1".to_string(), + turn_id: "turn-1".to_string(), + agent_id: subagent.agent_id, + agent_type: subagent.agent_type, + transcript_path: NullableString::from_path(/*path*/ None), + cwd: "/tmp".to_string(), + hook_event_name: "PreToolUse".to_string(), + model: "gpt-test".to_string(), + permission_mode: "default".to_string(), + tool_name: "Bash".to_string(), + tool_input: json!({ "command": "echo hello" }), + tool_use_id: "tool-1".to_string(), + }; + + assert_eq!( + serde_json::to_value(input).expect("serialize subagent hook input"), + json!({ + "session_id": "session-1", + "turn_id": "turn-1", + "agent_id": "agent-1", + "agent_type": "worker", + "transcript_path": null, + "cwd": "/tmp", + "hook_event_name": "PreToolUse", + "model": "gpt-test", + "permission_mode": "default", + "tool_name": "Bash", + "tool_input": { "command": "echo hello" }, + "tool_use_id": "tool-1", + }) + ); + + let root_input = PreToolUseCommandInput { + session_id: "session-1".to_string(), + turn_id: "turn-1".to_string(), + agent_id: None, + agent_type: None, + transcript_path: NullableString::from_path(/*path*/ None), + cwd: "/tmp".to_string(), + hook_event_name: "PreToolUse".to_string(), + model: "gpt-test".to_string(), + permission_mode: "default".to_string(), + tool_name: "Bash".to_string(), + tool_input: json!({ "command": "echo hello" }), + tool_use_id: "tool-1".to_string(), + }; + let root_input = serde_json::to_value(root_input).expect("serialize root hook input"); + assert_eq!(root_input.get("agent_id"), None); + assert_eq!(root_input.get("agent_type"), None); + } +} diff --git a/codex-rs/hooks/src/types.rs b/codex-rs/hooks/src/types.rs new file mode 100644 index 0000000000000000000000000000000000000000..6639308aa0a6637ed00fa2ceabe004efde450ba1 --- /dev/null +++ b/codex-rs/hooks/src/types.rs @@ -0,0 +1,152 @@ +use std::sync::Arc; + +use chrono::DateTime; +use chrono::SecondsFormat; +use chrono::Utc; +use codex_protocol::ThreadId; +use codex_utils_absolute_path::AbsolutePathBuf; +use futures::future::BoxFuture; +use serde::Serialize; +use serde::Serializer; + +pub type HookFn = Arc Fn(&'a HookPayload) -> BoxFuture<'a, HookResult> + Send + Sync>; + +#[derive(Debug)] +pub enum HookResult { + /// Success: hook completed successfully. + Success, + /// FailedContinue: hook failed, but other subsequent hooks should still execute and the + /// operation should continue. + FailedContinue(Box), + /// FailedAbort: hook failed, other subsequent hooks should not execute, and the operation + /// should be aborted. + FailedAbort(Box), +} + +impl HookResult { + pub fn should_abort_operation(&self) -> bool { + matches!(self, Self::FailedAbort(_)) + } +} + +#[derive(Debug)] +pub struct HookResponse { + pub hook_name: String, + pub result: HookResult, +} + +#[derive(Clone)] +pub struct Hook { + pub name: String, + pub func: HookFn, +} + +impl Default for Hook { + fn default() -> Self { + Self { + name: "default".to_string(), + func: Arc::new(|_| Box::pin(async { HookResult::Success })), + } + } +} + +impl Hook { + pub async fn execute(&self, payload: &HookPayload) -> HookResponse { + HookResponse { + hook_name: self.name.clone(), + result: (self.func)(payload).await, + } + } +} + +#[derive(Debug, Serialize, Clone)] +#[serde(rename_all = "snake_case")] +pub struct HookPayload { + pub session_id: ThreadId, + pub cwd: AbsolutePathBuf, + #[serde(skip_serializing_if = "Option::is_none")] + pub client: Option, + #[serde(serialize_with = "serialize_triggered_at")] + pub triggered_at: DateTime, + pub hook_event: HookEvent, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "snake_case")] +pub struct HookEventAfterAgent { + pub thread_id: ThreadId, + pub turn_id: String, + pub input_messages: Vec, + pub last_assistant_message: Option, +} + +fn serialize_triggered_at(value: &DateTime, serializer: S) -> Result +where + S: Serializer, +{ + serializer.serialize_str(&value.to_rfc3339_opts(SecondsFormat::Secs, true)) +} + +#[derive(Debug, Clone, Serialize)] +#[serde(tag = "event_type", rename_all = "snake_case")] +pub enum HookEvent { + AfterAgent { + #[serde(flatten)] + event: HookEventAfterAgent, + }, +} + +#[cfg(test)] +mod tests { + use chrono::TimeZone; + use chrono::Utc; + use codex_protocol::ThreadId; + use codex_utils_absolute_path::test_support::PathBufExt; + use codex_utils_absolute_path::test_support::test_path_buf; + use pretty_assertions::assert_eq; + use serde_json::json; + + use super::HookEvent; + use super::HookEventAfterAgent; + use super::HookPayload; + + #[test] + fn hook_payload_serializes_stable_wire_shape() { + let session_id = ThreadId::new(); + let thread_id = ThreadId::new(); + let cwd = test_path_buf("/tmp").abs(); + let payload = HookPayload { + session_id, + cwd: cwd.clone(), + client: None, + triggered_at: Utc + .with_ymd_and_hms(2025, 1, 1, 0, 0, 0) + .single() + .expect("valid timestamp"), + hook_event: HookEvent::AfterAgent { + event: HookEventAfterAgent { + thread_id, + turn_id: "turn-1".to_string(), + input_messages: vec!["hello".to_string()], + last_assistant_message: Some("hi".to_string()), + }, + }, + }; + + let actual = serde_json::to_value(payload).expect("serialize hook payload"); + let expected = json!({ + "session_id": session_id.to_string(), + "cwd": cwd.display().to_string(), + "triggered_at": "2025-01-01T00:00:00Z", + "hook_event": { + "event_type": "after_agent", + "thread_id": thread_id.to_string(), + "turn_id": "turn-1", + "input_messages": ["hello"], + "last_assistant_message": "hi", + }, + }); + + assert_eq!(actual, expected); + } +} diff --git a/codex-rs/http-client/src/bin/custom_ca_probe.rs b/codex-rs/http-client/src/bin/custom_ca_probe.rs new file mode 100644 index 0000000000000000000000000000000000000000..c8db21510d6a947c7fcaa419b2932b4198369537 --- /dev/null +++ b/codex-rs/http-client/src/bin/custom_ca_probe.rs @@ -0,0 +1,100 @@ +//! Helper binary for exercising shared custom CA environment handling in tests. +//! +//! The shared reqwest client honors `CODEX_CA_CERTIFICATE` and `SSL_CERT_FILE`, but those +//! environment variables are process-global and unsafe to mutate in parallel test execution. This +//! probe keeps the behavior under test while letting integration tests (`tests/ca_env.rs`) set +//! env vars per-process, proving: +//! +//! - env precedence is respected, +//! - multi-cert PEM bundles load, +//! - error messages guide users when CA files are invalid. +//! - optional HTTPS probes can complete a request through the constructed client. +//! +//! The detailed explanation of what "hermetic" means here lives in `codex_http_client::custom_ca`. +//! This binary exists so the tests can exercise +//! [`codex_http_client::build_reqwest_client_for_subprocess_tests`] in a separate process without +//! duplicating client-construction logic. + +use std::env; +use std::process; +use std::time::Duration; + +const PROBE_TLS13_ENV: &str = "CODEX_CUSTOM_CA_PROBE_TLS13"; +const PROBE_PROXY_ENV: &str = "CODEX_CUSTOM_CA_PROBE_PROXY"; +const PROBE_URL_ENV: &str = "CODEX_CUSTOM_CA_PROBE_URL"; + +fn main() { + let runtime = match tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + { + Ok(runtime) => runtime, + Err(error) => { + eprintln!("failed to create probe runtime: {error}"); + process::exit(1); + } + }; + + match runtime.block_on(run_probe()) { + Ok(()) => println!("ok"), + Err(error) => { + eprintln!("{error}"); + process::exit(1); + } + } +} + +async fn run_probe() -> Result<(), String> { + let proxy_url = env::var(PROBE_PROXY_ENV).ok(); + let target_url = env::var(PROBE_URL_ENV).ok(); + let mut builder = reqwest::Client::builder(); + if target_url.is_some() { + builder = builder.timeout(Duration::from_secs(5)); + } + if env::var_os(PROBE_TLS13_ENV).is_some() { + builder = builder.min_tls_version(reqwest::tls::Version::TLS_1_3); + } + + let client = build_probe_client(builder, proxy_url.as_deref())?; + if let Some(url) = target_url { + post_probe_request(&client, &url).await?; + } + Ok(()) +} + +fn build_probe_client( + builder: reqwest::ClientBuilder, + proxy_url: Option<&str>, +) -> Result { + if let Some(proxy_url) = proxy_url { + let proxy = reqwest::Proxy::https(proxy_url) + .map_err(|error| format!("failed to configure probe proxy {proxy_url}: {error}"))?; + return codex_http_client::build_reqwest_client_with_custom_ca(builder.proxy(proxy)) + .map_err(|error| error.to_string()); + } + + codex_http_client::build_reqwest_client_for_subprocess_tests(builder) + .map_err(|error| error.to_string()) +} + +async fn post_probe_request(client: &reqwest::Client, url: &str) -> Result<(), String> { + let response = client + .post(url) + .header("Content-Type", "application/x-www-form-urlencoded") + .body("grant_type=authorization_code&code=test") + .send() + .await + .map_err(|error| format!("probe request failed: {error:?}"))?; + let status = response.status(); + let body = response + .text() + .await + .map_err(|error| format!("failed to read probe response body: {error}"))?; + if !status.is_success() { + return Err(format!("probe request returned {status}: {body}")); + } + if body != "ok" { + return Err(format!("probe response body mismatch: {body}")); + } + Ok(()) +} diff --git a/codex-rs/http-client/src/chatgpt_cloudflare_cookies.rs b/codex-rs/http-client/src/chatgpt_cloudflare_cookies.rs new file mode 100644 index 0000000000000000000000000000000000000000..69124e031561e08d1e396f94eb2470891c297752 --- /dev/null +++ b/codex-rs/http-client/src/chatgpt_cloudflare_cookies.rs @@ -0,0 +1,409 @@ +use std::sync::Arc; +use std::sync::LazyLock; + +use reqwest::cookie::CookieStore; +use reqwest::cookie::Jar; +use reqwest::header::HeaderValue; + +use crate::chatgpt_hosts::is_allowed_chatgpt_host; + +// WARNING: this HTTP cookie store is process-global and may be shared across auth contexts. +// It must only ever contain Cloudflare infrastructure cookies and the `__oailb` +// routing cookie. Never extend this store to persist ChatGPT account, session, +// auth, or other user-specific cookie data. +static SHARED_CHATGPT_CLOUDFLARE_COOKIE_STORE: LazyLock> = + LazyLock::new(|| Arc::new(ChatGptCloudflareCookieStore::default())); + +#[derive(Debug, Default)] +struct ChatGptCloudflareCookieStore { + jar: Jar, +} + +pub(crate) struct ChatGptCookieStore { + cloudflare: Arc, + cookies: Vec, +} + +impl ChatGptCookieStore { + pub(crate) fn new(cookies: Vec) -> Self { + Self { + cloudflare: Arc::clone(&SHARED_CHATGPT_CLOUDFLARE_COOKIE_STORE), + cookies, + } + } + + pub(crate) fn configured_cookies(&self) -> &[HeaderValue] { + &self.cookies + } +} + +impl CookieStore for ChatGptCloudflareCookieStore { + fn set_cookies( + &self, + cookie_headers: &mut dyn Iterator, + url: &reqwest::Url, + ) { + if !is_chatgpt_cookie_url(url) { + return; + } + + let mut cloudflare_cookie_headers = + cookie_headers.filter(|header| is_allowed_cloudflare_set_cookie_header(header)); + self.jar.set_cookies(&mut cloudflare_cookie_headers, url); + } + + fn cookies(&self, url: &reqwest::Url) -> Option { + if is_chatgpt_cookie_url(url) { + self.jar.cookies(url).and_then(only_cloudflare_cookies) + } else { + None + } + } +} + +impl CookieStore for ChatGptCookieStore { + fn set_cookies( + &self, + cookie_headers: &mut dyn Iterator, + url: &reqwest::Url, + ) { + self.cloudflare.set_cookies(cookie_headers, url); + } + + fn cookies(&self, url: &reqwest::Url) -> Option { + if !is_chatgpt_cookie_url(url) { + return self.cloudflare.cookies(url); + } + let cloudflare = self.cloudflare.cookies(url); + let mut sensitive = cloudflare.as_ref().is_some_and(HeaderValue::is_sensitive); + let mut cookies = cloudflare + .map(|cloudflare| cloudflare.as_bytes().to_vec()) + .unwrap_or_default(); + for cookie in &self.cookies { + sensitive |= cookie.is_sensitive(); + if !cookies.is_empty() { + cookies.extend_from_slice(b"; "); + } + cookies.extend_from_slice(cookie.as_bytes()); + } + if cookies.is_empty() { + None + } else { + let mut header = HeaderValue::from_bytes(&cookies).ok()?; + header.set_sensitive(sensitive); + Some(header) + } + } +} + +/// Adds the process-local ChatGPT infrastructure cookie jar used by Codex HTTP clients. +/// +/// WARNING: this jar is global within the process. It is only acceptable because it hardcodes a +/// small allowlist of Cloudflare cookie names plus `__oailb` for routing, and refuses all other +/// ChatGPT cookies. Do not store ChatGPT account, session, auth, or other user-specific cookies +/// here. If a future caller needs those cookies, the store must be scoped to the auth/session +/// owner instead of shared globally. +pub fn with_chatgpt_cloudflare_cookie_store( + builder: reqwest::ClientBuilder, +) -> reqwest::ClientBuilder { + builder.cookie_provider(Arc::clone(&SHARED_CHATGPT_CLOUDFLARE_COOKIE_STORE)) +} + +fn is_chatgpt_cookie_url(url: &reqwest::Url) -> bool { + match url.scheme() { + "https" => {} + _ => return false, + } + + let Some(host) = url.host_str() else { + return false; + }; + + is_allowed_chatgpt_host(host) +} + +fn is_allowed_cloudflare_set_cookie_header(header: &HeaderValue) -> bool { + header + .to_str() + .ok() + .and_then(set_cookie_name) + .is_some_and(is_allowed_cloudflare_cookie_name) +} + +fn set_cookie_name(header: &str) -> Option<&str> { + let (name, _) = header.split_once('=')?; + let name = name.trim(); + (!name.is_empty()).then_some(name) +} + +fn only_cloudflare_cookies(header: HeaderValue) -> Option { + let header = header.to_str().ok()?; + let cookies = header + .split(';') + .filter_map(|cookie| { + let cookie = cookie.trim(); + let name = cookie.split_once('=')?.0.trim(); + is_allowed_cloudflare_cookie_name(name).then_some(cookie) + }) + .collect::>() + .join("; "); + + if cookies.is_empty() { + None + } else { + HeaderValue::from_str(&cookies).ok() + } +} + +fn is_allowed_cloudflare_cookie_name(name: &str) -> bool { + // `__oailb` is an OpenAI infrastructure routing cookie, not an authentication cookie. + // Keep this allowlist aligned with Cloudflare's documented service cookies: + // https://developers.cloudflare.com/fundamentals/reference/policies-compliances/cloudflare-cookies/ + matches!( + name, + "__cf_bm" + | "__cflb" + | "__cfruid" + | "__cfseq" + | "__cfwaitingroom" + | "__oailb" + | "_cfuvid" + | "cf_clearance" + | "cf_ob_info" + | "cf_use_ob" + ) || name.starts_with("cf_chl_") +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + use reqwest::cookie::CookieStore; + + #[test] + fn additional_cookies_use_current_path_scoped_cloudflare_cookies() { + let cloudflare = Arc::new(ChatGptCloudflareCookieStore::default()); + let mut sensitive_cookie = HeaderValue::from_static("additional=true"); + sensitive_cookie.set_sensitive(true); + let store = ChatGptCookieStore { + cloudflare: Arc::clone(&cloudflare), + cookies: vec![ + sensitive_cookie, + HeaderValue::from_static("another=enabled"), + ], + }; + let first_url = reqwest::Url::parse("https://cookies.chatgpt.com/a").unwrap(); + let second_url = reqwest::Url::parse("https://cookies.chatgpt.com/b").unwrap(); + let affinity = HeaderValue::from_static("__cflb=current; Path=/a; Secure"); + cloudflare.set_cookies(&mut std::iter::once(&affinity), &first_url); + + assert_eq!( + store.cookies(&first_url), + Some(HeaderValue::from_static( + "__cflb=current; additional=true; another=enabled" + )) + ); + assert!(store.cookies(&first_url).unwrap().is_sensitive()); + assert_eq!( + store.cookies(&second_url), + Some(HeaderValue::from_static("additional=true; another=enabled")) + ); + + let expired = HeaderValue::from_static("__cflb=; Max-Age=0; Path=/a; Secure"); + store.set_cookies(&mut std::iter::once(&expired), &first_url); + assert_eq!( + store.cookies(&first_url), + Some(HeaderValue::from_static("additional=true; another=enabled")) + ); + } + + #[test] + fn additional_cookies_are_restricted_to_https_chatgpt_hosts() { + let store = ChatGptCookieStore { + cloudflare: Arc::new(ChatGptCloudflareCookieStore::default()), + cookies: vec![HeaderValue::from_static("additional=true")], + }; + + for url in ["http://chatgpt.com/a", "https://api.openai.com/a"] { + let url = reqwest::Url::parse(url).unwrap(); + assert_eq!(store.cookies(&url), None); + } + let url = reqwest::Url::parse("https://chatgpt.com/a").unwrap(); + assert_eq!( + store.cookies(&url), + Some(HeaderValue::from_static("additional=true")) + ); + assert!(!store.cookies(&url).unwrap().is_sensitive()); + } + + #[test] + fn stores_and_returns_cloudflare_cookies_for_chatgpt_hosts() { + let store = ChatGptCloudflareCookieStore::default(); + let url = reqwest::Url::parse("https://chatgpt.com/backend-api/codex/responses").unwrap(); + let load_balancer = HeaderValue::from_static("__cflb=west; Path=/; Secure; HttpOnly"); + let cfuvid = HeaderValue::from_static("_cfuvid=visitor; Path=/; Secure; HttpOnly"); + let clearance = + HeaderValue::from_static("cf_clearance=clearance; Path=/; Secure; HttpOnly"); + + store.set_cookies(&mut [&load_balancer, &cfuvid, &clearance].into_iter(), &url); + + let mut cookies = store + .cookies(&url) + .and_then(|value| value.to_str().ok().map(str::to_string)) + .map(|header| { + header + .split("; ") + .map(str::to_string) + .collect::>() + }) + .unwrap_or_default(); + cookies.sort(); + assert_eq!( + cookies, + vec![ + "__cflb=west".to_string(), + "_cfuvid=visitor".to_string(), + "cf_clearance=clearance".to_string() + ] + ); + } + + #[test] + fn oailb_cookies_are_replayed_with_scope_and_expiration() { + let store = ChatGptCloudflareCookieStore::default(); + let url = reqwest::Url::parse("https://chatgpt.com/backend-api/codex/responses").unwrap(); + let cookie = HeaderValue::from_static( + "__oailb=route; Path=/backend-api; Max-Age=3600; Secure; HttpOnly; SameSite=Lax", + ); + store.set_cookies(&mut std::iter::once(&cookie), &url); + + let followup_url = reqwest::Url::parse("https://chatgpt.com/backend-api/ps/mcp").unwrap(); + assert_eq!( + store.cookies(&followup_url), + Some(HeaderValue::from_static("__oailb=route")) + ); + for outside_scope in [ + "https://chatgpt.com/", + "https://other.chatgpt.com/backend-api/ps/mcp", + "https://api.openai.com/backend-api/ps/mcp", + "http://chatgpt.com/backend-api/ps/mcp", + ] { + let outside_scope = reqwest::Url::parse(outside_scope).unwrap(); + assert_eq!(store.cookies(&outside_scope), None); + } + + let expired = HeaderValue::from_static( + "__oailb=; Path=/backend-api; Max-Age=0; Secure; HttpOnly; SameSite=Lax", + ); + store.set_cookies(&mut std::iter::once(&expired), &url); + assert_eq!(store.cookies(&followup_url), None); + } + + #[test] + fn ignores_non_chatgpt_cookies() { + let store = ChatGptCloudflareCookieStore::default(); + let url = reqwest::Url::parse("https://api.openai.com/v1/responses").unwrap(); + let set_cookie = HeaderValue::from_static("_cfuvid=visitor; Path=/; Secure; HttpOnly"); + + store.set_cookies(&mut std::iter::once(&set_cookie), &url); + + assert_eq!(store.cookies(&url), None); + } + + #[test] + fn ignores_non_cloudflare_cookies_for_chatgpt_hosts() { + let store = ChatGptCloudflareCookieStore::default(); + let url = reqwest::Url::parse("https://chatgpt.com/backend-api/codex/responses").unwrap(); + let set_cookie = HeaderValue::from_static( + "__Secure-next-auth.session-token=secret; Path=/; Secure; HttpOnly", + ); + + store.set_cookies(&mut std::iter::once(&set_cookie), &url); + + assert_eq!(store.cookies(&url), None); + } + + #[test] + fn ignores_mixed_non_cloudflare_cookies_for_chatgpt_hosts() { + let store = ChatGptCloudflareCookieStore::default(); + let url = reqwest::Url::parse("https://chatgpt.com/backend-api/codex/responses").unwrap(); + let cfuvid = HeaderValue::from_static("_cfuvid=visitor; Path=/; Secure; HttpOnly"); + let account_cookie = + HeaderValue::from_static("chatgpt_session=secret; Path=/; Secure; HttpOnly"); + + store.set_cookies(&mut [&cfuvid, &account_cookie].into_iter(), &url); + + assert_eq!( + store + .cookies(&url) + .and_then(|value| value.to_str().ok().map(str::to_string)), + Some("_cfuvid=visitor".to_string()) + ); + } + + #[test] + fn does_not_return_chatgpt_cloudflare_cookies_for_other_hosts() { + let store = ChatGptCloudflareCookieStore::default(); + let chatgpt_url = + reqwest::Url::parse("https://chatgpt.com/backend-api/codex/responses").unwrap(); + let api_url = reqwest::Url::parse("https://api.openai.com/v1/responses").unwrap(); + let set_cookie = HeaderValue::from_static("_cfuvid=visitor; Path=/; Secure; HttpOnly"); + + store.set_cookies(&mut std::iter::once(&set_cookie), &chatgpt_url); + + assert_eq!(store.cookies(&api_url), None); + } + + #[test] + fn rejects_plain_http_chatgpt_cookie_urls() { + let store = ChatGptCloudflareCookieStore::default(); + let http_url = reqwest::Url::parse("http://chatgpt.com/backend-api/codex/responses") + .expect("URL should parse"); + let https_url = reqwest::Url::parse("https://chatgpt.com/backend-api/codex/responses") + .expect("URL should parse"); + let set_cookie = HeaderValue::from_static("_cfuvid=visitor; Path=/; Secure; HttpOnly"); + + store.set_cookies(&mut std::iter::once(&set_cookie), &http_url); + + assert_eq!(store.cookies(&http_url), None); + assert_eq!(store.cookies(&https_url), None); + } + + #[test] + fn only_allows_https_urls() { + let url = reqwest::Url::parse("http://chatgpt.com/backend-api/codex/responses").unwrap(); + + assert!(!is_chatgpt_cookie_url(&url)); + + let url = reqwest::Url::parse("wss://chatgpt.com/backend-api/codex/responses").unwrap(); + + assert!(!is_chatgpt_cookie_url(&url)); + } + + #[test] + fn allows_only_known_cloudflare_cookie_names() { + for name in [ + "__cf_bm", + "__cflb", + "__cfruid", + "__cfseq", + "__cfwaitingroom", + "_cfuvid", + "cf_clearance", + "cf_ob_info", + "cf_use_ob", + "cf_chl_rc_i", + ] { + assert!(is_allowed_cloudflare_cookie_name(name)); + } + + for name in [ + "__Secure-next-auth.session-token", + "chatgpt_session", + "oai-auth-token", + "not_cf_clearance", + ] { + assert!(!is_allowed_cloudflare_cookie_name(name)); + } + } +} diff --git a/codex-rs/http-client/src/chatgpt_hosts.rs b/codex-rs/http-client/src/chatgpt_hosts.rs new file mode 100644 index 0000000000000000000000000000000000000000..0426c8bb36b29de0b79602aa1a176f5e76305398 --- /dev/null +++ b/codex-rs/http-client/src/chatgpt_hosts.rs @@ -0,0 +1,39 @@ +/// Returns whether `host` is one of the ChatGPT hosts Codex is allowed to treat +/// as first-party ChatGPT HTTP traffic. +pub fn is_allowed_chatgpt_host(host: &str) -> bool { + const EXACT_HOSTS: &[&str] = &["chatgpt.com", "chat.openai.com", "chatgpt-staging.com"]; + const SUBDOMAIN_SUFFIXES: &[&str] = &[".chatgpt.com", ".chatgpt-staging.com"]; + + EXACT_HOSTS.contains(&host) + || SUBDOMAIN_SUFFIXES + .iter() + .any(|suffix| host.ends_with(suffix)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn recognizes_chatgpt_hosts_without_suffix_tricks() { + for host in [ + "chatgpt.com", + "foo.chatgpt.com", + "staging.chatgpt.com", + "chat.openai.com", + "chatgpt-staging.com", + "api.chatgpt-staging.com", + ] { + assert!(is_allowed_chatgpt_host(host)); + } + + for host in [ + "evilchatgpt.com", + "chatgpt.com.evil.example", + "api.openai.com", + "foo.chat.openai.com", + ] { + assert!(!is_allowed_chatgpt_host(host)); + } + } +} diff --git a/codex-rs/http-client/src/client.rs b/codex-rs/http-client/src/client.rs new file mode 100644 index 0000000000000000000000000000000000000000..9cda749aa27581101a4f718644151f88c508100d --- /dev/null +++ b/codex-rs/http-client/src/client.rs @@ -0,0 +1,356 @@ +//! Reusable HTTP client and request-builder wrappers. + +use http::Error as HttpRequestBuildError; +use http::HeaderMap; +use http::HeaderName; +use http::HeaderValue; +use opentelemetry::global; +use opentelemetry::propagation::Injector; +use reqwest::IntoUrl; +use reqwest::Method; +use serde::Serialize; +use std::fmt::Display; +use std::time::Duration; +use tracing::Span; +use tracing_opentelemetry::OpenTelemetrySpanExt; + +pub type HttpError = reqwest::Error; +pub type HttpResponse = reqwest::Response; + +/// Reusable HTTP client wrapper with shared tracing and request-diagnostic behavior. +/// +/// Product callers should obtain this through [`crate::HttpClientFactory`] for a fixed +/// destination or use [`crate::RouteAwareClientPool`] when request and redirect URLs can vary. +#[derive(Clone, Debug)] +pub struct HttpClient { + inner: reqwest::Client, + request_logging: RequestLogging, +} + +impl HttpClient { + pub fn new(inner: reqwest::Client) -> Self { + Self::from_parts(inner, RequestLogging::Enabled) + } + + /// Creates a client that suppresses request URL and response-header diagnostics. + /// + /// Use this for endpoints whose URLs or headers may contain credentials that are redacted by + /// the caller above the HTTP transport boundary. + pub fn new_without_request_logging(inner: reqwest::Client) -> Self { + Self::from_parts(inner, RequestLogging::Disabled) + } + + pub(crate) fn from_parts(inner: reqwest::Client, request_logging: RequestLogging) -> Self { + Self { + inner, + request_logging, + } + } + + pub fn get(&self, url: U) -> RequestBuilder + where + U: IntoUrl, + { + self.request(Method::GET, url) + } + + pub fn head(&self, url: U) -> RequestBuilder + where + U: IntoUrl, + { + self.request(Method::HEAD, url) + } + + pub fn post(&self, url: U) -> RequestBuilder + where + U: IntoUrl, + { + self.request(Method::POST, url) + } + + pub fn delete(&self, url: U) -> RequestBuilder + where + U: IntoUrl, + { + self.request(Method::DELETE, url) + } + + pub fn request(&self, method: Method, url: U) -> RequestBuilder + where + U: IntoUrl, + { + let url_str = url.as_str().to_string(); + RequestBuilder::new( + self.inner.request(method.clone(), url), + method, + url_str, + self.request_logging, + ) + } + + pub(crate) async fn execute( + &self, + request: reqwest::Request, + ) -> Result { + let method = request.method().clone(); + let url = request.url().to_string(); + + match self.execute_without_request_logging(request).await { + Ok(response) => { + self.log_response(&method, &url, &response); + Ok(response) + } + Err(error) => { + self.log_error(&method, &url, &error); + Err(error) + } + } + } + + pub(crate) async fn execute_without_request_logging( + &self, + mut request: reqwest::Request, + ) -> Result { + request.headers_mut().extend(trace_headers()); + self.inner.execute(request).await + } + + pub(crate) fn log_response(&self, method: &Method, url: &str, response: &reqwest::Response) { + if self.request_logging == RequestLogging::Enabled { + tracing::debug!( + method = %method, + url = %url, + status = %response.status(), + headers = ?response.headers(), + version = ?response.version(), + "Request completed" + ); + } + } + + pub(crate) fn log_error(&self, method: &Method, url: &str, error: &reqwest::Error) { + if self.request_logging == RequestLogging::Enabled { + tracing::debug!( + method = %method, + url = %url, + status = error.status().map(|status| status.as_u16()), + error = %error, + "Request failed" + ); + } + } + pub(crate) fn log_error_summary(&self, method: &Method, url: &str, error: &reqwest::Error) { + if self.request_logging == RequestLogging::Enabled { + tracing::debug!( + method = %method, + url = %url, + status = error.status().map(|status| status.as_u16()), + is_timeout = error.is_timeout(), + is_connect = error.is_connect(), + "Request failed" + ); + } + } + + pub(crate) const fn request_logging_enabled(&self) -> bool { + matches!(self.request_logging, RequestLogging::Enabled) + } +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(crate) enum RequestLogging { + #[default] + Enabled, + Disabled, +} + +#[must_use = "requests are not sent unless `send` is awaited"] +#[derive(Debug)] +pub struct RequestBuilder { + builder: reqwest::RequestBuilder, + method: Method, + url: String, + request_logging: RequestLogging, +} + +impl RequestBuilder { + fn new( + builder: reqwest::RequestBuilder, + method: Method, + url: String, + request_logging: RequestLogging, + ) -> Self { + Self { + builder, + method, + url, + request_logging, + } + } + + fn map(self, f: impl FnOnce(reqwest::RequestBuilder) -> reqwest::RequestBuilder) -> Self { + Self { + builder: f(self.builder), + method: self.method, + url: self.url, + request_logging: self.request_logging, + } + } + + pub fn headers(self, headers: HeaderMap) -> Self { + self.map(|builder| builder.headers(headers)) + } + + pub fn header(self, key: K, value: V) -> Self + where + HeaderName: TryFrom, + >::Error: Into, + HeaderValue: TryFrom, + >::Error: Into, + { + self.map(|builder| builder.header(key, value)) + } + + pub fn bearer_auth(self, token: T) -> Self + where + T: Display, + { + self.map(|builder| builder.bearer_auth(token)) + } + + pub fn timeout(self, timeout: Duration) -> Self { + self.map(|builder| builder.timeout(timeout)) + } + + pub fn json(self, value: &T) -> Self + where + T: ?Sized + Serialize, + { + self.map(|builder| builder.json(value)) + } + + pub fn query(self, query: &T) -> Self + where + T: ?Sized + Serialize, + { + self.map(|builder| builder.query(query)) + } + + pub fn body(self, body: B) -> Self + where + B: Into, + { + self.map(|builder| builder.body(body)) + } + + pub async fn send(self) -> Result { + let headers = trace_headers(); + + match self.builder.headers(headers).send().await { + Ok(response) => { + if self.request_logging == RequestLogging::Enabled { + tracing::debug!( + method = %self.method, + url = %self.url, + status = %response.status(), + headers = ?response.headers(), + version = ?response.version(), + "Request completed" + ); + } + + Ok(response) + } + Err(error) => { + if self.request_logging == RequestLogging::Enabled { + let status = error.status(); + tracing::debug!( + method = %self.method, + url = %self.url, + status = status.map(|s| s.as_u16()), + error = %error, + "Request failed" + ); + } + Err(error) + } + } + } +} + +struct HeaderMapInjector<'a>(&'a mut HeaderMap); + +impl<'a> Injector for HeaderMapInjector<'a> { + fn set(&mut self, key: &str, value: String) { + if let (Ok(name), Ok(val)) = ( + HeaderName::from_bytes(key.as_bytes()), + HeaderValue::from_str(&value), + ) { + self.0.insert(name, val); + } + } +} + +pub(crate) fn trace_headers() -> HeaderMap { + let mut headers = HeaderMap::new(); + global::get_text_map_propagator(|prop| { + prop.inject_context( + &Span::current().context(), + &mut HeaderMapInjector(&mut headers), + ); + }); + headers +} + +#[cfg(test)] +mod tests { + use super::*; + use opentelemetry::propagation::Extractor; + use opentelemetry::propagation::TextMapPropagator; + use opentelemetry::trace::TraceContextExt; + use opentelemetry::trace::TracerProvider; + use opentelemetry_sdk::propagation::TraceContextPropagator; + use opentelemetry_sdk::trace::SdkTracerProvider; + use pretty_assertions::assert_eq; + use tracing::trace_span; + use tracing_subscriber::layer::SubscriberExt; + use tracing_subscriber::util::SubscriberInitExt; + + #[test] + fn inject_trace_headers_uses_current_span_context() { + global::set_text_map_propagator(TraceContextPropagator::new()); + + let provider = SdkTracerProvider::builder().build(); + let tracer = provider.tracer("test-tracer"); + let subscriber = + tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer)); + let _guard = subscriber.set_default(); + + let span = trace_span!("client_request"); + let _entered = span.enter(); + let span_context = span.context().span().span_context().clone(); + + let headers = trace_headers(); + + let extractor = HeaderMapExtractor(&headers); + let extracted = TraceContextPropagator::new().extract(&extractor); + let extracted_span = extracted.span(); + let extracted_context = extracted_span.span_context(); + + assert!(extracted_context.is_valid()); + assert_eq!(extracted_context.trace_id(), span_context.trace_id()); + assert_eq!(extracted_context.span_id(), span_context.span_id()); + } + + struct HeaderMapExtractor<'a>(&'a HeaderMap); + + impl<'a> Extractor for HeaderMapExtractor<'a> { + fn get(&self, key: &str) -> Option<&str> { + self.0.get(key).and_then(|value| value.to_str().ok()) + } + + fn keys(&self) -> Vec<&str> { + self.0.keys().map(HeaderName::as_str).collect() + } + } +} diff --git a/codex-rs/http-client/src/client_builder.rs b/codex-rs/http-client/src/client_builder.rs new file mode 100644 index 0000000000000000000000000000000000000000..9f374293a1025c875bf0cc24be48c7201993c50c --- /dev/null +++ b/codex-rs/http-client/src/client_builder.rs @@ -0,0 +1,323 @@ +//! HTTP client construction that makes outbound proxy policy explicit. +//! +//! Product traffic should normally enter through [`HttpClientFactory`] for a fixed destination or +//! [`crate::RouteAwareClientPool`] when request and redirect URLs can vary. The direct and +//! transport-default terminal methods exist only for narrow exceptional or legacy compatibility +//! paths. + +use http::HeaderMap; +use std::sync::Arc; +use std::time::Duration; + +use codex_utils_rustls_provider::ensure_rustls_crypto_provider; + +use crate::BuildCustomCaTransportError; +use crate::BuildRouteAwareHttpClientError; +use crate::ClientRouteClass; +use crate::HttpClient; +use crate::HttpClientFactory; +use crate::OutboundProxyRoute; +use crate::chatgpt_cloudflare_cookies::ChatGptCookieStore; +use crate::client::RequestLogging; +use crate::custom_ca::build_reqwest_client_with_custom_ca; +use crate::with_chatgpt_cloudflare_cookie_store; + +/// Configures an [`HttpClient`] without exposing the underlying HTTP implementation. +/// +/// Product traffic should prefer [`HttpClientFactory::build_client`] or finish this builder with +/// [`Self::build_respecting_outbound_proxy_policy`]. The other terminal methods deliberately +/// bypass the factory and are restricted to documented exceptional or legacy compatibility paths. +#[derive(Clone)] +pub struct HttpClientBuilder { + default_headers: Option, + follow_redirects: bool, + connect_timeout: Option, + chatgpt_cloudflare_cookie_store: bool, + chatgpt_cookie_store: Option>, + request_logging: RequestLogging, + tls_backend: TlsBackend, +} + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +enum TlsBackend { + #[default] + TransportDefault, + Rustls, +} + +impl HttpClientFactory { + /// Builds an HTTP client for one fixed destination using the configured proxy policy. + /// + /// This is the preferred construction path for product traffic that uses a fixed destination. + /// Use [`crate::RouteAwareClientPool`] instead when request or redirect URLs can vary. + pub fn build_client( + &self, + request_url: &str, + route_class: ClientRouteClass, + ) -> Result { + HttpClientBuilder::new().build_respecting_outbound_proxy_policy( + self, + request_url, + route_class, + ) + } + + /// Builds a policy-aware client without request URL or response-header diagnostics. + /// + /// This has the same routing guidance as [`Self::build_client`]. + pub fn build_client_without_request_logging( + &self, + request_url: &str, + route_class: ClientRouteClass, + ) -> Result { + HttpClientBuilder::new() + .without_request_logging() + .build_respecting_outbound_proxy_policy(self, request_url, route_class) + } +} + +impl HttpClientBuilder { + pub fn new() -> Self { + Self::default() + } + + pub fn default_headers(mut self, headers: HeaderMap) -> Self { + self.default_headers = Some(headers); + self + } + + pub fn without_redirects(mut self) -> Self { + self.follow_redirects = false; + self + } + + pub(crate) fn follows_redirects(&self) -> bool { + self.follow_redirects + } + + pub(crate) fn with_rustls_tls(mut self) -> Self { + self.tls_backend = TlsBackend::Rustls; + self + } + + /// Limits only connection establishment, not the request as a whole. + pub fn connect_timeout(mut self, timeout: Duration) -> Self { + self.connect_timeout = Some(timeout); + self + } + + pub fn with_chatgpt_cloudflare_cookie_store(mut self) -> Self { + self.chatgpt_cloudflare_cookie_store = true; + self + } + + /// Uses the factory's configured ChatGPT cookies without changing proxy behavior. + pub fn with_chatgpt_cookies(mut self, http_client_factory: &HttpClientFactory) -> Self { + self.chatgpt_cloudflare_cookie_store = true; + self.chatgpt_cookie_store = http_client_factory.chatgpt_cookie_store(); + self + } + + /// Suppresses request URL and response-header diagnostics. + pub fn without_request_logging(mut self) -> Self { + self.request_logging = RequestLogging::Disabled; + self + } + + /// Builds a client that honors the [`HttpClientFactory`] outbound proxy policy. + /// + /// This is the preferred terminal method for product traffic. The request URL is used to + /// resolve a concrete direct or proxy route when the factory is configured with + /// [`crate::OutboundProxyPolicy::RespectSystemProxy`]. + pub fn build_respecting_outbound_proxy_policy( + mut self, + http_client_factory: &HttpClientFactory, + request_url: &str, + route_class: ClientRouteClass, + ) -> Result { + self.chatgpt_cookie_store = http_client_factory.chatgpt_cookie_store(); + let (builder, request_logging) = self.into_reqwest_parts(); + let inner = http_client_factory.build_reqwest_client(builder, request_url, route_class)?; + Ok(HttpClient::from_parts(inner, request_logging)) + } + + /// Builds a client for a route that was already resolved by a route-aware caller. + pub(crate) fn build_for_resolved_route( + mut self, + http_client_factory: &HttpClientFactory, + route_class: ClientRouteClass, + route: &OutboundProxyRoute, + ) -> Result { + self.chatgpt_cookie_store = http_client_factory.chatgpt_cookie_store(); + let (builder, request_logging) = self.into_reqwest_parts(); + let inner = http_client_factory.build_reqwest_client_for_resolved_route( + builder, + route_class, + route, + )?; + Ok(HttpClient::from_parts(inner, request_logging)) + } + + /// Builds a client using the transport's default proxy behavior. + /// + /// # Legacy compatibility only + /// + /// This bypasses [`HttpClientFactory`] and therefore does not honor its configured outbound + /// proxy policy. New product traffic must use [`Self::build_respecting_outbound_proxy_policy`] + /// or [`HttpClientFactory::build_client`]. + #[deprecated( + note = "legacy compatibility only; use HttpClientFactory::build_client or build_respecting_outbound_proxy_policy" + )] + pub fn build_with_transport_default_proxy( + self, + ) -> Result { + self.build_with_proxy_routing(ProxyRouting::TransportDefault) + } + + /// Builds a client that connects directly without using a proxy. + /// + /// # Exceptional use only + /// + /// This bypasses [`HttpClientFactory`] and is appropriate only when bypassing proxy discovery + /// is itself required: for example, a hermetic local test fixture, a localhost callback, or + /// sandbox traffic whose egress routing is handled separately. Ordinary outbound product + /// traffic must use [`Self::build_respecting_outbound_proxy_policy`] or + /// [`HttpClientFactory::build_client`]. + pub fn build_direct(self) -> Result { + self.build_with_proxy_routing(ProxyRouting::Direct) + } + + /// Builds a transport-default client while preserving the legacy custom-CA fallback. + /// + /// # Legacy compatibility only + /// + /// This preserves call sites that historically logged a custom-CA error and continued with + /// system roots. New product traffic must propagate construction errors through + /// [`Self::build_respecting_outbound_proxy_policy`] or [`HttpClientFactory::build_client`]. + #[deprecated( + note = "legacy custom-CA fallback only; use HttpClientFactory::build_client or build_respecting_outbound_proxy_policy" + )] + pub fn build_with_transport_default_proxy_and_custom_ca_fallback(self) -> HttpClient { + self.build_with_custom_ca_fallback(ProxyRouting::TransportDefault) + } + + /// Builds a direct client while preserving the legacy custom-CA fallback. + /// + /// # Legacy compatibility only + /// + /// This combines the exceptional proxy bypass described by [`Self::build_direct`] with the + /// historical behavior of logging a custom-CA error and continuing with system roots. + #[deprecated( + note = "legacy custom-CA fallback only; use build_direct and propagate construction errors" + )] + pub fn build_direct_with_custom_ca_fallback(self) -> HttpClient { + self.build_with_custom_ca_fallback(ProxyRouting::Direct) + } + + fn build_with_proxy_routing( + self, + proxy_routing: ProxyRouting, + ) -> Result { + let request_logging = self.request_logging; + build_reqwest_client_with_custom_ca(self.reqwest_builder(proxy_routing)) + .map(|inner| HttpClient::from_parts(inner, request_logging)) + } + + fn build_with_custom_ca_fallback(self, proxy_routing: ProxyRouting) -> HttpClient { + self.build_with_custom_ca_fallback_using(proxy_routing, build_reqwest_client_with_custom_ca) + } + + fn build_with_custom_ca_fallback_using( + self, + proxy_routing: ProxyRouting, + build_with_custom_ca: impl FnOnce( + reqwest::ClientBuilder, + ) + -> Result, + ) -> HttpClient { + let request_logging = self.request_logging; + match build_with_custom_ca(self.clone().reqwest_builder(proxy_routing)) { + Ok(inner) => HttpClient::from_parts(inner, request_logging), + Err(error) => { + tracing::event!( + target: "codex_otel.log_only", + tracing::Level::WARN, + event.name = "codex.http_client.custom_ca_fallback", + "HTTP client fell back to system root certificates" + ); + tracing::warn!(error = %error, "failed to build HTTP client with custom CA"); + self.reqwest_builder(proxy_routing) + .build() + .map(|inner| HttpClient::from_parts(inner, request_logging)) + .unwrap_or_else(|fallback_error| { + tracing::warn!( + error = %fallback_error, + "failed to build fallback HTTP client" + ); + HttpClient::from_parts(reqwest::Client::new(), request_logging) + }) + } + } + } + + fn into_reqwest_parts(self) -> (reqwest::ClientBuilder, RequestLogging) { + let request_logging = self.request_logging; + (self.base_reqwest_builder(), request_logging) + } + + fn reqwest_builder(self, proxy_routing: ProxyRouting) -> reqwest::ClientBuilder { + let builder = self.base_reqwest_builder(); + match proxy_routing { + ProxyRouting::TransportDefault => builder, + ProxyRouting::Direct => builder.no_proxy(), + } + } + + fn base_reqwest_builder(self) -> reqwest::ClientBuilder { + let mut builder = reqwest::Client::builder(); + if self.tls_backend == TlsBackend::Rustls { + ensure_rustls_crypto_provider(); + builder = builder.use_rustls_tls(); + } + if let Some(default_headers) = self.default_headers { + builder = builder.default_headers(default_headers); + } + if !self.follow_redirects { + builder = builder.redirect(reqwest::redirect::Policy::none()); + } + if let Some(connect_timeout) = self.connect_timeout { + builder = builder.connect_timeout(connect_timeout); + } + if self.chatgpt_cloudflare_cookie_store { + builder = match self.chatgpt_cookie_store { + Some(store) => builder.cookie_provider(store), + None => with_chatgpt_cloudflare_cookie_store(builder), + }; + } + builder + } +} + +impl Default for HttpClientBuilder { + fn default() -> Self { + Self { + default_headers: None, + follow_redirects: true, + connect_timeout: None, + chatgpt_cloudflare_cookie_store: false, + chatgpt_cookie_store: None, + request_logging: RequestLogging::Enabled, + tls_backend: TlsBackend::TransportDefault, + } + } +} + +#[derive(Clone, Copy)] +enum ProxyRouting { + TransportDefault, + Direct, +} + +#[cfg(test)] +#[path = "client_builder_tests.rs"] +mod tests; diff --git a/codex-rs/http-client/src/client_builder_tests.rs b/codex-rs/http-client/src/client_builder_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..0b4cbb4b7d0230bf6e91a49a03b7421088cf9c04 --- /dev/null +++ b/codex-rs/http-client/src/client_builder_tests.rs @@ -0,0 +1,52 @@ +use super::*; +use http::HeaderValue; +use std::io::Read; +use std::io::Write; +use std::path::PathBuf; + +#[tokio::test] +async fn custom_ca_fallback_preserves_builder_configuration() { + let listener = + std::net::TcpListener::bind(("127.0.0.1", 0)).expect("HTTP listener should bind"); + let address = listener + .local_addr() + .expect("HTTP listener should have an address"); + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("HTTP listener should accept"); + let mut request = Vec::new(); + let mut chunk = [0_u8; 1024]; + while !request.windows(4).any(|window| window == b"\r\n\r\n") { + let bytes_read = stream.read(&mut chunk).expect("HTTP request should read"); + assert!(bytes_read > 0, "HTTP request should include headers"); + request.extend_from_slice(&chunk[..bytes_read]); + } + stream + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") + .expect("HTTP listener should write response"); + String::from_utf8(request).expect("HTTP request should be UTF-8") + }); + let mut headers = HeaderMap::new(); + headers.insert("x-builder-test", HeaderValue::from_static("preserved")); + let client = HttpClientBuilder::new() + .default_headers(headers) + .build_with_custom_ca_fallback_using(ProxyRouting::Direct, |_| { + Err(BuildCustomCaTransportError::InvalidCaFile { + source_env: "TEST_CA_ENV", + path: PathBuf::from("invalid-test-ca.pem"), + detail: "synthetic invalid CA".to_string(), + }) + }); + + let response = client + .get(format!("http://{address}/fallback")) + .send() + .await + .expect("fallback client should send request"); + assert!(response.status().is_success()); + let request = server.join().expect("HTTP listener should finish"); + assert!( + request + .lines() + .any(|line| line.eq_ignore_ascii_case("x-builder-test: preserved")) + ); +} diff --git a/codex-rs/http-client/src/custom_ca.rs b/codex-rs/http-client/src/custom_ca.rs new file mode 100644 index 0000000000000000000000000000000000000000..6ca705d0abf9f6672932f75b95fc7e41ee9cfc31 --- /dev/null +++ b/codex-rs/http-client/src/custom_ca.rs @@ -0,0 +1,824 @@ +//! Custom CA handling shared by Codex outbound HTTP and websocket clients. +//! +//! Codex constructs outbound reqwest clients and secure websocket connections in a few crates, but +//! they all need the same trust-store policy when enterprise proxies or gateways intercept TLS. +//! This module centralizes that policy so callers can start from an ordinary +//! `reqwest::ClientBuilder` or rustls client config, layer in custom CA support, and either get +//! back a configured transport or a user-facing error that explains how to fix a misconfigured CA +//! bundle. +//! +//! The module intentionally has a narrow responsibility: +//! +//! - read CA material from `CODEX_CA_CERTIFICATE`, falling back to `SSL_CERT_FILE` +//! - normalize PEM variants that show up in real deployments, including OpenSSL-style +//! `TRUSTED CERTIFICATE` labels and bundles that also contain CRLs +//! - return user-facing errors that explain how to fix misconfigured CA files +//! +//! Its production contract is narrow: produce a transport configuration whose root store contains +//! every parseable certificate block from the configured PEM bundle, or fail early with a precise +//! error before the caller starts network traffic. +//! +//! In this module's test setup, a hermetic test is one whose result depends only on the CA file +//! and environment variables that the test chose for itself. That matters here because the normal +//! reqwest client-construction path is not hermetic enough for environment-sensitive tests: +//! +//! - on macOS seatbelt runs, `reqwest::Client::builder().build()` can panic inside +//! `system-configuration` while probing platform proxy settings, which means the process can die +//! before the custom-CA code reports success or a structured error. That matters in practice +//! because Codex itself commonly runs spawned test processes under seatbelt, so this is not just +//! a hypothetical CI edge case. +//! - child processes inherit CA-related environment variables by default, which lets developer +//! shell state or CI configuration affect a test unless the test scrubs those variables first +//! +//! The tests in this crate therefore stay split across two layers: +//! +//! - unit tests in this module cover env-selection logic without constructing a real client +//! - subprocess integration tests under `tests/` cover real client construction through +//! [`build_reqwest_client_for_subprocess_tests`], which disables reqwest proxy autodetection so +//! the tests can observe custom-CA success and failure directly, including one TLS handshake +//! through a local HTTPS server +//! - those subprocess tests also scrub inherited CA environment variables before launch so their +//! result depends only on the test fixtures and env vars set by the test itself + +use std::env; +use std::fs; +use std::io; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; + +use codex_utils_rustls_provider::ensure_rustls_crypto_provider; +use rustls::ClientConfig; +use rustls::RootCertStore; +use rustls_pki_types::CertificateDer; +use rustls_pki_types::pem::PemObject; +use rustls_pki_types::pem::SectionKind; +use rustls_pki_types::pem::{self}; +use thiserror::Error; +use tracing::info; +use tracing::warn; + +pub const CODEX_CA_CERT_ENV: &str = "CODEX_CA_CERTIFICATE"; +pub const SSL_CERT_FILE_ENV: &str = "SSL_CERT_FILE"; +const CA_CERT_HINT: &str = "If you set CODEX_CA_CERTIFICATE or SSL_CERT_FILE, ensure it points to a PEM file containing one or more CERTIFICATE blocks, or unset it to use system roots."; +type PemSection = (SectionKind, Vec); + +/// Describes why a transport using shared custom CA support could not be constructed. +/// +/// These failure modes apply to both reqwest client construction and websocket TLS +/// configuration. A build can fail because the configured CA file could not be read, could not be +/// parsed as certificates, contained certs that the target TLS stack refused to register, or +/// because the final reqwest client builder failed. Callers that do not care about the +/// distinction can rely on the `From for io::Error` conversion. +#[derive(Debug, Error)] +pub enum BuildCustomCaTransportError { + /// Reading the selected CA file from disk failed before any PEM parsing could happen. + #[error( + "Failed to read CA certificate file {} selected by {}: {source}. {hint}", + path.display(), + source_env, + hint = CA_CERT_HINT + )] + ReadCaFile { + source_env: &'static str, + path: PathBuf, + source: io::Error, + }, + + /// The selected CA file was readable, but did not produce usable certificate material. + #[error( + "Failed to load CA certificates from {} selected by {}: {detail}. {hint}", + path.display(), + source_env, + hint = CA_CERT_HINT + )] + InvalidCaFile { + source_env: &'static str, + path: PathBuf, + detail: String, + }, + + /// One parsed certificate block could not be registered with the reqwest client builder. + #[error( + "Failed to parse certificate #{certificate_index} from {} selected by {}: {source}. {hint}", + path.display(), + source_env, + hint = CA_CERT_HINT + )] + RegisterCertificate { + source_env: &'static str, + path: PathBuf, + certificate_index: usize, + source: reqwest::Error, + }, + + /// Reqwest rejected the final client configuration after a custom CA bundle was loaded. + #[error( + "Failed to build HTTP client while using CA bundle from {} ({}): {source}", + source_env, + path.display() + )] + BuildClientWithCustomCa { + source_env: &'static str, + path: PathBuf, + #[source] + source: reqwest::Error, + }, + + /// Reqwest rejected the final client configuration while using only system roots. + #[error("Failed to build HTTP client while using system root certificates: {0}")] + BuildClientWithSystemRoots(#[source] reqwest::Error), + + /// One parsed certificate block could not be registered with the websocket TLS root store. + #[error( + "Failed to register certificate #{certificate_index} from {} selected by {} in rustls root store: {source}. {hint}", + path.display(), + source_env, + hint = CA_CERT_HINT + )] + RegisterRustlsCertificate { + source_env: &'static str, + path: PathBuf, + certificate_index: usize, + source: rustls::Error, + }, +} + +impl From for io::Error { + fn from(error: BuildCustomCaTransportError) -> Self { + match error { + BuildCustomCaTransportError::ReadCaFile { ref source, .. } => { + io::Error::new(source.kind(), error) + } + BuildCustomCaTransportError::InvalidCaFile { .. } + | BuildCustomCaTransportError::RegisterCertificate { .. } + | BuildCustomCaTransportError::RegisterRustlsCertificate { .. } => { + io::Error::new(io::ErrorKind::InvalidData, error) + } + BuildCustomCaTransportError::BuildClientWithCustomCa { .. } + | BuildCustomCaTransportError::BuildClientWithSystemRoots(_) => io::Error::other(error), + } + } +} + +/// Builds a reqwest client that honors Codex custom CA environment variables. +/// +/// Callers supply the baseline builder configuration they need, and this helper layers in custom +/// CA handling before finally constructing the client. `CODEX_CA_CERTIFICATE` takes precedence +/// over `SSL_CERT_FILE`, and empty values for either are treated as unset so callers do not +/// accidentally turn `VAR=""` into a bogus path lookup. +/// +/// Callers that build a raw `reqwest::Client` directly bypass this policy entirely. That is an +/// easy mistake to make when adding a new outbound Codex HTTP path, and the resulting bug only +/// shows up in environments where a proxy or gateway requires a custom root CA. +/// +/// # Errors +/// +/// Returns a [`BuildCustomCaTransportError`] when the configured CA file is unreadable, +/// malformed, or contains a certificate block that `reqwest` cannot register as a root. +pub fn build_reqwest_client_with_custom_ca( + builder: reqwest::ClientBuilder, +) -> Result { + build_reqwest_client_with_env(&ProcessEnv, builder) +} + +/// Builds a rustls client config when a Codex custom CA bundle is configured. +/// +/// This is the websocket-facing sibling of [`build_reqwest_client_with_custom_ca`]. When +/// `CODEX_CA_CERTIFICATE` or `SSL_CERT_FILE` selects a CA bundle, the returned config starts from +/// the platform native roots and then adds the configured custom CA certificates. When no custom +/// CA env var is set, this returns `Ok(None)` so websocket callers can keep using their ordinary +/// default connector path. +/// +/// Callers that let tungstenite build its default TLS connector directly bypass this policy +/// entirely. That bug only shows up in environments where secure websocket traffic needs the same +/// enterprise root CA bundle as HTTPS traffic. +pub fn maybe_build_rustls_client_config_with_custom_ca() +-> Result>, BuildCustomCaTransportError> { + maybe_build_rustls_client_config_with_env(&ProcessEnv) +} + +/// Builds a rustls client config using native roots and any configured Codex custom CA bundle. +/// +/// Unlike [`maybe_build_rustls_client_config_with_custom_ca`], this always returns a config. Use +/// this when the caller must perform TLS itself instead of delegating default configuration to a +/// transport library. +pub fn build_rustls_client_config_with_custom_ca() +-> Result, BuildCustomCaTransportError> { + build_rustls_client_config_with_env(&ProcessEnv) +} + +/// Builds a reqwest client for spawned subprocess tests that exercise CA behavior. +/// +/// This is the test-only client-construction path used by the subprocess coverage in `tests/`. +/// The module-level docs explain the hermeticity problem in full; this helper only addresses the +/// reqwest proxy-discovery panic side of that problem by disabling proxy autodetection. The tests +/// still scrub inherited CA environment variables themselves. Normal production callers should use +/// [`build_reqwest_client_with_custom_ca`] so test-only proxy behavior does not leak into +/// ordinary client construction. +pub fn build_reqwest_client_for_subprocess_tests( + builder: reqwest::ClientBuilder, +) -> Result { + build_reqwest_client_with_env(&ProcessEnv, builder.no_proxy()) +} + +fn maybe_build_rustls_client_config_with_env( + env_source: &dyn EnvSource, +) -> Result>, BuildCustomCaTransportError> { + let Some(bundle) = env_source.configured_ca_bundle() else { + return Ok(None); + }; + + build_rustls_client_config(Some(&bundle)).map(Some) +} + +fn build_rustls_client_config_with_env( + env_source: &dyn EnvSource, +) -> Result, BuildCustomCaTransportError> { + let bundle = env_source.configured_ca_bundle(); + build_rustls_client_config(bundle.as_ref()) +} + +fn build_rustls_client_config( + bundle: Option<&ConfiguredCaBundle>, +) -> Result, BuildCustomCaTransportError> { + ensure_rustls_crypto_provider(); + + // Start from the platform roots so websocket callers keep the same baseline trust behavior + // they would get from tungstenite's default rustls connector, then layer in the Codex custom + // CA bundle on top when configured. + let mut root_store = RootCertStore::empty(); + let rustls_native_certs::CertificateResult { certs, errors, .. } = + rustls_native_certs::load_native_certs(); + if !errors.is_empty() { + warn!( + native_root_error_count = errors.len(), + "encountered errors while loading native root certificates" + ); + } + let _ = root_store.add_parsable_certificates(certs); + + if let Some(bundle) = bundle { + let certificates = bundle.load_certificates()?; + for (idx, cert) in certificates.into_iter().enumerate() { + if let Err(source) = root_store.add(cert) { + warn!( + source_env = bundle.source_env, + ca_path = %bundle.path.display(), + certificate_index = idx + 1, + error = %source, + "failed to register CA certificate in rustls root store" + ); + return Err(BuildCustomCaTransportError::RegisterRustlsCertificate { + source_env: bundle.source_env, + path: bundle.path.clone(), + certificate_index: idx + 1, + source, + }); + } + } + } + + Ok(Arc::new( + ClientConfig::builder() + .with_root_certificates(root_store) + .with_no_client_auth(), + )) +} + +/// Builds a reqwest client using an injected environment source and reqwest builder. +/// +/// This exists so tests can exercise precedence behavior deterministically without mutating the +/// real process environment. It selects the CA bundle, delegates file parsing to +/// [`ConfiguredCaBundle::load_certificates`], preserves the caller's chosen `reqwest` builder +/// configuration, forces rustls when a custom CA is configured, and finally registers each parsed +/// certificate with that builder. +fn build_reqwest_client_with_env( + env_source: &dyn EnvSource, + mut builder: reqwest::ClientBuilder, +) -> Result { + if let Some(bundle) = env_source.configured_ca_bundle() { + ensure_rustls_crypto_provider(); + info!( + source_env = bundle.source_env, + ca_path = %bundle.path.display(), + "building HTTP client with rustls backend for custom CA bundle" + ); + builder = builder.use_rustls_tls(); + + let certificates = bundle.load_certificates()?; + + for (idx, cert) in certificates.iter().enumerate() { + let certificate = match reqwest::Certificate::from_der(cert.as_ref()) { + Ok(certificate) => certificate, + Err(source) => { + warn!( + source_env = bundle.source_env, + ca_path = %bundle.path.display(), + certificate_index = idx + 1, + error = %source, + "failed to register CA certificate" + ); + return Err(BuildCustomCaTransportError::RegisterCertificate { + source_env: bundle.source_env, + path: bundle.path.clone(), + certificate_index: idx + 1, + source, + }); + } + }; + builder = builder.add_root_certificate(certificate); + } + return match builder.build() { + Ok(client) => Ok(client), + Err(source) => { + warn!( + source_env = bundle.source_env, + ca_path = %bundle.path.display(), + error = %source, + "failed to build client after loading custom CA bundle" + ); + Err(BuildCustomCaTransportError::BuildClientWithCustomCa { + source_env: bundle.source_env, + path: bundle.path.clone(), + source, + }) + } + }; + } + + info!( + codex_ca_certificate_configured = false, + ssl_cert_file_configured = false, + "using system root certificates because no CA override environment variable was selected" + ); + + match builder.build() { + Ok(client) => Ok(client), + Err(source) => { + warn!( + error = %source, + "failed to build client while using system root certificates" + ); + Err(BuildCustomCaTransportError::BuildClientWithSystemRoots( + source, + )) + } + } +} + +/// Abstracts environment access so tests can cover precedence rules without mutating process-wide +/// variables. +trait EnvSource { + /// Returns the environment variable value for `key`, if this source considers it set. + /// + /// Implementations should return `None` for absent values and may also collapse unreadable + /// process-environment states into `None`, because the custom CA logic treats both cases as + /// "no override configured". Callers build precedence and empty-string handling on top of this + /// method, so implementations should not trim or normalize the returned string. + fn var(&self, key: &str) -> Option; + + /// Returns a non-empty environment variable value interpreted as a filesystem path. + /// + /// Empty strings are treated as unset because presence here acts as a boolean "custom CA + /// override requested" signal. This keeps the precedence logic from treating `VAR=""` as an + /// attempt to open the current working directory or some other platform-specific oddity once + /// it is converted into a path. + fn non_empty_path(&self, key: &str) -> Option { + self.var(key) + .filter(|value| !value.is_empty()) + .map(PathBuf::from) + } + + /// Returns the configured CA bundle and which environment variable selected it. + /// + /// `CODEX_CA_CERTIFICATE` wins over `SSL_CERT_FILE` because it is the Codex-specific override. + /// Keeping the winning variable name with the path lets later logging explain not only which + /// file was used but also why that file was chosen. + fn configured_ca_bundle(&self) -> Option { + self.non_empty_path(CODEX_CA_CERT_ENV) + .map(|path| ConfiguredCaBundle { + source_env: CODEX_CA_CERT_ENV, + path, + }) + .or_else(|| { + self.non_empty_path(SSL_CERT_FILE_ENV) + .map(|path| ConfiguredCaBundle { + source_env: SSL_CERT_FILE_ENV, + path, + }) + }) + } +} + +/// Reads CA configuration from the real process environment. +/// +/// This is the production `EnvSource` implementation used by +/// [`build_reqwest_client_with_custom_ca`]. Tests substitute in-memory env maps so they can +/// exercise precedence and empty-value behavior without mutating process-global variables. +struct ProcessEnv; + +impl EnvSource for ProcessEnv { + fn var(&self, key: &str) -> Option { + env::var(key).ok() + } +} + +/// Identifies the CA bundle selected for a client and the policy decision that selected it. +/// +/// This is the concrete output of the environment-precedence logic. Callers use `source_env` for +/// logging and diagnostics, while `path` is the bundle that will actually be loaded. +struct ConfiguredCaBundle { + /// The environment variable that won the precedence check for this bundle. + source_env: &'static str, + /// The filesystem path that should be read as PEM certificate input. + path: PathBuf, +} + +impl ConfiguredCaBundle { + /// Loads certificates from this selected CA bundle. + /// + /// The bundle already represents the output of environment-precedence selection, so this is + /// the natural point where the file-loading phase begins. The method owns the high-level + /// success/failure logs for that phase and keeps the source env and path together for lower- + /// level parsing and error shaping. + fn load_certificates( + &self, + ) -> Result>, BuildCustomCaTransportError> { + match self.parse_certificates() { + Ok(certificates) => { + info!( + source_env = self.source_env, + ca_path = %self.path.display(), + certificate_count = certificates.len(), + "loaded certificates from custom CA bundle" + ); + Ok(certificates) + } + Err(error) => { + warn!( + source_env = self.source_env, + ca_path = %self.path.display(), + error = %error, + "failed to load custom CA bundle" + ); + Err(error) + } + } + } + + /// Loads every certificate block from a PEM file intended for Codex CA overrides. + /// + /// This accepts a few common real-world variants so Codex behaves like other CA-aware tooling: + /// leading comments are preserved, `TRUSTED CERTIFICATE` labels are normalized to standard + /// certificate labels, and embedded CRLs are ignored when they are well-formed enough for the + /// section iterator to classify them. + fn parse_certificates( + &self, + ) -> Result>, BuildCustomCaTransportError> { + let pem_data = self.read_pem_data()?; + let normalized_pem = NormalizedPem::from_pem_data(self.source_env, &self.path, &pem_data); + + let mut certificates = Vec::new(); + let mut logged_crl_presence = false; + for section_result in normalized_pem.sections() { + // Known limitation: if `rustls-pki-types` fails while parsing a malformed CRL section, + // that error is reported here before we can classify the block as ignorable. A bundle + // containing valid certificates plus a malformed `X509 CRL` therefore still fails to + // load today, even though well-formed CRLs are ignored. + let (section_kind, der) = match section_result { + Ok(section) => section, + Err(error) => return Err(self.pem_parse_error(&error)), + }; + match section_kind { + SectionKind::Certificate => { + // Standard CERTIFICATE blocks already decode to the exact DER bytes reqwest + // wants. Only OpenSSL TRUSTED CERTIFICATE blocks need trimming to drop any + // trailing X509_AUX trust metadata before registration. + let cert_der = normalized_pem.certificate_der(&der).ok_or_else(|| { + self.invalid_ca_file( + "failed to extract certificate data from TRUSTED CERTIFICATE: invalid DER length", + ) + })?; + certificates.push(CertificateDer::from(cert_der.to_vec())); + } + SectionKind::Crl if !logged_crl_presence => { + info!( + source_env = self.source_env, + ca_path = %self.path.display(), + "ignoring X509 CRL entries found in custom CA bundle" + ); + logged_crl_presence = true; + } + _ => {} + } + } + + if certificates.is_empty() { + return Err(self.pem_parse_error(&pem::Error::NoItemsFound)); + } + + Ok(certificates) + } + + /// Reads the CA bundle bytes while preserving the original filesystem error kind. + /// + /// The caller wants a user-facing error that includes the bundle path and remediation hint, but + /// higher-level surfaces still benefit from distinguishing "not found" from other I/O + /// failures. This helper keeps both pieces together. + fn read_pem_data(&self) -> Result, BuildCustomCaTransportError> { + fs::read(&self.path).map_err(|source| BuildCustomCaTransportError::ReadCaFile { + source_env: self.source_env, + path: self.path.clone(), + source, + }) + } + + /// Rewrites PEM parsing failures into user-facing configuration errors. + /// + /// The underlying parser knows whether the file was empty, malformed, or contained unsupported + /// PEM content, but callers need a message that also points them back to the relevant + /// environment variables and the expected remediation. + fn pem_parse_error(&self, error: &pem::Error) -> BuildCustomCaTransportError { + let detail = match error { + pem::Error::NoItemsFound => "no certificates found in PEM file".to_string(), + _ => format!("failed to parse PEM file: {error}"), + }; + + self.invalid_ca_file(detail) + } + + /// Creates an invalid-CA error tied to this file path. + /// + /// Most parse-time failures in this module eventually collapse to "the configured CA bundle is + /// not usable", but the detailed reason still matters for operator debugging. Centralizing that + /// formatting keeps the path and hint text consistent across the different parser branches. + fn invalid_ca_file(&self, detail: impl std::fmt::Display) -> BuildCustomCaTransportError { + BuildCustomCaTransportError::InvalidCaFile { + source_env: self.source_env, + path: self.path.clone(), + detail: detail.to_string(), + } + } +} + +/// The PEM text shape after OpenSSL compatibility normalization. +/// +/// `Standard` means the input already used ordinary PEM certificate labels. `TrustedCertificate` +/// means the input used OpenSSL's `TRUSTED CERTIFICATE` labels, so callers must also be prepared +/// to trim trailing `X509_AUX` bytes from decoded certificate sections. +enum NormalizedPem { + /// PEM contents that already used ordinary `CERTIFICATE` labels. + Standard(String), + /// PEM contents rewritten from OpenSSL `TRUSTED CERTIFICATE` labels to `CERTIFICATE`. + TrustedCertificate(String), +} + +impl NormalizedPem { + /// Normalizes PEM text from a CA bundle into the label shape this module expects. + /// + /// Codex only needs certificate DER bytes to seed `reqwest`'s root store, but operators may + /// point it at CA files that came from OpenSSL tooling rather than from a minimal certificate + /// bundle. OpenSSL's `TRUSTED CERTIFICATE` form is one such variant: it is still certificate + /// material, but it uses a different PEM label and may carry auxiliary trust metadata that + /// this crate does not consume. This constructor rewrites only the PEM labels so the mixed- + /// section parser can keep treating the file as certificate input. The rustls ecosystem does + /// not currently accept `TRUSTED CERTIFICATE` as a standard certificate label upstream, so + /// this remains a local compatibility shim rather than behavior delegated to + /// `rustls-pki-types`. + /// + /// See also: + /// - rustls/pemfile issue #52, closed as not planned, documenting that + /// `BEGIN TRUSTED CERTIFICATE` blocks are ignored upstream: + /// + /// - OpenSSL `x509 -trustout`, which emits `TRUSTED CERTIFICATE` PEM blocks: + /// + /// - OpenSSL PEM readers, which document that plain `PEM_read_bio_X509()` discards auxiliary + /// trust settings: + /// + /// - `openssl s_server`, a real OpenSSL-based server/test tool that operates in this + /// ecosystem: + /// + fn from_pem_data(source_env: &'static str, path: &Path, pem_data: &[u8]) -> Self { + let pem = String::from_utf8_lossy(pem_data); + if pem.contains("TRUSTED CERTIFICATE") { + info!( + source_env, + ca_path = %path.display(), + "normalizing OpenSSL TRUSTED CERTIFICATE labels in custom CA bundle" + ); + Self::TrustedCertificate( + pem.replace("BEGIN TRUSTED CERTIFICATE", "BEGIN CERTIFICATE") + .replace("END TRUSTED CERTIFICATE", "END CERTIFICATE"), + ) + } else { + Self::Standard(pem.into_owned()) + } + } + + /// Returns the normalized PEM contents regardless of the label shape that produced them. + fn contents(&self) -> &str { + match self { + Self::Standard(contents) | Self::TrustedCertificate(contents) => contents, + } + } + + /// Iterates over every recognized PEM section in this normalized PEM text. + /// + /// `rustls-pki-types` exposes mixed-section parsing through a `PemObject` implementation on the + /// `(SectionKind, Vec)` tuple. Keeping that type-directed API here lets callers iterate in + /// terms of normalized sections rather than trait plumbing. + fn sections(&self) -> impl Iterator> + '_ { + PemSection::pem_slice_iter(self.contents().as_bytes()) + } + + /// Returns the certificate DER bytes for one parsed PEM certificate section. + /// + /// Standard PEM certificates already decode to the exact DER bytes `reqwest` wants. OpenSSL + /// `TRUSTED CERTIFICATE` sections may append `X509_AUX` bytes after the certificate, so those + /// sections need to be trimmed down to their first DER object before registration. + fn certificate_der<'a>(&self, der: &'a [u8]) -> Option<&'a [u8]> { + match self { + Self::Standard(_) => Some(der), + Self::TrustedCertificate(_) => first_der_item(der), + } + } +} + +/// Returns the first DER-encoded ASN.1 object in `der`, ignoring any trailing OpenSSL metadata. +/// +/// A PEM `CERTIFICATE` block usually decodes to exactly one DER blob: the certificate itself. +/// OpenSSL's `TRUSTED CERTIFICATE` variant is different. It starts with that same certificate +/// blob, but may append extra `X509_AUX` bytes after it to describe OpenSSL-specific trust +/// settings. `reqwest::Certificate::from_der` only understands the certificate object, not those +/// trailing OpenSSL extensions. +/// +/// This helper therefore asks a narrower question than "is this a valid certificate?": where does +/// the first top-level DER object end? If that boundary can be found, the caller keeps only that +/// prefix and discards the trailing trust metadata. If it cannot be found, the input is treated as +/// malformed CA data. +fn first_der_item(der: &[u8]) -> Option<&[u8]> { + der_item_length(der).map(|length| &der[..length]) +} + +/// Returns the byte length of the first DER item in `der`. +/// +/// DER is a binary encoding for ASN.1 objects. Each object begins with: +/// +/// - a tag byte describing what kind of object follows +/// - one or more length bytes describing how many content bytes belong to that object +/// - the content bytes themselves +/// +/// For this module, the important fact is that a certificate is stored as one complete top-level +/// DER object. Once we know that object's declared length, we know exactly where the certificate +/// ends and where any trailing OpenSSL `X509_AUX` data begins. +/// +/// This helper intentionally parses only that outer length field. It does not validate the inner +/// certificate structure, the meaning of the tag, or every nested ASN.1 value. That narrower scope +/// is deliberate: the caller only needs a safe slice boundary for the leading certificate object +/// before handing those bytes to `reqwest`, which performs the real certificate parsing. +/// +/// The implementation supports the DER length forms needed here: +/// +/// - short form, where the length is stored directly in the second byte +/// - long form, where the second byte says how many following bytes make up the length value +/// +/// Indefinite lengths are rejected because DER does not permit them, and any declared length that +/// would run past the end of the input is treated as malformed. +fn der_item_length(der: &[u8]) -> Option { + let &length_octet = der.get(1)?; + if length_octet & 0x80 == 0 { + return Some(2 + usize::from(length_octet)).filter(|length| *length <= der.len()); + } + + let length_octets = usize::from(length_octet & 0x7f); + if length_octets == 0 { + return None; + } + + let length_start = 2usize; + let length_end = length_start.checked_add(length_octets)?; + let length_bytes = der.get(length_start..length_end)?; + let mut content_length = 0usize; + for &byte in length_bytes { + content_length = content_length + .checked_mul(256)? + .checked_add(usize::from(byte))?; + } + + length_end + .checked_add(content_length) + .filter(|length| *length <= der.len()) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::fs; + use std::path::PathBuf; + + use pretty_assertions::assert_eq; + use tempfile::TempDir; + + use super::BuildCustomCaTransportError; + use super::CODEX_CA_CERT_ENV; + use super::EnvSource; + use super::SSL_CERT_FILE_ENV; + use super::maybe_build_rustls_client_config_with_env; + + const TEST_CERT: &str = include_str!("../tests/fixtures/test-ca.pem"); + + struct MapEnv { + values: HashMap, + } + + impl EnvSource for MapEnv { + fn var(&self, key: &str) -> Option { + self.values.get(key).cloned() + } + } + + fn map_env(pairs: &[(&str, &str)]) -> MapEnv { + MapEnv { + values: pairs + .iter() + .map(|(key, value)| ((*key).to_string(), (*value).to_string())) + .collect(), + } + } + + fn write_cert_file(temp_dir: &TempDir, name: &str, contents: &str) -> PathBuf { + let path = temp_dir.path().join(name); + fs::write(&path, contents).unwrap_or_else(|error| { + panic!("write cert fixture failed for {}: {error}", path.display()) + }); + path + } + + #[test] + fn ca_path_prefers_codex_env() { + let env = map_env(&[ + (CODEX_CA_CERT_ENV, "/tmp/codex.pem"), + (SSL_CERT_FILE_ENV, "/tmp/fallback.pem"), + ]); + + assert_eq!( + env.configured_ca_bundle().map(|bundle| bundle.path), + Some(PathBuf::from("/tmp/codex.pem")) + ); + } + + #[test] + fn ca_path_falls_back_to_ssl_cert_file() { + let env = map_env(&[(SSL_CERT_FILE_ENV, "/tmp/fallback.pem")]); + + assert_eq!( + env.configured_ca_bundle().map(|bundle| bundle.path), + Some(PathBuf::from("/tmp/fallback.pem")) + ); + } + + #[test] + fn ca_path_ignores_empty_values() { + let env = map_env(&[ + (CODEX_CA_CERT_ENV, ""), + (SSL_CERT_FILE_ENV, "/tmp/fallback.pem"), + ]); + + assert_eq!( + env.configured_ca_bundle().map(|bundle| bundle.path), + Some(PathBuf::from("/tmp/fallback.pem")) + ); + } + + #[test] + fn rustls_config_uses_custom_ca_bundle_when_configured() { + let temp_dir = TempDir::new().expect("tempdir"); + let cert_path = write_cert_file(&temp_dir, "ca.pem", TEST_CERT); + let env = map_env(&[(CODEX_CA_CERT_ENV, cert_path.to_string_lossy().as_ref())]); + + let config = maybe_build_rustls_client_config_with_env(&env) + .expect("rustls config") + .expect("custom CA config should be present"); + + assert!(config.enable_sni); + } + + #[test] + fn rustls_config_reports_invalid_ca_file() { + let temp_dir = TempDir::new().expect("tempdir"); + let cert_path = write_cert_file(&temp_dir, "empty.pem", ""); + let env = map_env(&[(CODEX_CA_CERT_ENV, cert_path.to_string_lossy().as_ref())]); + + let error = maybe_build_rustls_client_config_with_env(&env).expect_err("invalid CA"); + + assert!(matches!( + error, + BuildCustomCaTransportError::InvalidCaFile { .. } + )); + } +} + +#[cfg(test)] +#[path = "custom_ca_tls_tests.rs"] +mod tls_tests; diff --git a/codex-rs/http-client/src/custom_ca_tls_tests.rs b/codex-rs/http-client/src/custom_ca_tls_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..fb3c78fe14701cb78dd85f5aab3c4b0932bf6d4e --- /dev/null +++ b/codex-rs/http-client/src/custom_ca_tls_tests.rs @@ -0,0 +1,79 @@ +//! Handshake coverage for the custom trust path retained by Windows realtime connections. + +use super::CODEX_CA_CERT_ENV; +use super::ConfiguredCaBundle; +use super::build_rustls_client_config; +use pretty_assertions::assert_eq; +use rcgen::BasicConstraints; +use rcgen::CertificateParams; +use rcgen::CertifiedIssuer; +use rcgen::IsCa; +use rcgen::KeyPair; +use std::sync::Arc; + +#[test] +fn custom_intermediate_trust_preserves_hostname_validation() { + let mut params = CertificateParams::default(); + params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + let root = CertifiedIssuer::self_signed(params.clone(), KeyPair::generate().unwrap()).unwrap(); + let intermediate = + CertifiedIssuer::signed_by(params, KeyPair::generate().unwrap(), &root).unwrap(); + let leaf_key = KeyPair::generate().unwrap(); + let leaf = CertificateParams::new(vec!["localhost".to_string()]) + .unwrap() + .signed_by(&leaf_key, &intermediate) + .unwrap(); + let temp = tempfile::TempDir::new().unwrap(); + let path = temp.path().join("intermediate.pem"); + std::fs::write(&path, intermediate.pem()).unwrap(); + let config = build_rustls_client_config(Some(&ConfiguredCaBundle { + source_env: CODEX_CA_CERT_ENV, + path, + })) + .unwrap(); + let server = Arc::new( + rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(vec![leaf.der().clone()], leaf_key.into()) + .unwrap(), + ); + + assert_eq!( + handshake(config.clone(), server.clone(), "localhost"), + Ok(()) + ); + let error = handshake(config, server, "wrong.example").unwrap_err(); + assert!( + matches!( + error, + rustls::Error::InvalidCertificate( + rustls::CertificateError::NotValidForName + | rustls::CertificateError::NotValidForNameContext { .. } + ) + ), + "{error:?}" + ); +} + +fn handshake( + config: Arc, + server: Arc, + hostname: &'static str, +) -> Result<(), rustls::Error> { + let mut client = rustls::ClientConnection::new(config, hostname.try_into().unwrap()).unwrap(); + let mut server = rustls::ServerConnection::new(server).unwrap(); + for _ in 0..10 { + let mut bytes = Vec::new(); + client.write_tls(&mut bytes).unwrap(); + server.read_tls(&mut bytes.as_slice()).unwrap(); + server.process_new_packets()?; + bytes.clear(); + server.write_tls(&mut bytes).unwrap(); + client.read_tls(&mut bytes.as_slice()).unwrap(); + client.process_new_packets()?; + if !client.is_handshaking() { + return Ok(()); + } + } + panic!("handshake did not complete"); +} diff --git a/codex-rs/http-client/src/error.rs b/codex-rs/http-client/src/error.rs new file mode 100644 index 0000000000000000000000000000000000000000..43ace4f7a78ecab1153e548a1205b77b93d60349 --- /dev/null +++ b/codex-rs/http-client/src/error.rs @@ -0,0 +1,37 @@ +//! Errors returned by the shared Codex HTTP transport. + +use crate::client::HttpError; +use http::HeaderMap; +use http::StatusCode; +use thiserror::Error; + +#[derive(Debug, Error)] +pub enum TransportError { + #[error("http {status}: {body:?}")] + Http { + status: StatusCode, + url: Option, + headers: Option, + body: Option, + }, + #[error("retry limit reached")] + RetryLimit, + #[error("timeout")] + Timeout, + #[error("connection failed: {0}")] + Connection(#[source] HttpError), + #[error("network error: {0}")] + Network(String), + #[error("request build error: {0}")] + Build(String), + #[error("response body exceeds the {max_bytes} byte limit")] + ResponseTooLarge { max_bytes: usize }, +} + +#[derive(Debug, Error)] +pub enum StreamError { + #[error("stream failed: {0}")] + Stream(String), + #[error("timeout")] + Timeout, +} diff --git a/codex-rs/http-client/src/lib.rs b/codex-rs/http-client/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..f9cdb1fb2066c479764d919c1381d15ab206bd93 --- /dev/null +++ b/codex-rs/http-client/src/lib.rs @@ -0,0 +1,72 @@ +mod chatgpt_cloudflare_cookies; +mod chatgpt_hosts; +mod client; +mod client_builder; +mod custom_ca; +mod error; +mod network_policy; +mod outbound_proxy; +mod request; +mod route_aware_client_pool; +mod route_aware_redirect; +mod tls_backend_fallback; +mod transport; + +pub use crate::chatgpt_cloudflare_cookies::with_chatgpt_cloudflare_cookie_store; +pub use crate::chatgpt_hosts::is_allowed_chatgpt_host; +pub use crate::client::HttpClient; +pub use crate::client::HttpError; +pub use crate::client::HttpResponse; +pub use crate::client::RequestBuilder; +pub use crate::client_builder::HttpClientBuilder; +pub use crate::custom_ca::BuildCustomCaTransportError; +/// Test-only subprocess hook for custom CA coverage. +/// +/// This stays public only so the `custom_ca_probe` binary target can reuse the shared helper. It +/// is hidden from normal docs because ordinary callers should use +/// [`build_reqwest_client_with_custom_ca`] instead. +#[doc(hidden)] +pub use crate::custom_ca::build_reqwest_client_for_subprocess_tests; +pub use crate::custom_ca::build_reqwest_client_with_custom_ca; +pub use crate::custom_ca::build_rustls_client_config_with_custom_ca; +pub use crate::custom_ca::maybe_build_rustls_client_config_with_custom_ca; +pub use crate::error::StreamError; +pub use crate::error::TransportError; +pub use crate::network_policy::DestinationPolicy; +pub use crate::network_policy::NetworkPermit; +pub use crate::network_policy::NetworkPolicy; +pub use crate::network_policy::NetworkPolicyController; +pub use crate::network_policy::NetworkPolicyDenied; +pub use crate::network_policy::NetworkPolicyRevision; +pub use crate::outbound_proxy::BuildRouteAwareHttpClientError; +pub use crate::outbound_proxy::ClientRouteClass; +pub use crate::outbound_proxy::HttpClientFactory; +#[cfg(target_os = "macos")] +pub use crate::outbound_proxy::MacosSystemProxyConfiguration; +pub use crate::outbound_proxy::OutboundProxyPolicy; +pub use crate::outbound_proxy::OutboundProxyRoute; +pub use crate::outbound_proxy::RouteFailureClass; +#[doc(hidden)] +pub use crate::outbound_proxy::cache_system_proxy_route_for_test; +#[cfg(target_os = "macos")] +pub use crate::outbound_proxy::macos_system_proxy_configuration; +pub use crate::request::EncodedJsonBody; +pub use crate::request::PreparedRequestBody; +pub use crate::request::Request; +pub use crate::request::RequestBody; +pub use crate::request::RequestCompression; +pub use crate::request::Response; +pub use crate::route_aware_client_pool::RouteAwareClientPool; +pub use crate::route_aware_client_pool::RouteAwareClientPoolError; +pub use crate::route_aware_client_pool::RouteAwareRequestBuilder; +pub use crate::route_aware_client_pool::RouteAwareRequestError; +pub use crate::transport::ByteStream; +pub use crate::transport::HttpTransport; +pub use crate::transport::ReqwestTransport; +pub use crate::transport::StreamResponse; + +#[cfg(windows)] +mod windows_tls; + +#[cfg(windows)] +pub use crate::windows_tls::build_windows_platform_tls_config; diff --git a/codex-rs/http-client/src/network_policy.rs b/codex-rs/http-client/src/network_policy.rs new file mode 100644 index 0000000000000000000000000000000000000000..fa92421d60d60d1cf0415f25a404f33001177bb5 --- /dev/null +++ b/codex-rs/http-client/src/network_policy.rs @@ -0,0 +1,353 @@ +//! Application destination checks and revocation of outstanding network operations. +//! +//! The configuration owner publishes an already-composed policy. Transports retain only its +//! read side and a permit for each request; invalidation cannot leave an old request authorized. + +use std::collections::BTreeSet; +use std::fmt; +use std::future::Future; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::MutexGuard; +use std::sync::Weak; + +use reqwest::Url; +use tokio::sync::watch; + +/// Effective application destinations, after the requirements owner has applied precedence. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum DestinationPolicy { + Unrestricted, + /// Only HTTPS and WSS requests to these exact, normalized hosts are permitted. + Restricted { + allowed_hosts: BTreeSet, + }, +} + +impl DestinationPolicy { + fn allows(&self, url: &Url) -> bool { + match self { + Self::Unrestricted => true, + Self::Restricted { allowed_hosts } => { + matches!(url.scheme(), "https" | "wss") + && url.host_str().is_some_and(|host| { + allowed_hosts.contains(host.strip_suffix('.').unwrap_or(host)) + }) + } + } + } +} + +/// A deterministic denial, distinct from retryable network failures. +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum NetworkPolicyDenied { + #[error("application network policy is unavailable")] + Unavailable, + #[error("destination denied by application network policy")] + Destination, + #[error("application network permission was revoked")] + Revoked, + #[error("this SDK transport is disabled by application network restrictions")] + UnsupportedTransport, +} + +/// Read access to application policy. Clones observe the same updates. +#[derive(Clone, Default)] +pub struct NetworkPolicy { + state: Option>>, + endpoints: Option>>, + account: Option, +} + +impl fmt::Debug for NetworkPolicy { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("NetworkPolicy") + .field("managed", &self.state.is_some()) + .finish_non_exhaustive() + } +} + +impl PartialEq for NetworkPolicy { + fn eq(&self, other: &Self) -> bool { + self.endpoints == other.endpoints + && self.account == other.account + && match (&self.state, &other.state) { + (Some(left), Some(right)) => Arc::ptr_eq(left, right), + (None, None) => true, + (Some(_), None) | (None, Some(_)) => false, + } + } +} + +impl Eq for NetworkPolicy {} + +struct State { + policy: Option, + revision: u64, + account: u64, + permits: Vec>, + changes: watch::Sender<()>, +} + +struct Permit { + destination: PermitDestination, + revoked: watch::Sender, +} + +enum PermitDestination { + Url(Url), + UnrestrictedSdk, +} + +impl PermitDestination { + fn allowed_by(&self, policy: &DestinationPolicy) -> bool { + match self { + Self::Url(url) => policy.allows(url), + Self::UnrestrictedSdk => matches!(policy, DestinationPolicy::Unrestricted), + } + } +} + +/// Publication access retained by the account/configuration owner. +#[derive(Clone)] +pub struct NetworkPolicyController { + state: Arc>, +} + +/// Identifies the account/configuration generation that a load is allowed to publish into. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct NetworkPolicyRevision(u64); + +impl Default for NetworkPolicyController { + fn default() -> Self { + Self { + state: Arc::new(Mutex::new(State { + policy: None, + revision: 0, + account: 0, + permits: Vec::new(), + changes: watch::channel(()).0, + })), + } + } +} + +fn lock_state(state: &Mutex) -> MutexGuard<'_, State> { + state + .lock() + .unwrap_or_else(|_| panic!("application network policy lock poisoned")) +} + +impl NetworkPolicyController { + pub fn policy(&self) -> NetworkPolicy { + NetworkPolicy { + state: Some(self.state.clone()), + endpoints: None, + account: None, + } + } + + /// Fails a requirements load without changing the current account identity. + pub fn unavailable(&self, revision: NetworkPolicyRevision) { + let mut state = lock_state(&self.state); + if state.revision != revision.0 { + return; + } + if state.policy.take().is_some() { + state.changes.send_replace(()); + } + for permit in state + .permits + .drain(..) + .filter_map(|permit| permit.upgrade()) + { + permit.revoked.send_replace(/*value*/ true); + } + } + + /// Publishes a successful load only if its account/configuration generation is still current. + pub fn publish(&self, revision: NetworkPolicyRevision, policy: DestinationPolicy) -> bool { + let mut state = lock_state(&self.state); + if state.revision != revision.0 { + return false; + } + state.permits.retain(|permit| { + let Some(permit) = permit.upgrade() else { + return false; + }; + if !permit.destination.allowed_by(&policy) { + permit.revoked.send_replace(/*value*/ true); + } + true + }); + if state.policy.as_ref() != Some(&policy) { + state.policy = Some(policy); + state.changes.send_replace(()); + } + true + } +} + +impl NetworkPolicy { + pub const fn unmanaged() -> Self { + Self { + state: None, + endpoints: None, + account: None, + } + } + + /// Binds a content client to the current account so retained credentials cannot + /// be used after another workspace's policy has been installed. + pub fn for_current_account(mut self) -> Self { + self.account = self.state.as_ref().map(|state| lock_state(state).account); + self + } + + /// Whether this transport participates in the application policy lifecycle. + pub fn is_managed(&self) -> bool { + self.state.is_some() + } + + /// Watches effective changes; unmanaged callers have no application policy owner. + pub fn changes(&self) -> Option> { + self.state + .as_ref() + .map(|state| lock_state(state).changes.subscribe()) + } + + pub fn revision(&self) -> NetworkPolicyRevision { + NetworkPolicyRevision( + self.state + .as_ref() + .map_or(0, |state| lock_state(state).revision), + ) + } + + /// Revokes outstanding operations before an account change or a failed requirements load. + /// This read-side operation can only remove access; publication stays with the controller. + pub fn invalidate(&self) { + if let Some(state) = &self.state { + let mut state = lock_state(state); + state.revision += 1; + state.account += 1; + state.policy = None; + state.changes.send_replace(()); + for permit in state + .permits + .drain(..) + .filter_map(|permit| permit.upgrade()) + { + permit.revoked.send_replace(/*value*/ true); + } + } + } + + /// Checks a destination before proxy resolution, DNS, or transport work starts. + pub fn acquire(&self, url: &Url) -> Result { + if self + .endpoints + .as_ref() + .is_some_and(|endpoints| !endpoints.contains(url)) + { + return Err(NetworkPolicyDenied::Destination); + } + self.acquire_destination(PermitDestination::Url(url.clone())) + } + + /// Disables SDKs without destination enforcement whenever restrictions apply. + /// The returned permit also cancels work if an unrestricted policy changes. + pub fn acquire_for_unsupported_sdk(&self) -> Result { + if self.endpoints.is_some() { + return Err(NetworkPolicyDenied::UnsupportedTransport); + } + self.acquire_destination(PermitDestination::UnrestrictedSdk) + } + + fn acquire_destination( + &self, + destination: PermitDestination, + ) -> Result { + let Some(state) = &self.state else { + return Ok(NetworkPermit { inner: None }); + }; + let mut state = lock_state(state); + if self.account.is_some_and(|account| account != state.account) { + return Err(NetworkPolicyDenied::Revoked); + } + let policy = state + .policy + .as_ref() + .ok_or(NetworkPolicyDenied::Unavailable)?; + if !destination.allowed_by(policy) { + return Err(match destination { + PermitDestination::Url(_) => NetworkPolicyDenied::Destination, + PermitDestination::UnrestrictedSdk => NetworkPolicyDenied::UnsupportedTransport, + }); + } + let (revoked, _) = watch::channel(/*init*/ false); + let permit = Arc::new(Permit { + destination, + revoked, + }); + state.permits.retain(|permit| permit.strong_count() != 0); + state.permits.push(Arc::downgrade(&permit)); + Ok(NetworkPermit { + inner: Some(permit), + }) + } +} + +/// Authorization retained for the entire operation, including response or WebSocket streaming. +#[derive(Clone)] +pub struct NetworkPermit { + inner: Option>, +} + +impl fmt::Debug for NetworkPermit { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("NetworkPermit").finish_non_exhaustive() + } +} + +impl NetworkPermit { + pub fn check(&self) -> Result<(), NetworkPolicyDenied> { + if self + .inner + .as_ref() + .is_some_and(|permit| *permit.revoked.borrow()) + { + Err(NetworkPolicyDenied::Revoked) + } else { + Ok(()) + } + } + + pub async fn revoked(&self) { + match &self.inner { + Some(permit) => { + let _ = permit + .revoked + .subscribe() + .wait_for(|revoked| *revoked) + .await; + } + None => std::future::pending().await, + } + } + + pub async fn run( + &self, + operation: impl Future, + ) -> Result { + tokio::select! { + biased; + _ = self.revoked() => Err(NetworkPolicyDenied::Revoked), + result = operation => Ok(result), + } + } +} + +#[cfg(test)] +#[path = "network_policy_tests.rs"] +mod tests; diff --git a/codex-rs/http-client/src/network_policy_tests.rs b/codex-rs/http-client/src/network_policy_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..dd9787136501deedecb64fd88697d8ad00caca80 --- /dev/null +++ b/codex-rs/http-client/src/network_policy_tests.rs @@ -0,0 +1,110 @@ +//! Destination matching, bootstrap scope, and account revocation behavior. + +use std::collections::BTreeSet; + +use pretty_assertions::assert_eq; +use reqwest::Url; + +use super::*; + +#[test] +fn restricted_policy_matches_exact_secure_hosts() { + let controller = NetworkPolicyController::default(); + let policy = controller.policy(); + assert!(controller.publish( + policy.revision(), + DestinationPolicy::Restricted { + allowed_hosts: BTreeSet::from(["example.com".to_string()]), + } + )); + for (url, expected) in [ + ("https://EXAMPLE.com/path", true), + ("https://example.com.:8443/path", true), + ("wss://example.com/path", true), + ("http://example.com/path", false), + ("ws://example.com/path", false), + ("https://sub.example.com/path", false), + ("https://example.com.evil.test/path", false), + ("https://example.com@evil.test/path", false), + ("https://127.0.0.1/path", false), + ] { + assert_eq!( + policy.acquire(&Url::parse(url).unwrap()).is_ok(), + expected, + "{url}" + ); + } +} + +#[tokio::test] +async fn invalidation_revokes_existing_permits_and_rejects_stale_publication() { + let controller = NetworkPolicyController::default(); + let policy = controller.policy(); + let url = Url::parse("https://example.com").unwrap(); + assert_eq!( + policy.acquire(&url).unwrap_err(), + NetworkPolicyDenied::Unavailable + ); + let previous = policy.revision(); + assert!(controller.publish(previous, DestinationPolicy::Unrestricted)); + let permit = policy.acquire(&url).unwrap(); + policy.invalidate(); + assert!(!controller.publish(previous, DestinationPolicy::Unrestricted)); + assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); + assert_eq!( + permit.run(async { "should not run" }).await, + Err(NetworkPolicyDenied::Revoked) + ); + assert!(policy.acquire(&url).is_ok()); +} + +#[tokio::test] +async fn load_failure_recovers_current_account_but_cannot_revive_previous_account() { + let controller = NetworkPolicyController::default(); + let policy = controller.policy(); + let account = policy.clone().for_current_account(); + let url = Url::parse("https://example.com").unwrap(); + let mut changes = policy.changes().unwrap(); + controller.publish(policy.revision(), DestinationPolicy::Unrestricted); + changes.changed().await.unwrap(); + controller.publish(policy.revision(), DestinationPolicy::Unrestricted); + assert!(!changes.has_changed().unwrap()); + let old_request = account.acquire(&url).unwrap(); + controller.unavailable(policy.revision()); + assert_eq!( + account.acquire(&url).unwrap_err(), + NetworkPolicyDenied::Unavailable + ); + controller.publish(policy.revision(), DestinationPolicy::Unrestricted); + assert!(account.acquire(&url).is_ok()); + assert_eq!(old_request.check(), Err(NetworkPolicyDenied::Revoked)); + policy.invalidate(); + controller.publish(policy.revision(), DestinationPolicy::Unrestricted); + assert_eq!( + account.acquire(&url).unwrap_err(), + NetworkPolicyDenied::Revoked + ); + assert!(policy.for_current_account().acquire(&url).is_ok()); +} + +#[tokio::test] +async fn unsupported_sdk_work_is_denied_and_running_work_is_cancelled() { + let controller = NetworkPolicyController::default(); + let policy = controller.policy(); + controller.publish(policy.revision(), DestinationPolicy::Unrestricted); + let permit = policy.acquire_for_unsupported_sdk().unwrap(); + controller.publish( + policy.revision(), + DestinationPolicy::Restricted { + allowed_hosts: BTreeSet::from(["example.com".to_string()]), + }, + ); + assert_eq!( + policy.acquire_for_unsupported_sdk().unwrap_err(), + NetworkPolicyDenied::UnsupportedTransport + ); + assert_eq!( + permit.run(std::future::pending::<()>()).await, + Err(NetworkPolicyDenied::Revoked) + ); +} diff --git a/codex-rs/http-client/src/outbound_proxy.rs b/codex-rs/http-client/src/outbound_proxy.rs new file mode 100644 index 0000000000000000000000000000000000000000..64713478ec079cde90997f97209c87808a3b7e8b --- /dev/null +++ b/codex-rs/http-client/src/outbound_proxy.rs @@ -0,0 +1,896 @@ +//! Conservative outbound proxy selection for resolver-aware HTTP clients. +//! +//! When enabled, platform system discovery is tried first, explicit environment +//! proxies are the fallback, and the final fallback is a direct connection. +//! When disabled, callers retain the existing reqwest builder behavior. + +use std::borrow::Cow; +use std::collections::HashMap; +use std::fmt; +use std::io; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::OnceLock; +use std::time::Duration; +use std::time::Instant; +#[cfg(any(target_os = "windows", target_os = "macos"))] +use tokio::sync::Semaphore; + +use http::HeaderValue; + +use crate::NetworkPolicy; +use crate::chatgpt_cloudflare_cookies::ChatGptCookieStore; +use crate::custom_ca::BuildCustomCaTransportError; +use crate::custom_ca::build_reqwest_client_with_custom_ca; +use sha2::Digest; +use sha2::Sha256; +use thiserror::Error; + +const SYSTEM_PROXY_SUCCESS_CACHE_TTL: Duration = Duration::from_secs(60); +const SYSTEM_PROXY_UNAVAILABLE_CACHE_TTL: Duration = Duration::from_secs(5); +const SYSTEM_PROXY_CACHE_MAX_ENTRIES: usize = 256; +#[cfg(any(target_os = "windows", target_os = "macos"))] +static ASYNC_SYSTEM_PROXY_RESOLUTION_PERMIT: Semaphore = Semaphore::const_new(1); + +#[cfg(target_os = "macos")] +mod macos; +#[cfg(target_os = "windows")] +mod windows; + +/// Coarse semantic bucket for the HTTP or WebSocket client being constructed. +/// +/// This is not the selected proxy route or a concrete endpoint. It labels the +/// product path that owns the client so proxy-resolution diagnostics can +/// distinguish auth, API, WebSocket, and miscellaneous traffic without exposing +/// endpoint details. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ClientRouteClass { + /// Login, token refresh/revoke, PAT, and agent identity auth traffic. + Auth, + /// First-party API traffic that is not part of the auth flow. + Api, + /// WebSocket traffic. + WebSocket, + /// Call sites without a more specific route class. + Other, +} + +impl fmt::Display for ClientRouteClass { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::Auth => "auth", + Self::Api => "api", + Self::WebSocket => "wss", + Self::Other => "other", + }) + } +} + +/// Coarse failure class for route selection errors. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RouteFailureClass { + ProxyResolutionUnavailable, + ConnectTimeout, + ProxyAuthenticationRequired, + TlsError, + InvalidProxyConfig, + UnsupportedProxyScheme, + ResolverError, +} + +impl fmt::Display for RouteFailureClass { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::ProxyResolutionUnavailable => "proxy_resolution_unavailable", + Self::ConnectTimeout => "connect_timeout", + Self::ProxyAuthenticationRequired => "proxy_407", + Self::TlsError => "tls_error", + Self::InvalidProxyConfig => "invalid_proxy_config", + Self::UnsupportedProxyScheme => "unsupported_proxy_scheme", + Self::ResolverError => "resolver_error", + }) + } +} + +/// Resolved outbound proxy behavior for HTTP clients. +/// +/// Callers must choose a policy explicitly so omitting feature resolution cannot silently select +/// legacy behavior. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OutboundProxyPolicy { + /// Preserve reqwest's built-in proxy behavior. + ReqwestDefault, + /// Resolve system/PAC/WPAD settings, then environment settings, then direct routing. + RespectSystemProxy, +} + +/// Privacy-safe macOS system proxy configuration for one outbound destination. +#[cfg(target_os = "macos")] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MacosSystemProxyConfiguration { + /// A PAC script or automatic proxy-discovery route applies to the destination. + Automatic, + /// An explicitly configured HTTP or HTTPS proxy applies. + Manual, + /// macOS selected a direct connection for the destination. + Direct, + /// The destination or system proxy settings could not be inspected. + Unavailable, +} + +/// Inspects macOS proxy configuration without executing PAC scripts or exposing proxy URLs. +#[cfg(target_os = "macos")] +pub fn macos_system_proxy_configuration(request_url: &str) -> MacosSystemProxyConfiguration { + macos::configuration(request_url) +} + +/// Resolved proxy route for a concrete outbound destination. +/// +/// `TransportDefault` preserves the underlying transport behavior only when system-proxy support +/// is disabled. When system resolution is enabled, environment and direct fallbacks are resolved +/// explicitly so the transport cannot repeat system discovery. Proxy URLs and no-proxy settings +/// are intentionally redacted from `Debug` output because they may contain credentials or private +/// hostnames. +#[derive(Clone, Hash, PartialEq, Eq)] +pub enum OutboundProxyRoute { + /// Preserve the underlying transport's existing proxy behavior. + TransportDefault, + /// Connect directly and bypass transport-level proxy discovery. + Direct, + /// Connect through the selected proxy URL. + Proxy { + url: String, + no_proxy: Option, + }, +} + +impl fmt::Debug for OutboundProxyRoute { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::TransportDefault => f.write_str("TransportDefault"), + Self::Direct => f.write_str("Direct"), + Self::Proxy { .. } => f + .debug_struct("Proxy") + .field("url", &"") + .field("no_proxy", &"") + .finish(), + } + } +} + +/// Builds route-specific HTTP clients using one resolved outbound proxy policy. +/// +/// Construct this once from the effective application configuration and carry it with the +/// session or component that owns outbound requests. Individual request paths should supply only +/// their destination and route class rather than resolving feature state themselves. +#[derive(Clone)] +pub struct HttpClientFactory { + outbound_proxy_policy: OutboundProxyPolicy, + chatgpt_cookie_store: Option>, + network_policy: NetworkPolicy, +} + +impl PartialEq for HttpClientFactory { + fn eq(&self, other: &Self) -> bool { + self.outbound_proxy_policy == other.outbound_proxy_policy + && self.network_policy == other.network_policy + && self + .chatgpt_cookie_store + .as_ref() + .map(|store| store.configured_cookies()) + == other + .chatgpt_cookie_store + .as_ref() + .map(|store| store.configured_cookies()) + } +} + +impl Eq for HttpClientFactory {} + +impl fmt::Debug for HttpClientFactory { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("HttpClientFactory") + .field("outbound_proxy_policy", &self.outbound_proxy_policy) + .field("network_policy", &self.network_policy) + .finish() + } +} + +impl HttpClientFactory { + /// Creates a factory from the outbound proxy policy resolved by the application. + pub const fn new(outbound_proxy_policy: OutboundProxyPolicy) -> Self { + Self { + outbound_proxy_policy, + chatgpt_cookie_store: None, + network_policy: NetworkPolicy::unmanaged(), + } + } + + /// Carries the account/configuration owner's application policy into every transport. + pub fn with_network_policy(mut self, policy: NetworkPolicy) -> Self { + self.network_policy = policy; + self + } + + pub fn network_policy(&self) -> &NetworkPolicy { + &self.network_policy + } + + /// Adds process-scoped cookies to requests made by ChatGPT cookie-store clients. + pub fn with_chatgpt_cookies(mut self, cookies: impl IntoIterator) -> Self { + let cookies = cookies.into_iter().collect::>(); + self.chatgpt_cookie_store = + (!cookies.is_empty()).then(|| Arc::new(ChatGptCookieStore::new(cookies))); + self + } + + /// Returns whether ChatGPT cookie-store clients have additional configured cookies. + pub fn has_chatgpt_cookies(&self) -> bool { + self.chatgpt_cookie_store.is_some() + } + + pub(crate) fn chatgpt_cookie_store(&self) -> Option> { + self.chatgpt_cookie_store.clone() + } + + /// Returns the outbound proxy policy used for clients built by this factory. + pub const fn outbound_proxy_policy(&self) -> OutboundProxyPolicy { + self.outbound_proxy_policy + } + + /// Resolves the proxy route for a concrete destination. + /// + /// WebSocket schemes are resolved through their HTTP equivalents so platform PAC and system + /// proxy APIs apply the same policy to `ws`/`wss` and `http`/`https` destinations. When system + /// resolution is unavailable, explicit environment settings are resolved before falling back + /// to a direct route. + pub fn resolve_proxy_route(&self, request_url: &str) -> OutboundProxyRoute { + resolve_proxy_route( + &ProcessEnv, + request_url, + self.outbound_proxy_policy, + resolve_system_proxy, + ) + } + + /// Resolves the proxy route for a concrete destination without blocking a Tokio worker. + pub async fn resolve_proxy_route_async( + &self, + request_url: String, + ) -> io::Result { + if matches!( + self.outbound_proxy_policy, + OutboundProxyPolicy::ReqwestDefault + ) { + return Ok(OutboundProxyRoute::TransportDefault); + } + + if let Some(route) = self.cached_proxy_route(&request_url) { + return Ok(route); + } + + #[cfg(not(any(target_os = "windows", target_os = "macos")))] + return Ok(self.resolve_proxy_route(&request_url)); + + #[cfg(any(target_os = "windows", target_os = "macos"))] + { + let permit = ASYNC_SYSTEM_PROXY_RESOLUTION_PERMIT + .acquire() + .await + .map_err(io::Error::other)?; + let factory = self.clone(); + tokio::task::spawn_blocking(move || { + // Keep the permit with the blocking task: cancelling the caller must not allow a + // second PAC/WinHTTP lookup to start while this one is still running. + let _permit = permit; + factory.resolve_proxy_route(&request_url) + }) + .await + .map_err(io::Error::other) + } + } + + fn cached_proxy_route(&self, request_url: &str) -> Option { + let env_proxy_kind = EnvProxyKind::from_request_url(request_url); + let request_url = proxy_resolution_url(request_url); + if RequestOrigin::parse(&request_url).is_none() { + return Some(OutboundProxyRoute::Direct); + } + cached_system_proxy_decision(&request_url) + .map(|decision| route_from_system_decision(&ProcessEnv, env_proxy_kind, decision)) + } + + /// Builds a reqwest client for a concrete outbound route. + pub fn build_reqwest_client( + &self, + builder: reqwest::ClientBuilder, + request_url: &str, + route_class: ClientRouteClass, + ) -> Result { + build_reqwest_client_for_route( + builder, + request_url, + route_class, + self.outbound_proxy_policy, + ) + } + + pub(crate) fn build_reqwest_client_for_resolved_route( + &self, + builder: reqwest::ClientBuilder, + route_class: ClientRouteClass, + route: &OutboundProxyRoute, + ) -> Result { + let builder = configure_builder_for_resolved_route(builder, route_class, route)?; + build_reqwest_client_with_custom_ca(builder).map_err(Into::into) + } +} + +fn resolve_proxy_route( + env: &dyn EnvSource, + request_url: &str, + outbound_proxy_policy: OutboundProxyPolicy, + resolve_system_proxy: impl FnOnce(&str, &RequestOrigin) -> SystemProxyDecision, +) -> OutboundProxyRoute { + if matches!(outbound_proxy_policy, OutboundProxyPolicy::ReqwestDefault) { + return OutboundProxyRoute::TransportDefault; + } + + let env_proxy_kind = EnvProxyKind::from_request_url(request_url); + let request_url = proxy_resolution_url(request_url); + let Some(origin) = RequestOrigin::parse(&request_url) else { + return OutboundProxyRoute::Direct; + }; + + route_from_system_decision( + env, + env_proxy_kind, + resolve_system_proxy(&request_url, &origin), + ) +} + +fn route_from_system_decision( + env: &dyn EnvSource, + env_proxy_kind: EnvProxyKind, + decision: SystemProxyDecision, +) -> OutboundProxyRoute { + match decision { + SystemProxyDecision::Direct => OutboundProxyRoute::Direct, + SystemProxyDecision::Proxy { url } => OutboundProxyRoute::Proxy { + url, + no_proxy: None, + }, + SystemProxyDecision::Unavailable { .. } => resolve_env_proxy_route(env, env_proxy_kind), + } +} + +fn resolve_env_proxy_route( + env: &dyn EnvSource, + env_proxy_kind: EnvProxyKind, +) -> OutboundProxyRoute { + let proxy_url = match env_proxy_kind { + EnvProxyKind::Https => { + proxy_env_value(env, "HTTPS_PROXY").or_else(|| proxy_env_value(env, "ALL_PROXY")) + } + EnvProxyKind::SecureWebSocket => proxy_env_value(env, "HTTPS_PROXY") + .or_else(|| proxy_env_value(env, "HTTP_PROXY")) + .or_else(|| proxy_env_value(env, "ALL_PROXY")), + EnvProxyKind::Http => { + proxy_env_value(env, "HTTP_PROXY").or_else(|| proxy_env_value(env, "ALL_PROXY")) + } + EnvProxyKind::Other => proxy_env_value(env, "ALL_PROXY"), + }; + match proxy_url { + Some(url) => OutboundProxyRoute::Proxy { + url, + no_proxy: proxy_env_value(env, "NO_PROXY"), + }, + None => OutboundProxyRoute::Direct, + } +} + +#[derive(Clone, Copy)] +enum EnvProxyKind { + Http, + Https, + SecureWebSocket, + Other, +} + +impl EnvProxyKind { + fn from_request_url(request_url: &str) -> Self { + let scheme = request_url + .parse::() + .ok() + .and_then(|uri| uri.scheme_str().map(str::to_ascii_lowercase)); + match scheme.as_deref() { + Some("http" | "ws") => Self::Http, + Some("https") => Self::Https, + Some("wss") => Self::SecureWebSocket, + Some(_) | None => Self::Other, + } + } +} + +fn proxy_resolution_url(request_url: &str) -> Cow<'_, str> { + if let Some(suffix) = request_url.strip_prefix("wss://") { + Cow::Owned(format!("https://{suffix}")) + } else if let Some(suffix) = request_url.strip_prefix("ws://") { + Cow::Owned(format!("http://{suffix}")) + } else { + Cow::Borrowed(request_url) + } +} + +/// Error while building a resolver-aware reqwest client. +#[derive(Debug, Error)] +pub enum BuildRouteAwareHttpClientError { + #[error(transparent)] + CustomCa(#[from] BuildCustomCaTransportError), + + #[error("Failed to configure outbound proxy selected for {route_class}")] + InvalidProxyConfig { route_class: ClientRouteClass }, +} + +impl From for io::Error { + fn from(error: BuildRouteAwareHttpClientError) -> Self { + match error { + BuildRouteAwareHttpClientError::CustomCa(error) => error.into(), + BuildRouteAwareHttpClientError::InvalidProxyConfig { .. } => io::Error::other(error), + } + } +} + +/// Builds a reqwest client with conservative route selection and shared CA handling. +/// +/// Unavailable platform resolution falls back to environment proxies and then direct. Errors after +/// a route is selected are returned without trying another route. Ordered PAC candidates are +/// currently collapsed to one route on both Windows and macOS; later proxy or `DIRECT` candidates +/// are not retried after a connection failure. +fn build_reqwest_client_for_route( + builder: reqwest::ClientBuilder, + request_url: &str, + route_class: ClientRouteClass, + outbound_proxy_policy: OutboundProxyPolicy, +) -> Result { + let builder = configure_proxy_for_route( + &ProcessEnv, + builder, + request_url, + route_class, + outbound_proxy_policy, + resolve_system_proxy, + )?; + build_reqwest_client_with_custom_ca(builder).map_err(Into::into) +} + +fn configure_proxy_for_route( + env: &dyn EnvSource, + builder: reqwest::ClientBuilder, + request_url: &str, + route_class: ClientRouteClass, + outbound_proxy_policy: OutboundProxyPolicy, + resolve_system_proxy: impl FnOnce(&str, &RequestOrigin) -> SystemProxyDecision, +) -> Result { + let route = resolve_proxy_route( + env, + request_url, + outbound_proxy_policy, + resolve_system_proxy, + ); + configure_builder_for_resolved_route(builder, route_class, &route) +} + +fn configure_builder_for_resolved_route( + builder: reqwest::ClientBuilder, + route_class: ClientRouteClass, + route: &OutboundProxyRoute, +) -> Result { + match route { + OutboundProxyRoute::TransportDefault => Ok(builder), + OutboundProxyRoute::Direct => Ok(builder.no_proxy()), + OutboundProxyRoute::Proxy { url, no_proxy } => { + let no_proxy = no_proxy.as_deref().and_then(reqwest::NoProxy::from_string); + configure_concrete_proxy(builder, route_class, url, no_proxy) + } + } +} + +fn configure_concrete_proxy( + builder: reqwest::ClientBuilder, + route_class: ClientRouteClass, + proxy_url: &str, + no_proxy: Option, +) -> Result { + let proxy = match reqwest::Proxy::all(proxy_url) { + Ok(proxy) => proxy, + Err(_source) => { + return Err(BuildRouteAwareHttpClientError::InvalidProxyConfig { route_class }); + } + }; + Ok(builder.proxy(proxy.no_proxy(no_proxy))) +} + +#[derive(Debug, Clone, PartialEq, Eq)] +#[allow(dead_code)] +struct RequestOrigin { + scheme: String, + host: String, + port: u16, +} + +impl RequestOrigin { + fn parse(request_url: &str) -> Option { + let uri = request_url.parse::().ok()?; + let scheme = uri.scheme_str()?.to_ascii_lowercase(); + let host = uri.host()?.trim_matches(['[', ']']).to_ascii_lowercase(); + let port = uri.port_u16().or(match scheme.as_str() { + "http" | "ws" => Some(80), + "https" | "wss" => Some(443), + _ => None, + })?; + Some(Self { scheme, host, port }) + } +} + +#[cfg_attr( + not(any(target_os = "windows", target_os = "macos")), + allow( + dead_code, + reason = "Direct and Proxy are constructed only by platform-specific resolvers" + ) +)] +#[derive(Debug, Clone, PartialEq, Eq)] +enum SystemProxyDecision { + Direct, + Proxy { url: String }, + Unavailable { failure: RouteFailureClass }, +} + +fn resolve_system_proxy(request_url: &str, origin: &RequestOrigin) -> SystemProxyDecision { + let cache = SYSTEM_PROXY_CACHE.get_or_init(|| Mutex::new(HashMap::new())); + resolve_system_proxy_with(cache, request_url, origin, resolve_platform_system_proxy) +} + +fn resolve_system_proxy_with( + cache: &Mutex>, + request_url: &str, + origin: &RequestOrigin, + resolve_platform_system_proxy: impl FnOnce(&str, &RequestOrigin) -> SystemProxyDecision, +) -> SystemProxyDecision { + let mut cache = match cache.lock() { + Ok(cache) => cache, + Err(error) => panic!("system proxy cache lock should not be poisoned: {error}"), + }; + let cache_key = system_proxy_cache_key(request_url); + if let Some(decision) = + cached_system_proxy_decision_from_cache(&mut cache, &cache_key, Instant::now()) + { + return decision; + } + + // Keep cache misses single-flight. Platform PAC/WPAD APIs are synchronous, so async callers + // run this work on the blocking pool; serializing misses prevents concurrent requests from + // consuming an unbounded number of blocking workers while system lookup is pending. + let decision = resolve_platform_system_proxy(request_url, origin); + insert_system_proxy_cache_entry(&mut cache, &cache_key, decision.clone(), Instant::now()); + decision +} + +#[cfg(target_os = "macos")] +fn resolve_platform_system_proxy(request_url: &str, origin: &RequestOrigin) -> SystemProxyDecision { + macos::resolve(request_url, origin) +} + +#[cfg(target_os = "windows")] +fn resolve_platform_system_proxy(request_url: &str, origin: &RequestOrigin) -> SystemProxyDecision { + windows::resolve(request_url, origin) +} + +#[cfg(not(any(target_os = "windows", target_os = "macos")))] +fn resolve_platform_system_proxy( + _request_url: &str, + _origin: &RequestOrigin, +) -> SystemProxyDecision { + SystemProxyDecision::Unavailable { + failure: RouteFailureClass::ProxyResolutionUnavailable, + } +} + +#[derive(Debug, Clone)] +struct CachedSystemProxyDecision { + decision: SystemProxyDecision, + expires_at: Instant, +} + +static SYSTEM_PROXY_CACHE: OnceLock>> = + OnceLock::new(); + +fn cached_system_proxy_decision(request_url: &str) -> Option { + let cache = SYSTEM_PROXY_CACHE.get_or_init(|| Mutex::new(HashMap::new())); + let mut cache = cache.lock().ok()?; + let key = system_proxy_cache_key(request_url); + cached_system_proxy_decision_from_cache(&mut cache, &key, Instant::now()) +} + +fn cached_system_proxy_decision_from_cache( + cache: &mut HashMap, + cache_key: &str, + now: Instant, +) -> Option { + let cached = cache.get(cache_key)?; + if cached.expires_at > now { + return Some(cached.decision.clone()); + } + cache.remove(cache_key); + None +} + +fn cache_system_proxy_decision(request_url: &str, decision: SystemProxyDecision) { + let cache = SYSTEM_PROXY_CACHE.get_or_init(|| Mutex::new(HashMap::new())); + if let Ok(mut cache) = cache.lock() { + let cache_key = system_proxy_cache_key(request_url); + insert_system_proxy_cache_entry(&mut cache, &cache_key, decision, Instant::now()); + } +} + +/// Primes one proxy decision for cross-crate integration tests. +/// +/// This is public only so tests in HTTP-client consumers can exercise system-proxy routing +/// deterministically on every supported platform. +pub fn cache_system_proxy_route_for_test(request_url: &str, proxy_url: String) { + cache_system_proxy_decision(request_url, SystemProxyDecision::Proxy { url: proxy_url }); +} + +fn insert_system_proxy_cache_entry( + cache: &mut HashMap, + cache_key: &str, + decision: SystemProxyDecision, + now: Instant, +) { + let ttl = match &decision { + SystemProxyDecision::Direct | SystemProxyDecision::Proxy { .. } => { + SYSTEM_PROXY_SUCCESS_CACHE_TTL + } + SystemProxyDecision::Unavailable { .. } => SYSTEM_PROXY_UNAVAILABLE_CACHE_TTL, + }; + + cache.retain(|_, cached| cached.expires_at > now); + if cache.len() >= SYSTEM_PROXY_CACHE_MAX_ENTRIES + && !cache.contains_key(cache_key) + && let Some(cache_key_to_evict) = cache + .iter() + .min_by_key(|(_, cached)| cached.expires_at) + .map(|(cache_key, _)| cache_key.clone()) + { + cache.remove(&cache_key_to_evict); + } + cache.insert( + cache_key.to_string(), + CachedSystemProxyDecision { + decision, + expires_at: now + ttl, + }, + ); +} + +fn system_proxy_cache_key(request_url: &str) -> String { + // Keep URL-specific PAC decisions without retaining the raw routed URL. + let mut hasher = Sha256::new(); + hasher.update(b"system-proxy-cache-v1\0"); + hasher.update(request_url.as_bytes()); + format!("{:x}", hasher.finalize()) +} + +#[cfg(any(test, target_os = "windows"))] +fn no_proxy_matches_origin(no_proxy: &str, origin: &RequestOrigin) -> bool { + no_proxy + .split(',') + .map(str::trim) + .filter(|entry| !entry.is_empty()) + .any(|entry| no_proxy_entry_matches_origin(entry, origin)) +} + +#[cfg(any(test, target_os = "windows"))] +fn no_proxy_entry_matches_origin(entry: &str, origin: &RequestOrigin) -> bool { + if entry == "*" { + return true; + } + + let mut entry = entry + .strip_prefix("http://") + .or_else(|| entry.strip_prefix("https://")) + .unwrap_or(entry) + .trim_matches(['[', ']']) + .to_ascii_lowercase(); + let mut port = None; + let parsed_host_port = entry.rsplit_once(':').and_then(|(host, candidate_port)| { + if host.contains(':') { + return None; + } + candidate_port + .parse::() + .ok() + .map(|parsed_port| (host.to_string(), parsed_port)) + }); + if let Some((host, parsed_port)) = parsed_host_port { + entry = host; + port = Some(parsed_port); + } + if port.is_some_and(|port| port != origin.port) { + return false; + } + + if let Some(suffix) = entry.strip_prefix('.') { + return origin.host == suffix || origin.host.ends_with(&format!(".{suffix}")); + } + + if entry.contains('*') { + return wildcard_host_match(&entry, &origin.host); + } + + origin.host == entry +} + +#[cfg(any(test, target_os = "windows"))] +fn wildcard_host_match(pattern: &str, host: &str) -> bool { + let mut remaining = host; + let mut first = true; + for part in pattern.split('*') { + if part.is_empty() { + continue; + } + if first && !pattern.starts_with('*') { + let Some(stripped) = remaining.strip_prefix(part) else { + return false; + }; + remaining = stripped; + } else { + let Some(index) = remaining.find(part) else { + return false; + }; + remaining = &remaining[index + part.len()..]; + } + first = false; + } + pattern.ends_with('*') || remaining.is_empty() +} + +#[cfg(any(test, target_os = "windows"))] +#[derive(Debug, Clone, PartialEq, Eq)] +enum ParsedProxyListDecision { + Direct, + Proxy(String), + UnsupportedScheme, + Unavailable, +} + +#[cfg(any(test, target_os = "windows"))] +fn parse_proxy_list(input: &str, target_scheme: &str) -> ParsedProxyListDecision { + let mut saw_unsupported = false; + + { + let mut process_token = |token: &str| { + let decision = parse_proxy_token(token, target_scheme); + match decision { + ParsedProxyListDecision::Direct => Some(ParsedProxyListDecision::Direct), + ParsedProxyListDecision::Proxy(url) => Some(ParsedProxyListDecision::Proxy(url)), + ParsedProxyListDecision::UnsupportedScheme => { + saw_unsupported = true; + None + } + ParsedProxyListDecision::Unavailable => None, + } + }; + + for segment in input + .split(';') + .map(str::trim) + .filter(|segment| !segment.is_empty()) + { + let mut parts = segment.split_whitespace(); + let directive = parts.next(); + let hostport = parts.next(); + let extra = parts.next(); + let is_proxy_directive = matches!( + directive.map(str::to_ascii_lowercase).as_deref(), + Some("proxy" | "http" | "https" | "socks" | "socks4" | "socks5") + ) && hostport.is_some() + && extra.is_none(); + + if is_proxy_directive { + if let Some(decision) = process_token(segment) { + return decision; + } + } else { + for token in segment.split_whitespace() { + if let Some(decision) = process_token(token) { + return decision; + } + } + } + } + } + + if saw_unsupported { + ParsedProxyListDecision::UnsupportedScheme + } else { + ParsedProxyListDecision::Unavailable + } +} + +#[cfg(any(test, target_os = "windows"))] +fn parse_proxy_token(token: &str, target_scheme: &str) -> ParsedProxyListDecision { + if token.eq_ignore_ascii_case("DIRECT") { + return ParsedProxyListDecision::Direct; + } + + if let Some(decision) = parse_proxy_key_token(token, target_scheme) { + return decision; + } + if token.contains('=') { + return ParsedProxyListDecision::Unavailable; + } + + let mut parts = token.split_whitespace(); + let directive = parts.next(); + let hostport = parts.next(); + if let (Some(directive), Some(hostport), None) = (directive, hostport, parts.next()) { + return match directive.to_ascii_lowercase().as_str() { + "proxy" | "http" => proxy_url_from_hostport("http", hostport), + "https" => proxy_url_from_hostport("https", hostport), + "socks" | "socks4" | "socks5" => ParsedProxyListDecision::UnsupportedScheme, + _ => ParsedProxyListDecision::Unavailable, + }; + } + + proxy_url_from_hostport("http", token) +} + +#[cfg(any(test, target_os = "windows"))] +fn parse_proxy_key_token(token: &str, target_scheme: &str) -> Option { + let (key, value) = token.split_once('=')?; + if key.trim().eq_ignore_ascii_case(target_scheme) { + Some(proxy_url_from_hostport("http", value.trim())) + } else { + Some(ParsedProxyListDecision::Unavailable) + } +} + +#[cfg(any(test, target_os = "windows"))] +fn proxy_url_from_hostport(proxy_scheme: &str, hostport: &str) -> ParsedProxyListDecision { + if hostport.is_empty() { + return ParsedProxyListDecision::Unavailable; + } + if hostport.contains("://") { + return ParsedProxyListDecision::Proxy(hostport.to_string()); + } + ParsedProxyListDecision::Proxy(format!("{proxy_scheme}://{hostport}")) +} + +trait EnvSource { + fn var(&self, key: &str) -> Option; +} + +struct ProcessEnv; + +impl EnvSource for ProcessEnv { + fn var(&self, key: &str) -> Option { + std::env::var(key).ok() + } +} + +fn proxy_env_value(env: &dyn EnvSource, upper: &str) -> Option { + let lower = upper.to_ascii_lowercase(); + env.var(upper) + .or_else(|| env.var(&lower)) + .filter(|value| !value.is_empty()) +} + +#[cfg(test)] +#[path = "route_aware_redirect_integration_tests.rs"] +mod redirect_integration_tests; + +#[cfg(test)] +#[path = "outbound_proxy_tests.rs"] +mod tests; diff --git a/codex-rs/http-client/src/outbound_proxy/macos.rs b/codex-rs/http-client/src/outbound_proxy/macos.rs new file mode 100644 index 0000000000000000000000000000000000000000..4f516517092ebcc0e9c5a6a58bda5639f8d48180 --- /dev/null +++ b/codex-rs/http-client/src/outbound_proxy/macos.rs @@ -0,0 +1,418 @@ +//! macOS system proxy resolution through SystemConfiguration and CFNetwork. + +use std::ffi::c_void; +use std::ptr; +use std::time::Duration; +use std::time::Instant; + +use super::MacosSystemProxyConfiguration; +use super::RequestOrigin; +use super::RouteFailureClass; +use super::SystemProxyDecision; +use system_configuration::core_foundation::array::CFArray; +use system_configuration::core_foundation::array::CFArrayRef; +use system_configuration::core_foundation::base::CFEqual; +use system_configuration::core_foundation::base::CFGetTypeID; +use system_configuration::core_foundation::base::CFIndex; +use system_configuration::core_foundation::base::CFType; +use system_configuration::core_foundation::base::CFTypeRef; +use system_configuration::core_foundation::base::TCFType; +use system_configuration::core_foundation::base::kCFAllocatorDefault; +use system_configuration::core_foundation::dictionary::CFDictionary; +use system_configuration::core_foundation::dictionary::CFDictionaryRef; +use system_configuration::core_foundation::error::CFErrorRef; +use system_configuration::core_foundation::number::CFNumber; +use system_configuration::core_foundation::runloop::CFRunLoop; +use system_configuration::core_foundation::runloop::CFRunLoopSource; +use system_configuration::core_foundation::runloop::CFRunLoopSourceInvalidate; +use system_configuration::core_foundation::runloop::CFRunLoopSourceRef; +use system_configuration::core_foundation::runloop::kCFRunLoopDefaultMode; +use system_configuration::core_foundation::string::CFString; +use system_configuration::core_foundation::string::CFStringRef; +use system_configuration::core_foundation::url::CFURL; +use system_configuration::core_foundation::url::CFURLCreateWithString; +use system_configuration::core_foundation::url::CFURLGetTypeID; +use system_configuration::core_foundation::url::CFURLRef; +use system_configuration::dynamic_store::SCDynamicStoreBuilder; + +const PAC_EXECUTION_TIMEOUT: Duration = Duration::from_secs(5); + +type ProxyDictionary = CFDictionary; +type ProxyArray = CFArray; + +#[repr(C)] +struct CFStreamClientContext { + version: CFIndex, + info: *mut c_void, + retain: Option *mut c_void>, + release: Option, + copy_description: Option CFStringRef>, +} + +type CFProxyAutoConfigurationResultCallback = + unsafe extern "C" fn(*mut c_void, CFArrayRef, CFErrorRef); + +#[link(name = "CFNetwork", kind = "framework")] +unsafe extern "C" { + static kCFProxyTypeKey: CFStringRef; + static kCFProxyHostNameKey: CFStringRef; + static kCFProxyPortNumberKey: CFStringRef; + static kCFProxyAutoConfigurationURLKey: CFStringRef; + static kCFProxyAutoConfigurationJavaScriptKey: CFStringRef; + static kCFProxyTypeNone: CFStringRef; + static kCFProxyTypeHTTP: CFStringRef; + static kCFProxyTypeHTTPS: CFStringRef; + static kCFProxyTypeSOCKS: CFStringRef; + static kCFProxyTypeAutoConfigurationURL: CFStringRef; + static kCFProxyTypeAutoConfigurationJavaScript: CFStringRef; + + fn CFNetworkCopyProxiesForURL(url: CFURLRef, proxy_settings: CFDictionaryRef) -> CFArrayRef; + fn CFNetworkExecuteProxyAutoConfigurationURL( + proxy_auto_config_url: CFURLRef, + target_url: CFURLRef, + callback: CFProxyAutoConfigurationResultCallback, + client_context: *mut CFStreamClientContext, + ) -> CFRunLoopSourceRef; + fn CFNetworkExecuteProxyAutoConfigurationScript( + proxy_auto_config_script: CFStringRef, + target_url: CFURLRef, + callback: CFProxyAutoConfigurationResultCallback, + client_context: *mut CFStreamClientContext, + ) -> CFRunLoopSourceRef; +} + +pub(super) fn resolve(request_url: &str, origin: &RequestOrigin) -> SystemProxyDecision { + let Some(target_url) = cf_url(request_url) else { + return SystemProxyDecision::Unavailable { + failure: RouteFailureClass::InvalidProxyConfig, + }; + }; + + let Some(settings) = system_proxy_settings() else { + return SystemProxyDecision::Unavailable { + failure: RouteFailureClass::ProxyResolutionUnavailable, + }; + }; + + let Some(proxies) = copy_proxies_for_url(&target_url, &settings) else { + return SystemProxyDecision::Unavailable { + failure: RouteFailureClass::ProxyResolutionUnavailable, + }; + }; + + proxy_array_decision(&proxies, &target_url, origin) +} + +pub(super) fn configuration(request_url: &str) -> MacosSystemProxyConfiguration { + let Some(target_url) = cf_url(request_url) else { + return MacosSystemProxyConfiguration::Unavailable; + }; + let Some(settings) = system_proxy_settings() else { + return MacosSystemProxyConfiguration::Unavailable; + }; + let Some(proxies) = copy_proxies_for_url(&target_url, &settings) else { + return MacosSystemProxyConfiguration::Unavailable; + }; + + (&proxies) + .into_iter() + .find_map(|proxy| { + let proxy_type = cf_string_value(&proxy, unsafe { kCFProxyTypeKey })?; + if cf_string_equals(&proxy_type, unsafe { kCFProxyTypeAutoConfigurationURL }) + || cf_string_equals(&proxy_type, unsafe { + kCFProxyTypeAutoConfigurationJavaScript + }) + { + Some(MacosSystemProxyConfiguration::Automatic) + } else if cf_string_equals(&proxy_type, unsafe { kCFProxyTypeHTTP }) + || cf_string_equals(&proxy_type, unsafe { kCFProxyTypeHTTPS }) + { + Some(MacosSystemProxyConfiguration::Manual) + } else { + cf_string_equals(&proxy_type, unsafe { kCFProxyTypeNone }) + .then_some(MacosSystemProxyConfiguration::Direct) + } + }) + .unwrap_or(MacosSystemProxyConfiguration::Unavailable) +} + +fn system_proxy_settings() -> Option> { + let store = SCDynamicStoreBuilder::new("Codex").build()?; + store.get_proxies() +} + +fn copy_proxies_for_url( + target_url: &CFURL, + settings: &CFDictionary, +) -> Option { + let proxies = unsafe { + CFNetworkCopyProxiesForURL( + target_url.as_concrete_TypeRef(), + settings.as_concrete_TypeRef(), + ) + }; + if proxies.is_null() { + None + } else { + Some(unsafe { ProxyArray::wrap_under_create_rule(proxies) }) + } +} + +fn proxy_array_decision( + proxies: &ProxyArray, + target_url: &CFURL, + origin: &RequestOrigin, +) -> SystemProxyDecision { + let mut saw_unsupported = false; + let mut saw_unavailable = false; + + // CFNetwork returns candidates in failover order, but the shared resolver currently carries + // only one route. This matches the Windows limitation; cross-platform retry requires request + // replay semantics and is intentionally deferred. + for proxy in proxies { + match proxy_entry_decision(&proxy, target_url, origin) { + ProxyEntryDecision::Direct => return SystemProxyDecision::Direct, + ProxyEntryDecision::Proxy { url } => return SystemProxyDecision::Proxy { url }, + ProxyEntryDecision::UnsupportedScheme => saw_unsupported = true, + ProxyEntryDecision::Unavailable => saw_unavailable = true, + } + } + + if saw_unsupported { + SystemProxyDecision::Unavailable { + failure: RouteFailureClass::UnsupportedProxyScheme, + } + } else if saw_unavailable { + SystemProxyDecision::Unavailable { + failure: RouteFailureClass::ProxyResolutionUnavailable, + } + } else { + SystemProxyDecision::Direct + } +} + +fn proxy_entry_decision( + proxy: &ProxyDictionary, + target_url: &CFURL, + origin: &RequestOrigin, +) -> ProxyEntryDecision { + let Some(proxy_type) = cf_string_value(proxy, unsafe { kCFProxyTypeKey }) else { + return ProxyEntryDecision::Unavailable; + }; + + if cf_string_equals(&proxy_type, unsafe { kCFProxyTypeNone }) { + return ProxyEntryDecision::Direct; + } + + if cf_string_equals(&proxy_type, unsafe { kCFProxyTypeHTTP }) { + return concrete_proxy_entry(proxy, "http"); + } + + if cf_string_equals(&proxy_type, unsafe { kCFProxyTypeHTTPS }) { + // CFNetwork's HTTPS proxy type is a tunneling proxy for HTTPS destinations; it does not + // preserve an explicit TLS-to-proxy transport. See https://developer.apple.com/documentation/cfnetwork/kcfproxytypehttps. + return concrete_proxy_entry(proxy, "http"); + } + + if cf_string_equals(&proxy_type, unsafe { kCFProxyTypeSOCKS }) { + return ProxyEntryDecision::UnsupportedScheme; + } + + if cf_string_equals(&proxy_type, unsafe { kCFProxyTypeAutoConfigurationURL }) { + let Some(pac_url) = cf_url_value(proxy, unsafe { kCFProxyAutoConfigurationURLKey }) else { + return ProxyEntryDecision::Unavailable; + }; + return pac_decision(execute_pac_url(&pac_url, target_url), target_url, origin); + } + + if cf_string_equals(&proxy_type, unsafe { + kCFProxyTypeAutoConfigurationJavaScript + }) { + let Some(script) = + cf_string_value(proxy, unsafe { kCFProxyAutoConfigurationJavaScriptKey }) + else { + return ProxyEntryDecision::Unavailable; + }; + return pac_decision( + execute_pac(|callback, context| unsafe { + CFNetworkExecuteProxyAutoConfigurationScript( + script.as_concrete_TypeRef(), + target_url.as_concrete_TypeRef(), + callback, + context, + ) + }), + target_url, + origin, + ); + } + + ProxyEntryDecision::Unavailable +} + +fn pac_decision( + result: Result, + target_url: &CFURL, + origin: &RequestOrigin, +) -> ProxyEntryDecision { + let proxies = match result { + Ok(proxies) => proxies, + Err(RouteFailureClass::UnsupportedProxyScheme) => { + return ProxyEntryDecision::UnsupportedScheme; + } + Err(_) => return ProxyEntryDecision::Unavailable, + }; + + match proxy_array_decision(&proxies, target_url, origin) { + SystemProxyDecision::Direct => ProxyEntryDecision::Direct, + SystemProxyDecision::Proxy { url } => ProxyEntryDecision::Proxy { url }, + SystemProxyDecision::Unavailable { + failure: RouteFailureClass::UnsupportedProxyScheme, + } => ProxyEntryDecision::UnsupportedScheme, + SystemProxyDecision::Unavailable { failure: _ } => ProxyEntryDecision::Unavailable, + } +} + +fn execute_pac_url(pac_url: &CFURL, target_url: &CFURL) -> Result { + execute_pac(|callback, context| unsafe { + CFNetworkExecuteProxyAutoConfigurationURL( + pac_url.as_concrete_TypeRef(), + target_url.as_concrete_TypeRef(), + callback, + context, + ) + }) +} + +fn execute_pac( + create_source: impl FnOnce( + CFProxyAutoConfigurationResultCallback, + *mut CFStreamClientContext, + ) -> CFRunLoopSourceRef, +) -> Result { + let mut state = PacRunLoopState { result: None }; + let mut context = CFStreamClientContext { + version: 0, + info: (&mut state as *mut PacRunLoopState).cast::(), + retain: None, + release: None, + copy_description: None, + }; + + let source = create_source(pac_result_callback, &mut context); + if source.is_null() { + return Err(RouteFailureClass::ProxyResolutionUnavailable); + } + + let source = unsafe { CFRunLoopSource::wrap_under_create_rule(source) }; + let run_loop = CFRunLoop::get_current(); + let mode = unsafe { kCFRunLoopDefaultMode }; + run_loop.add_source(&source, mode); + + let started_at = Instant::now(); + while state.result.is_none() && started_at.elapsed() < PAC_EXECUTION_TIMEOUT { + CFRunLoop::run_in_mode(mode, Duration::from_millis(50), true); + } + + if state.result.is_none() { + unsafe { CFRunLoopSourceInvalidate(source.as_concrete_TypeRef()) }; + } + run_loop.remove_source(&source, mode); + state + .result + .unwrap_or(Err(RouteFailureClass::ConnectTimeout)) +} + +unsafe extern "C" fn pac_result_callback( + client: *mut c_void, + proxies: CFArrayRef, + error: CFErrorRef, +) { + let state = unsafe { &mut *client.cast::() }; + state.result = if !error.is_null() || proxies.is_null() { + Some(Err(RouteFailureClass::ProxyResolutionUnavailable)) + } else { + Some(Ok(unsafe { ProxyArray::wrap_under_get_rule(proxies) })) + }; + CFRunLoop::get_current().stop(); +} + +struct PacRunLoopState { + result: Option>, +} + +fn concrete_proxy_entry(proxy: &ProxyDictionary, proxy_scheme: &str) -> ProxyEntryDecision { + let Some(host) = cf_string_value(proxy, unsafe { kCFProxyHostNameKey }) + .map(|host| host.to_string()) + .filter(|host| !host.is_empty()) + else { + return ProxyEntryDecision::Unavailable; + }; + + let host = bracket_ipv6_host(&host); + let url = match cf_i32_value(proxy, unsafe { kCFProxyPortNumberKey }) { + Some(port) if port > 0 => format!("{proxy_scheme}://{host}:{port}"), + _ => format!("{proxy_scheme}://{host}"), + }; + ProxyEntryDecision::Proxy { url } +} + +fn bracket_ipv6_host(host: &str) -> String { + if host.contains(':') && !host.starts_with('[') { + format!("[{host}]") + } else { + host.to_string() + } +} + +fn cf_string_value(proxy: &ProxyDictionary, key: CFStringRef) -> Option { + proxy + .find(key) + .and_then(|value| value.downcast::()) +} + +fn cf_i32_value(proxy: &ProxyDictionary, key: CFStringRef) -> Option { + proxy + .find(key) + .and_then(|value| value.downcast::()) + .and_then(|value| value.to_i32()) +} + +fn cf_url_value(proxy: &ProxyDictionary, key: CFStringRef) -> Option { + proxy.find(key).and_then(|value| { + if unsafe { CFGetTypeID(value.as_CFTypeRef()) == CFURLGetTypeID() } { + Some(unsafe { CFURL::wrap_under_get_rule(value.as_CFTypeRef() as CFURLRef) }) + } else { + value + .downcast::() + .and_then(|value| cf_url(value.to_string().as_str())) + } + }) +} + +fn cf_string_equals(value: &CFString, expected: CFStringRef) -> bool { + unsafe { CFEqual(value.as_CFTypeRef(), expected as CFTypeRef) != 0 } +} + +fn cf_url(value: &str) -> Option { + let value = CFString::new(value); + let url = unsafe { + CFURLCreateWithString( + kCFAllocatorDefault, + value.as_concrete_TypeRef(), + ptr::null(), + ) + }; + if url.is_null() { + None + } else { + Some(unsafe { CFURL::wrap_under_create_rule(url) }) + } +} + +enum ProxyEntryDecision { + Direct, + Proxy { url: String }, + UnsupportedScheme, + Unavailable, +} diff --git a/codex-rs/http-client/src/outbound_proxy/windows.rs b/codex-rs/http-client/src/outbound_proxy/windows.rs new file mode 100644 index 0000000000000000000000000000000000000000..4ff927723affa2bc99a91372794e14e49c914089 --- /dev/null +++ b/codex-rs/http-client/src/outbound_proxy/windows.rs @@ -0,0 +1,364 @@ +//! Windows system proxy resolution through WinHTTP. + +use std::ffi::c_void; +use std::ptr; + +use super::ParsedProxyListDecision; +use super::RequestOrigin; +use super::RouteFailureClass; +use super::SystemProxyDecision; +use super::no_proxy_matches_origin; +use super::parse_proxy_list; +use windows_sys::Win32::Foundation::ERROR_FILE_NOT_FOUND; +use windows_sys::Win32::Foundation::FALSE; +use windows_sys::Win32::Foundation::GetLastError; +use windows_sys::Win32::Foundation::GlobalFree; +use windows_sys::Win32::Foundation::TRUE; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_AUTODETECTION_FAILED; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_BAD_AUTO_PROXY_SCRIPT; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_CANNOT_CONNECT; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_CONNECTION_ERROR; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_INVALID_URL; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_LOGIN_FAILURE; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_NAME_NOT_RESOLVED; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_SCRIPT_EXECUTION_ERROR; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_SECURE_CERT_CN_INVALID; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_SECURE_CERT_DATE_INVALID; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_SECURE_CERT_REV_FAILED; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_SECURE_CERT_REVOKED; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_SECURE_CERT_WRONG_USAGE; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_SECURE_CHANNEL_ERROR; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_SECURE_FAILURE; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_SECURE_INVALID_CA; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_SECURE_INVALID_CERT; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_TIMEOUT; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_UNABLE_TO_DOWNLOAD_SCRIPT; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_UNHANDLED_SCRIPT_TYPE; +use windows_sys::Win32::Networking::WinHttp::ERROR_WINHTTP_UNRECOGNIZED_SCHEME; +use windows_sys::Win32::Networking::WinHttp::WINHTTP_ACCESS_TYPE_NAMED_PROXY; +use windows_sys::Win32::Networking::WinHttp::WINHTTP_ACCESS_TYPE_NO_PROXY; +use windows_sys::Win32::Networking::WinHttp::WINHTTP_AUTO_DETECT_TYPE_DHCP; +use windows_sys::Win32::Networking::WinHttp::WINHTTP_AUTO_DETECT_TYPE_DNS_A; +use windows_sys::Win32::Networking::WinHttp::WINHTTP_AUTOPROXY_AUTO_DETECT; +use windows_sys::Win32::Networking::WinHttp::WINHTTP_AUTOPROXY_CONFIG_URL; +use windows_sys::Win32::Networking::WinHttp::WINHTTP_AUTOPROXY_OPTIONS; +use windows_sys::Win32::Networking::WinHttp::WINHTTP_CURRENT_USER_IE_PROXY_CONFIG; +use windows_sys::Win32::Networking::WinHttp::WINHTTP_PROXY_INFO; +use windows_sys::Win32::Networking::WinHttp::WinHttpCloseHandle; +use windows_sys::Win32::Networking::WinHttp::WinHttpGetIEProxyConfigForCurrentUser; +use windows_sys::Win32::Networking::WinHttp::WinHttpGetProxyForUrl; +use windows_sys::Win32::Networking::WinHttp::WinHttpOpen; +use windows_sys::core::PWSTR; + +pub(super) fn resolve(request_url: &str, origin: &RequestOrigin) -> SystemProxyDecision { + let ie_config = match current_user_ie_proxy_config() { + Ok(config) => config, + Err(failure) => { + return SystemProxyDecision::Unavailable { failure }; + } + }; + + if let Some(pac_url) = ie_config.auto_config_url.as_deref() { + let decision = resolve_with_pac_url(request_url, origin, pac_url); + if !matches!(decision, SystemProxyDecision::Unavailable { .. }) { + return decision; + } + } + + if ie_config.auto_detect { + let decision = resolve_with_auto_detect(request_url, origin); + if !matches!(decision, SystemProxyDecision::Unavailable { .. }) { + return decision; + } + } + + if let Some(proxy) = ie_config.static_proxy.as_deref() { + if ie_config + .proxy_bypass + .as_deref() + .is_some_and(|bypass| proxy_bypass_matches_origin(bypass, origin)) + { + return SystemProxyDecision::Direct; + } + return proxy_list_decision(proxy, origin); + } + + if ie_config.auto_config_url.is_some() || ie_config.auto_detect { + SystemProxyDecision::Unavailable { + failure: RouteFailureClass::ProxyResolutionUnavailable, + } + } else { + SystemProxyDecision::Direct + } +} + +fn resolve_with_pac_url( + request_url: &str, + origin: &RequestOrigin, + pac_url: &str, +) -> SystemProxyDecision { + let pac_url = wide_null(pac_url); + let options = WINHTTP_AUTOPROXY_OPTIONS { + dwFlags: WINHTTP_AUTOPROXY_CONFIG_URL, + dwAutoDetectFlags: 0, + lpszAutoConfigUrl: pac_url.as_ptr(), + lpvReserved: ptr::null_mut(), + dwReserved: 0, + fAutoLogonIfChallenged: TRUE, + }; + resolve_with_winhttp_options(request_url, origin, options) +} + +fn resolve_with_auto_detect(request_url: &str, origin: &RequestOrigin) -> SystemProxyDecision { + let options = WINHTTP_AUTOPROXY_OPTIONS { + dwFlags: WINHTTP_AUTOPROXY_AUTO_DETECT, + dwAutoDetectFlags: WINHTTP_AUTO_DETECT_TYPE_DHCP | WINHTTP_AUTO_DETECT_TYPE_DNS_A, + lpszAutoConfigUrl: ptr::null(), + lpvReserved: ptr::null_mut(), + dwReserved: 0, + fAutoLogonIfChallenged: FALSE, + }; + resolve_with_winhttp_options(request_url, origin, options) +} + +fn resolve_with_winhttp_options( + request_url: &str, + origin: &RequestOrigin, + mut options: WINHTTP_AUTOPROXY_OPTIONS, +) -> SystemProxyDecision { + let Some(session) = WinHttpSession::open() else { + return SystemProxyDecision::Unavailable { + failure: classify_winhttp_error(last_error()), + }; + }; + + let request_url = wide_null(request_url); + let mut proxy_info = WINHTTP_PROXY_INFO { + dwAccessType: WINHTTP_ACCESS_TYPE_NO_PROXY, + lpszProxy: ptr::null_mut(), + lpszProxyBypass: ptr::null_mut(), + }; + let ok = unsafe { + WinHttpGetProxyForUrl( + session.0, + request_url.as_ptr(), + &mut options, + &mut proxy_info, + ) + }; + if ok == FALSE { + return SystemProxyDecision::Unavailable { + failure: classify_winhttp_error(last_error()), + }; + } + + let proxy_info = ProxyInfo::from_raw(proxy_info); + if proxy_info.access_type == WINHTTP_ACCESS_TYPE_NO_PROXY { + return SystemProxyDecision::Direct; + } + if proxy_info.access_type != WINHTTP_ACCESS_TYPE_NAMED_PROXY { + return SystemProxyDecision::Unavailable { + failure: RouteFailureClass::ProxyResolutionUnavailable, + }; + } + let Some(proxy) = proxy_info.proxy.as_deref() else { + return SystemProxyDecision::Unavailable { + failure: RouteFailureClass::ProxyResolutionUnavailable, + }; + }; + proxy_list_decision(proxy, origin) +} + +fn proxy_list_decision(proxy_list: &str, origin: &RequestOrigin) -> SystemProxyDecision { + match parse_proxy_list(proxy_list, &origin.scheme) { + ParsedProxyListDecision::Direct => SystemProxyDecision::Direct, + ParsedProxyListDecision::Proxy(url) => SystemProxyDecision::Proxy { url }, + ParsedProxyListDecision::UnsupportedScheme => SystemProxyDecision::Unavailable { + failure: RouteFailureClass::UnsupportedProxyScheme, + }, + ParsedProxyListDecision::Unavailable => SystemProxyDecision::Unavailable { + failure: RouteFailureClass::ProxyResolutionUnavailable, + }, + } +} + +fn current_user_ie_proxy_config() -> Result { + let mut raw = WINHTTP_CURRENT_USER_IE_PROXY_CONFIG { + fAutoDetect: FALSE, + lpszAutoConfigUrl: ptr::null_mut(), + lpszProxy: ptr::null_mut(), + lpszProxyBypass: ptr::null_mut(), + }; + let ok = unsafe { WinHttpGetIEProxyConfigForCurrentUser(&mut raw) }; + if ok == FALSE { + let error = last_error(); + if error == ERROR_FILE_NOT_FOUND { + // Match WinHTTP's fallback by attempting WPAD when no IE proxy settings exist. + return Ok(IeProxyConfig { + auto_detect: true, + ..Default::default() + }); + } + return Err(classify_winhttp_error(error)); + } + + let auto_config_url = GlobalWideString::from_raw(raw.lpszAutoConfigUrl).into_string(); + let static_proxy = GlobalWideString::from_raw(raw.lpszProxy).into_string(); + let proxy_bypass = GlobalWideString::from_raw(raw.lpszProxyBypass).into_string(); + + Ok(IeProxyConfig { + auto_detect: raw.fAutoDetect != FALSE, + auto_config_url, + static_proxy, + proxy_bypass, + }) +} + +#[derive(Debug, Default)] +struct IeProxyConfig { + auto_detect: bool, + auto_config_url: Option, + static_proxy: Option, + proxy_bypass: Option, +} + +struct ProxyInfo { + access_type: u32, + proxy: Option, + _proxy_bypass: Option, +} + +impl ProxyInfo { + fn from_raw(raw: WINHTTP_PROXY_INFO) -> Self { + Self { + access_type: raw.dwAccessType, + proxy: GlobalWideString::from_raw(raw.lpszProxy).into_string(), + _proxy_bypass: GlobalWideString::from_raw(raw.lpszProxyBypass).into_string(), + } + } +} + +struct GlobalWideString(PWSTR); + +impl GlobalWideString { + fn from_raw(ptr: PWSTR) -> Self { + Self(ptr) + } + + fn into_string(self) -> Option { + if self.0.is_null() { + return None; + } + let string = unsafe { wide_ptr_to_string(self.0) }; + if string.is_empty() { + None + } else { + Some(string) + } + } +} + +impl Drop for GlobalWideString { + fn drop(&mut self) { + if !self.0.is_null() { + unsafe { + GlobalFree(self.0.cast::()); + } + } + } +} + +struct WinHttpSession(*mut c_void); + +impl WinHttpSession { + fn open() -> Option { + let agent = wide_null("Codex"); + let handle = unsafe { + WinHttpOpen( + agent.as_ptr(), + WINHTTP_ACCESS_TYPE_NO_PROXY, + ptr::null(), + ptr::null(), + 0, + ) + }; + if handle.is_null() { + None + } else { + Some(Self(handle)) + } + } +} + +impl Drop for WinHttpSession { + fn drop(&mut self) { + if !self.0.is_null() { + unsafe { + WinHttpCloseHandle(self.0); + } + } + } +} + +fn proxy_bypass_matches_origin(proxy_bypass: &str, origin: &RequestOrigin) -> bool { + proxy_bypass + .split(|ch: char| ch == ';' || ch == ',' || ch.is_whitespace()) + .map(str::trim) + .filter(|entry| !entry.is_empty()) + .any(|entry| { + if entry.eq_ignore_ascii_case("") { + !origin.host.contains('.') + } else { + no_proxy_matches_origin(entry, origin) + } + }) +} + +#[cfg(test)] +#[path = "windows_tests.rs"] +mod tests; + +fn wide_null(value: &str) -> Vec { + value.encode_utf16().chain(std::iter::once(0)).collect() +} + +unsafe fn wide_ptr_to_string(ptr: PWSTR) -> String { + let mut len = 0; + while unsafe { *ptr.add(len) } != 0 { + len += 1; + } + let slice = unsafe { std::slice::from_raw_parts(ptr, len) }; + String::from_utf16_lossy(slice) +} + +fn last_error() -> u32 { + unsafe { GetLastError() } +} + +fn classify_winhttp_error(code: u32) -> RouteFailureClass { + match code { + ERROR_WINHTTP_TIMEOUT => RouteFailureClass::ConnectTimeout, + ERROR_WINHTTP_LOGIN_FAILURE => RouteFailureClass::ProxyAuthenticationRequired, + ERROR_WINHTTP_AUTODETECTION_FAILED + | ERROR_WINHTTP_BAD_AUTO_PROXY_SCRIPT + | ERROR_WINHTTP_SCRIPT_EXECUTION_ERROR + | ERROR_WINHTTP_UNABLE_TO_DOWNLOAD_SCRIPT + | ERROR_WINHTTP_UNHANDLED_SCRIPT_TYPE => RouteFailureClass::ProxyResolutionUnavailable, + ERROR_WINHTTP_SECURE_CERT_CN_INVALID + | ERROR_WINHTTP_SECURE_CERT_DATE_INVALID + | ERROR_WINHTTP_SECURE_CERT_REVOKED + | ERROR_WINHTTP_SECURE_CERT_REV_FAILED + | ERROR_WINHTTP_SECURE_CERT_WRONG_USAGE + | ERROR_WINHTTP_SECURE_CHANNEL_ERROR + | ERROR_WINHTTP_SECURE_FAILURE + | ERROR_WINHTTP_SECURE_INVALID_CA + | ERROR_WINHTTP_SECURE_INVALID_CERT => RouteFailureClass::TlsError, + ERROR_WINHTTP_INVALID_URL | ERROR_WINHTTP_UNRECOGNIZED_SCHEME => { + RouteFailureClass::InvalidProxyConfig + } + ERROR_WINHTTP_CANNOT_CONNECT + | ERROR_WINHTTP_CONNECTION_ERROR + | ERROR_WINHTTP_NAME_NOT_RESOLVED => RouteFailureClass::ResolverError, + _ => RouteFailureClass::ResolverError, + } +} diff --git a/codex-rs/http-client/src/outbound_proxy/windows_tests.rs b/codex-rs/http-client/src/outbound_proxy/windows_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..589c06661cb3d10c0bf7f7d48ef75efce2cc5168 --- /dev/null +++ b/codex-rs/http-client/src/outbound_proxy/windows_tests.rs @@ -0,0 +1,20 @@ +//! Windows proxy parsing tests. + +use super::*; + +#[test] +fn proxy_bypass_matches_whitespace_separated_winhttp_entries() { + let local_origin = RequestOrigin { + scheme: "https".to_string(), + host: "intranet".to_string(), + port: 443, + }; + assert!(proxy_bypass_matches_origin(" *.corp", &local_origin)); + + let corp_origin = RequestOrigin { + scheme: "https".to_string(), + host: "service.corp".to_string(), + port: 443, + }; + assert!(proxy_bypass_matches_origin(" *.corp", &corp_origin)); +} diff --git a/codex-rs/http-client/src/outbound_proxy_redirect_coverage_tests.rs b/codex-rs/http-client/src/outbound_proxy_redirect_coverage_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..8aefa93f1ea995cd85d638f2bb24f9a8854617dc --- /dev/null +++ b/codex-rs/http-client/src/outbound_proxy_redirect_coverage_tests.rs @@ -0,0 +1,221 @@ +use super::*; +use pretty_assertions::assert_eq; +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; +use tracing_subscriber::Layer; +use tracing_subscriber::layer::SubscriberExt; + +#[tokio::test] +async fn route_aware_pool_strips_credentials_on_cross_origin_redirect() { + let (destination_addr, destination_thread) = spawn_proxy_listener(); + let destination_url = format!("http://{destination_addr}/final"); + let (redirect_addr, redirect_thread) = spawn_redirect_listener(&destination_url); + let initial_url = format!("http://{redirect_addr}/start"); + cache_system_proxy_decision(&initial_url, SystemProxyDecision::Direct); + cache_system_proxy_decision(&destination_url, SystemProxyDecision::Direct); + let pool = crate::RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy), + ClientRouteClass::Api, + ); + + let response = tokio::time::timeout( + Duration::from_secs(2), + pool.get(&initial_url) + .header(AUTHORIZATION, "Bearer origin-secret") + .header(COOKIE, "session=origin-secret") + .header(PROXY_AUTHORIZATION, "Basic proxy-secret") + .send(), + ) + .await + .expect("redirected request should finish") + .expect("cross-origin redirect should succeed"); + let initial_request = only_request(redirect_thread, "redirect"); + let destination_request = only_request(destination_thread, "destination"); + + assert_eq!(response.url().as_str(), destination_url); + assert_eq!( + [ + credential_headers(&initial_request), + credential_headers(&destination_request), + ], + [(true, true, true), (false, false, false)] + ); +} + +#[tokio::test] +async fn route_aware_pool_retains_credentials_for_same_origin_and_route() { + let (proxy_addr, proxy_thread) = spawn_http_listener(vec![ + "HTTP/1.1 302 Found\r\nLocation: /final\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + .to_string(), + "HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok".to_string(), + ]); + let initial_url = "http://same-route.test/start"; + let redirected_url = "http://same-route.test/final"; + let route = SystemProxyDecision::Proxy { + url: format!("http://{proxy_addr}"), + }; + cache_system_proxy_decision(initial_url, route.clone()); + cache_system_proxy_decision(redirected_url, route); + let pool = crate::RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy), + ClientRouteClass::Api, + ); + + let response = tokio::time::timeout( + Duration::from_secs(2), + pool.get(initial_url) + .header(AUTHORIZATION, "Bearer origin-secret") + .header(COOKIE, "session=origin-secret") + .header(PROXY_AUTHORIZATION, "Basic proxy-secret") + .send(), + ) + .await + .expect("redirected request should finish") + .expect("same-route redirect should succeed"); + let requests = proxy_thread.join().expect("proxy thread should finish"); + + assert_eq!(response.url().as_str(), redirected_url); + assert_eq!( + requests + .iter() + .map(|request| credential_headers(request)) + .collect::>(), + vec![(true, true, true), (true, true, true)] + ); +} + +#[tokio::test] +async fn route_aware_pool_sanitizes_redirected_failure_logs() { + let log_buffer = Arc::new(Mutex::new(Vec::new())); + let subscriber = tracing_subscriber::registry().with( + tracing_subscriber::fmt::layer() + .with_ansi(false) + .with_writer(TestLogWriter { + buffer: Arc::clone(&log_buffer), + }) + .with_filter( + tracing_subscriber::filter::Targets::new() + .with_target("codex_http_client", tracing::Level::TRACE), + ), + ); + let _guard = tracing::subscriber::set_default(subscriber); + tracing::debug!(target: "codex_http_client", "log capture sentinel"); + + let (enabled_failure_addr, enabled_failure_thread) = spawn_failing_listener(); + let enabled_target_url = + format!("http://{enabled_failure_addr}/final?token=enabled-target-secret"); + let (enabled_redirect_addr, enabled_redirect_thread) = + spawn_redirect_listener(&enabled_target_url); + let enabled_initial_url = format!("http://{enabled_redirect_addr}/start"); + cache_system_proxy_decision(&enabled_initial_url, SystemProxyDecision::Direct); + cache_system_proxy_decision(&enabled_target_url, SystemProxyDecision::Direct); + let enabled_pool = crate::RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy), + ClientRouteClass::Api, + ); + + enabled_pool + .get(&enabled_initial_url) + .timeout(Duration::from_secs(2)) + .send() + .await + .expect_err("redirect target should fail"); + only_request(enabled_redirect_thread, "enabled redirect"); + enabled_failure_thread + .join() + .expect("enabled failure thread should finish"); + + let (disabled_failure_addr, disabled_failure_thread) = spawn_failing_listener(); + let disabled_target_url = + format!("http://{disabled_failure_addr}/final?token=disabled-target-secret"); + let (disabled_redirect_addr, disabled_redirect_thread) = + spawn_redirect_listener(&disabled_target_url); + let disabled_initial_url = format!("http://{disabled_redirect_addr}/start"); + cache_system_proxy_decision(&disabled_initial_url, SystemProxyDecision::Direct); + cache_system_proxy_decision(&disabled_target_url, SystemProxyDecision::Direct); + let disabled_pool = crate::RouteAwareClientPool::new_without_request_logging( + HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy), + ClientRouteClass::Api, + ); + + disabled_pool + .get(&disabled_initial_url) + .timeout(Duration::from_secs(2)) + .send() + .await + .expect_err("redirect target should fail"); + only_request(disabled_redirect_thread, "disabled redirect"); + disabled_failure_thread + .join() + .expect("disabled failure thread should finish"); + + let logs = String::from_utf8(log_buffer.lock().expect("log buffer lock").clone()) + .expect("logs should be UTF-8"); + assert!(logs.contains("log capture sentinel")); + assert!(logs.contains(&enabled_initial_url)); + assert!(!logs.contains(&disabled_initial_url)); + assert_eq!(logs.matches("Request failed").count(), 1); + assert!(logs.contains("is_timeout")); + assert!(logs.contains("is_connect")); + for secret in [ + "enabled-target-secret", + "disabled-target-secret", + enabled_target_url.as_str(), + disabled_target_url.as_str(), + "error=", + ] { + assert!(!logs.contains(secret), "logs exposed {secret}:\n{logs}"); + } +} + +fn credential_headers(request: &str) -> (bool, bool, bool) { + let names = request + .lines() + .filter_map(|line| { + line.split_once(':') + .map(|(name, _)| name.to_ascii_lowercase()) + }) + .collect::>(); + ( + names.iter().any(|name| name == "authorization"), + names.iter().any(|name| name == "cookie"), + names.iter().any(|name| name == "proxy-authorization"), + ) +} + +fn spawn_failing_listener() -> (std::net::SocketAddr, std::thread::JoinHandle<()>) { + let listener = + std::net::TcpListener::bind(("127.0.0.1", 0)).expect("failing listener should bind"); + let address = listener + .local_addr() + .expect("failing listener should have an address"); + listener + .set_nonblocking(true) + .expect("failing listener should become nonblocking"); + let thread = std::thread::spawn(move || { + let deadline = Instant::now() + Duration::from_secs(2); + let (mut stream, _) = loop { + match listener.accept() { + Ok(connection) => break connection, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => { + assert!( + Instant::now() < deadline, + "failing listener should receive a request" + ); + std::thread::sleep(Duration::from_millis(10)); + } + Err(error) => panic!("failing listener should accept: {error}"), + } + }; + // Accepted sockets inherit the listener's nonblocking mode on macOS. + stream + .set_nonblocking(false) + .expect("failing stream should become blocking"); + stream + .set_read_timeout(Some(Duration::from_secs(2))) + .expect("failing stream should get a read timeout"); + read_http_message(&mut stream); + }); + (address, thread) +} diff --git a/codex-rs/http-client/src/outbound_proxy_tests.rs b/codex-rs/http-client/src/outbound_proxy_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..22b3716d9cf9cfad119d7256362d90b2907c5da2 --- /dev/null +++ b/codex-rs/http-client/src/outbound_proxy_tests.rs @@ -0,0 +1,657 @@ +//! Shared outbound proxy policy tests. + +use super::*; +use crate::HttpClientBuilder; +use http::HeaderMap; +use http::HeaderValue; +use http::header::AUTHORIZATION; +use http::header::COOKIE; +use http::header::PROXY_AUTHORIZATION; +use pretty_assertions::assert_eq; +use std::io::Read; +use std::io::Write; +use std::sync::Arc; +use std::sync::Mutex; +use tracing_subscriber::Layer; +use tracing_subscriber::layer::SubscriberExt; + +#[path = "outbound_proxy_redirect_coverage_tests.rs"] +mod redirect_coverage_tests; + +struct MapEnv { + values: HashMap, +} + +fn spawn_proxy_listener() -> (std::net::SocketAddr, std::thread::JoinHandle>) { + spawn_http_listener(vec![ + "HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok".to_string(), + ]) +} + +fn spawn_redirect_listener( + location: &str, +) -> (std::net::SocketAddr, std::thread::JoinHandle>) { + spawn_http_listener(vec![format!( + "HTTP/1.1 302 Found\r\nLocation: {location}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + )]) +} + +fn spawn_http_listener( + responses: Vec, +) -> (std::net::SocketAddr, std::thread::JoinHandle>) { + let listener = + std::net::TcpListener::bind(("127.0.0.1", 0)).expect("HTTP listener should bind"); + let address = listener + .local_addr() + .expect("HTTP listener should have an address"); + listener + .set_nonblocking(true) + .expect("HTTP listener should become nonblocking"); + let thread = std::thread::spawn(move || { + let mut requests = Vec::new(); + for response in responses { + let deadline = Instant::now() + Duration::from_secs(10); + let (mut stream, _) = loop { + match listener.accept() { + Ok(connection) => break connection, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => { + assert!( + Instant::now() < deadline, + "HTTP listener should receive the next request" + ); + std::thread::sleep(Duration::from_millis(10)); + } + Err(error) => panic!("HTTP listener should accept: {error}"), + } + }; + stream + .set_nonblocking(false) + .expect("HTTP stream should become blocking"); + stream + .set_read_timeout(Some(Duration::from_secs(10))) + .expect("HTTP stream should get a read timeout"); + requests.push(read_http_message(&mut stream)); + stream + .write_all(response.as_bytes()) + .expect("HTTP listener should write response"); + } + requests + }); + (address, thread) +} + +fn only_request(thread: std::thread::JoinHandle>, source: &str) -> String { + let requests = thread + .join() + .unwrap_or_else(|_| panic!("{source} thread should finish")); + let [request]: [String; 1] = requests.try_into().unwrap_or_else(|requests: Vec| { + panic!( + "{source} should receive one request, got {}", + requests.len() + ) + }); + request +} + +fn read_http_message(stream: &mut impl Read) -> String { + let mut buffer = Vec::new(); + let mut chunk = [0_u8; 1024]; + loop { + let bytes_read = stream.read(&mut chunk).expect("HTTP message should read"); + if bytes_read == 0 { + break; + } + buffer.extend_from_slice(&chunk[..bytes_read]); + if let Some(header_end) = buffer.windows(4).position(|window| window == b"\r\n\r\n") { + let body_start = header_end + 4; + let headers = String::from_utf8_lossy(&buffer[..body_start]); + let content_length = headers + .lines() + .filter_map(|line| line.split_once(':')) + .find_map(|(name, value)| { + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + if buffer.len() >= body_start + content_length { + break; + } + } + } + String::from_utf8_lossy(&buffer).into_owned() +} + +#[test] +fn cloned_factories_share_chatgpt_cookie_stores_without_changing_value_equality() { + let factory = + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault).with_chatgpt_cookies([ + HeaderValue::from_static("first=true"), + HeaderValue::from_static("second=true"), + ]); + let cloned = factory.clone(); + + assert!(Arc::ptr_eq( + &factory.chatgpt_cookie_store().expect("configured store"), + &cloned.chatgpt_cookie_store().expect("shared store"), + )); + assert_eq!( + factory, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault).with_chatgpt_cookies([ + HeaderValue::from_static("first=true"), + HeaderValue::from_static("second=true"), + ]) + ); +} + +#[test] +fn websocket_route_uses_http_equivalent_for_system_resolution() { + let env = MapEnv { + values: HashMap::new(), + }; + let route = resolve_proxy_route( + &env, + "wss://api.openai.com/v1/responses", + OutboundProxyPolicy::RespectSystemProxy, + |request_url, origin| { + assert_eq!(request_url, "https://api.openai.com/v1/responses"); + assert_eq!(origin.scheme, "https"); + assert_eq!(origin.host, "api.openai.com"); + assert_eq!(origin.port, 443); + SystemProxyDecision::Proxy { + url: "http://proxy.example:8080".to_string(), + } + }, + ); + + assert_eq!( + route, + OutboundProxyRoute::Proxy { + url: "http://proxy.example:8080".to_string(), + no_proxy: None, + } + ); +} + +#[test] +fn reqwest_default_route_preserves_transport_proxy_behavior() { + let env = MapEnv { + values: HashMap::new(), + }; + let route = resolve_proxy_route( + &env, + "wss://api.openai.com/v1/responses", + OutboundProxyPolicy::ReqwestDefault, + |_, _| panic!("default policy should not resolve system proxy settings"), + ); + + assert_eq!(route, OutboundProxyRoute::TransportDefault); +} + +#[cfg(target_os = "macos")] +#[test] +fn macos_proxy_configuration_rejects_invalid_destination() { + assert_eq!( + macos_system_proxy_configuration("not a valid destination"), + MacosSystemProxyConfiguration::Unavailable + ); +} + +impl EnvSource for MapEnv { + fn var(&self, key: &str) -> Option { + self.values.get(key).cloned() + } +} + +#[test] +fn proxy_env_value_matches_reqwest_casing_precedence() { + let env = MapEnv { + values: HashMap::from([ + ("HTTPS_PROXY".to_string(), "upper".to_string()), + ("https_proxy".to_string(), "lower".to_string()), + ("http_proxy".to_string(), "lower-only".to_string()), + ("ALL_PROXY".to_string(), String::new()), + ("all_proxy".to_string(), "masked".to_string()), + ]), + }; + + assert_eq!( + proxy_env_value(&env, "HTTPS_PROXY"), + Some("upper".to_string()) + ); + assert_eq!( + proxy_env_value(&env, "HTTP_PROXY"), + Some("lower-only".to_string()) + ); + assert_eq!(proxy_env_value(&env, "ALL_PROXY"), None); +} + +#[test] +fn environment_fallback_reads_injected_proxy_environment() { + let env = MapEnv { + values: HashMap::from([("HTTPS_PROXY".to_string(), "://invalid".to_string())]), + }; + let route = resolve_env_proxy_route(&env, EnvProxyKind::Https); + let result = configure_builder_for_resolved_route( + reqwest::Client::builder(), + ClientRouteClass::Auth, + &route, + ); + + assert!(matches!( + result, + Err(BuildRouteAwareHttpClientError::InvalidProxyConfig { + route_class: ClientRouteClass::Auth, + }) + )); +} + +#[test] +fn unavailable_system_route_resolves_environment_or_direct_explicitly() { + let env = MapEnv { + values: HashMap::from([ + ( + "HTTPS_PROXY".to_string(), + "http://proxy.example:8080".to_string(), + ), + ("NO_PROXY".to_string(), "localhost,.internal".to_string()), + ]), + }; + + assert_eq!( + route_from_system_decision( + &env, + EnvProxyKind::Https, + SystemProxyDecision::Unavailable { + failure: RouteFailureClass::ProxyResolutionUnavailable, + }, + ), + OutboundProxyRoute::Proxy { + url: "http://proxy.example:8080".to_string(), + no_proxy: Some("localhost,.internal".to_string()), + } + ); + assert_eq!( + route_from_system_decision( + &MapEnv { + values: HashMap::new(), + }, + EnvProxyKind::Https, + SystemProxyDecision::Unavailable { + failure: RouteFailureClass::ProxyResolutionUnavailable, + }, + ), + OutboundProxyRoute::Direct + ); +} + +#[test] +fn unavailable_system_route_preserves_wss_http_proxy_fallback() { + let env = MapEnv { + values: HashMap::from([( + "HTTP_PROXY".to_string(), + "http://proxy.example:8080".to_string(), + )]), + }; + + let route = resolve_proxy_route( + &env, + "wss://api.openai.com/v1/responses", + OutboundProxyPolicy::RespectSystemProxy, + |_, _| SystemProxyDecision::Unavailable { + failure: RouteFailureClass::ProxyResolutionUnavailable, + }, + ); + + assert_eq!( + route, + OutboundProxyRoute::Proxy { + url: "http://proxy.example:8080".to_string(), + no_proxy: None, + } + ); +} + +#[cfg(any(target_os = "windows", target_os = "macos"))] +#[tokio::test] +async fn async_resolution_uses_cached_route_before_global_permit() { + let request_url = "https://cached-fast-path.test/request"; + cache_system_proxy_decision(request_url, SystemProxyDecision::Direct); + let factory = HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy); + let permit = ASYNC_SYSTEM_PROXY_RESOLUTION_PERMIT + .acquire() + .await + .expect("global proxy permit should stay open"); + + let route = tokio::time::timeout( + Duration::from_secs(2), + factory.resolve_proxy_route_async(request_url.to_string()), + ) + .await + .expect("cached resolution should not wait for the global permit") + .expect("cached route should resolve"); + drop(permit); + + assert_eq!(route, OutboundProxyRoute::Direct); +} + +#[tokio::test] +async fn enabled_environment_proxy_routes_request_through_proxy() { + let (proxy_addr, proxy_thread) = spawn_proxy_listener(); + let env = MapEnv { + values: HashMap::from([("HTTP_PROXY".to_string(), format!("http://{proxy_addr}"))]), + }; + let request_url = "http://enabled-proxy.test/proxy-check"; + let builder = configure_proxy_for_route( + &env, + reqwest::Client::builder().timeout(Duration::from_secs(2)), + request_url, + ClientRouteClass::Auth, + OutboundProxyPolicy::RespectSystemProxy, + |_, _| SystemProxyDecision::Unavailable { + failure: RouteFailureClass::ProxyResolutionUnavailable, + }, + ) + .expect("enabled proxy route should configure"); + + let response = builder + .build() + .expect("proxy client should build") + .get(request_url) + .send() + .await + .expect("request should use local proxy"); + let proxy_request = only_request(proxy_thread, "proxy"); + + assert_eq!(response.status(), reqwest::StatusCode::OK); + assert_eq!( + proxy_request.lines().next(), + Some("GET http://enabled-proxy.test/proxy-check HTTP/1.1") + ); +} + +#[tokio::test] +async fn route_aware_builder_preserves_default_headers() { + let (server_addr, server_thread) = spawn_proxy_listener(); + let request_url = format!("http://{server_addr}/builder-check"); + cache_system_proxy_decision(&request_url, SystemProxyDecision::Direct); + let mut headers = HeaderMap::new(); + headers.insert("x-builder-test", HeaderValue::from_static("preserved")); + let factory = HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy); + let client = HttpClientBuilder::new() + .default_headers(headers) + .build_respecting_outbound_proxy_policy(&factory, &request_url, ClientRouteClass::Api) + .expect("route-aware client should build"); + + let response = client + .get(&request_url) + .send() + .await + .expect("request should use direct route"); + let request = only_request(server_thread, "server"); + + assert!(response.status().is_success()); + assert!( + request + .lines() + .any(|line| line.eq_ignore_ascii_case("x-builder-test: preserved")) + ); +} + +#[tokio::test] +async fn route_aware_pool_uses_respect_system_proxy_route_for_exact_url() { + let (proxy_addr, proxy_thread) = spawn_proxy_listener(); + let request_url = "http://route-aware-proxy.test/proxy-check?pac=exact"; + cache_system_proxy_decision( + request_url, + SystemProxyDecision::Proxy { + url: format!("http://{proxy_addr}"), + }, + ); + let pool = crate::RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy), + ClientRouteClass::Api, + ); + + let response = tokio::time::timeout(Duration::from_secs(2), pool.get(request_url).send()) + .await + .expect("proxy request should finish") + .expect("request should use local proxy"); + let proxy_request = only_request(proxy_thread, "proxy"); + + assert_eq!(response.status(), reqwest::StatusCode::OK); + assert_eq!( + proxy_request.lines().next(), + Some("GET http://route-aware-proxy.test/proxy-check?pac=exact HTTP/1.1") + ); +} + +#[tokio::test] +async fn route_aware_pool_logs_only_the_final_redirect_outcome() { + let (proxy_addr, proxy_thread) = spawn_proxy_listener(); + let redirected_url = "http://redirect-target.test/final?token=redirect-target-secret-value"; + let (redirect_addr, redirect_thread) = spawn_redirect_listener(redirected_url); + let initial_url = format!("http://{redirect_addr}/start"); + cache_system_proxy_decision(&initial_url, SystemProxyDecision::Direct); + cache_system_proxy_decision( + redirected_url, + SystemProxyDecision::Proxy { + url: format!("http://{proxy_addr}"), + }, + ); + let pool = crate::RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy), + ClientRouteClass::Api, + ); + let log_buffer = Arc::new(Mutex::new(Vec::new())); + let subscriber = tracing_subscriber::registry().with( + tracing_subscriber::fmt::layer() + .with_ansi(false) + .with_writer(TestLogWriter { + buffer: Arc::clone(&log_buffer), + }) + .with_filter( + tracing_subscriber::filter::Targets::new() + .with_target("codex_http_client", tracing::Level::TRACE), + ), + ); + let _guard = tracing::subscriber::set_default(subscriber); + tracing::debug!(target: "codex_http_client", "log capture sentinel"); + + let response = tokio::time::timeout(Duration::from_secs(2), pool.get(&initial_url).send()) + .await + .expect("redirected request should finish") + .expect("redirected request should use selected routes"); + only_request(redirect_thread, "redirect"); + only_request(proxy_thread, "proxy"); + + assert_eq!(response.status(), reqwest::StatusCode::OK); + let logs = String::from_utf8(log_buffer.lock().expect("log buffer lock").clone()) + .expect("logs should be UTF-8"); + assert!(logs.contains("log capture sentinel")); + assert!(logs.contains(&initial_url)); + assert_eq!(logs.matches("Request completed").count(), 1); + for secret in ["redirect-target-secret-value", redirected_url, "location"] { + assert!( + !logs + .to_ascii_lowercase() + .contains(&secret.to_ascii_lowercase()) + ); + } +} + +#[test] +fn parses_pac_proxy_tokens() { + assert_eq!( + parse_proxy_list("PROXY proxy.internal:8080; DIRECT", "https"), + ParsedProxyListDecision::Proxy("http://proxy.internal:8080".to_string()) + ); + assert_eq!( + parse_proxy_list("HTTPS proxy.internal:8443", "https"), + ParsedProxyListDecision::Proxy("https://proxy.internal:8443".to_string()) + ); +} + +#[test] +fn unavailable_system_proxy_decision_is_cached() { + let request_url = "https://unavailable-cache.test/oauth/token"; + let decision = SystemProxyDecision::Unavailable { + failure: RouteFailureClass::ProxyResolutionUnavailable, + }; + + cache_system_proxy_decision(request_url, decision.clone()); + + assert_eq!(cached_system_proxy_decision(request_url), Some(decision)); +} + +#[derive(Clone)] +struct TestLogWriter { + buffer: Arc>>, +} + +struct TestLogSink { + buffer: Arc>>, +} + +impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for TestLogWriter { + type Writer = TestLogSink; + + fn make_writer(&'a self) -> Self::Writer { + TestLogSink { + buffer: Arc::clone(&self.buffer), + } + } +} + +impl Write for TestLogSink { + fn write(&mut self, buffer: &[u8]) -> io::Result { + let mut log_buffer = self + .buffer + .lock() + .map_err(|_| io::Error::other("log buffer lock was poisoned"))?; + log_buffer.extend(buffer); + Ok(buffer.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} + +#[test] +fn system_proxy_resolution_is_single_flight() { + let cache = Arc::new(Mutex::new(HashMap::new())); + let request_url = "https://single-flight.test/models"; + let origin = RequestOrigin::parse(request_url).expect("valid request URL"); + let (started_tx, started_rx) = std::sync::mpsc::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + let worker_cache = Arc::clone(&cache); + let worker_origin = origin.clone(); + + let worker = std::thread::spawn(move || { + resolve_system_proxy_with(&worker_cache, request_url, &worker_origin, |_, _| { + started_tx.send(()).expect("test should still be running"); + release_rx.recv().expect("test should release resolver"); + SystemProxyDecision::Direct + }) + }); + + started_rx.recv().expect("resolver should start"); + assert!(matches!( + cache.try_lock(), + Err(std::sync::TryLockError::WouldBlock) + )); + release_tx + .send(()) + .expect("resolver should still be running"); + assert_eq!( + worker.join().expect("resolver should finish"), + SystemProxyDecision::Direct + ); + assert_eq!( + resolve_system_proxy_with(&cache, request_url, &origin, |_, _| { + panic!("cached waiter should not resolve the platform proxy again") + }), + SystemProxyDecision::Direct + ); +} + +#[test] +fn system_proxy_cache_is_bounded() { + let mut cache = HashMap::new(); + let now = Instant::now(); + + for index in 0..=SYSTEM_PROXY_CACHE_MAX_ENTRIES { + insert_system_proxy_cache_entry( + &mut cache, + &format!("https://bounded-cache.test/{index}"), + SystemProxyDecision::Direct, + now, + ); + } + + assert_eq!(cache.len(), SYSTEM_PROXY_CACHE_MAX_ENTRIES); +} + +#[test] +fn parses_static_winhttp_proxy_entries_for_target_scheme() { + assert_eq!( + parse_proxy_list("http=web-proxy:8080;https=secure-proxy:8443", "https"), + ParsedProxyListDecision::Proxy("http://secure-proxy:8443".to_string()) + ); + assert_eq!( + parse_proxy_list("http=web-proxy:8080 https=secure-proxy:8443", "https"), + ParsedProxyListDecision::Proxy("http://secure-proxy:8443".to_string()) + ); + assert_eq!( + parse_proxy_list("http=web-proxy:8080", "https"), + ParsedProxyListDecision::Unavailable + ); + assert_eq!( + parse_proxy_list("proxy.internal:8080", "https"), + ParsedProxyListDecision::Proxy("http://proxy.internal:8080".to_string()) + ); +} + +#[test] +fn reports_direct_and_unsupported_proxy_tokens() { + assert_eq!( + parse_proxy_list("DIRECT; PROXY proxy.internal:8080", "https"), + ParsedProxyListDecision::Direct + ); + assert_eq!( + parse_proxy_list("DIRECT", "https"), + ParsedProxyListDecision::Direct + ); + assert_eq!( + parse_proxy_list("SOCKS proxy.internal:1080", "https"), + ParsedProxyListDecision::UnsupportedScheme + ); +} + +#[test] +fn no_proxy_matches_exact_suffix_wildcard_and_port() { + let origin = RequestOrigin { + scheme: "https".to_string(), + host: "auth.openai.com".to_string(), + port: 443, + }; + assert!(no_proxy_matches_origin("auth.openai.com", &origin)); + assert!(!no_proxy_matches_origin("openai.com", &origin)); + assert!(no_proxy_matches_origin(".openai.com", &origin)); + assert!(no_proxy_matches_origin("*.openai.com", &origin)); + assert!(no_proxy_matches_origin("auth.openai.com:443", &origin)); + assert!(!no_proxy_matches_origin("auth.openai.com:8443", &origin)); +} + +#[test] +fn system_proxy_cache_key_preserves_url_specific_pac_decisions() { + let request_url = "https://auth.openai.com/oauth/token?access_token=secret"; + let cache_key = system_proxy_cache_key(request_url); + + assert_ne!( + cache_key, + system_proxy_cache_key("https://auth.openai.com/oauth/revoke") + ); + assert!(!cache_key.contains(request_url)); +} diff --git a/codex-rs/http-client/src/request.rs b/codex-rs/http-client/src/request.rs new file mode 100644 index 0000000000000000000000000000000000000000..b21ce6280e5fb243e56149c06938b797035433a8 --- /dev/null +++ b/codex-rs/http-client/src/request.rs @@ -0,0 +1,333 @@ +use bytes::Bytes; +use http::HeaderMap; +use http::HeaderValue; +use http::Method; +use serde::Serialize; +use serde_json::Value; +use std::time::Duration; + +/// A JSON request body serialized once into reference-counted bytes. +/// +/// Clones share the encoded allocation. Internally, the body can also hold the +/// final compressed wire bytes while retaining the original JSON only when +/// request-body trace logging is enabled. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct EncodedJsonBody { + bytes: Bytes, + trace_bytes: Option, + prepared: bool, +} + +impl EncodedJsonBody { + /// Serializes `value` into a reusable JSON body. + pub fn encode(value: &T) -> Result { + serde_json::to_vec(value).map(|bytes| Self { + bytes: Bytes::from(bytes), + trace_bytes: None, + prepared: false, + }) + } + + /// Returns the encoded bytes currently stored by this body. + pub fn as_bytes(&self) -> &[u8] { + &self.bytes + } + + pub(crate) fn trace_bytes(&self) -> &[u8] { + self.trace_bytes.as_ref().unwrap_or(&self.bytes) + } +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum RequestCompression { + #[default] + None, + Zstd, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RequestBody { + Json(Value), + EncodedJson(EncodedJsonBody), + Raw(Bytes), +} + +impl RequestBody { + pub fn json(&self) -> Option<&Value> { + match self { + Self::Json(value) => Some(value), + Self::EncodedJson(_) | Self::Raw(_) => None, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PreparedRequestBody { + pub headers: HeaderMap, + pub body: Option, +} + +impl PreparedRequestBody { + pub fn body_bytes(&self) -> Bytes { + self.body.clone().unwrap_or_default() + } +} + +#[derive(Debug, Clone)] +pub struct Request { + pub method: Method, + pub url: String, + pub headers: HeaderMap, + pub body: Option, + pub compression: RequestCompression, + pub timeout: Option, + /// Maximum accepted response-body bytes, including streaming and error bodies. + /// `None` leaves the response unbounded. + pub response_body_limit_bytes: Option, +} + +impl Request { + pub fn new(method: Method, url: String) -> Self { + Self { + method, + url, + headers: HeaderMap::new(), + body: None, + compression: RequestCompression::None, + timeout: None, + response_body_limit_bytes: None, + } + } + + pub fn with_json(mut self, body: &T) -> Self { + self.body = serde_json::to_value(body).ok().map(RequestBody::Json); + self + } + + pub fn with_raw_body(mut self, body: impl Into) -> Self { + self.body = Some(RequestBody::Raw(body.into())); + self + } + + pub fn with_compression(mut self, compression: RequestCompression) -> Self { + self.compression = compression; + self + } + + /// Prepares the body once and stores the exact bytes that will be sent. + /// + /// Cloning the returned request shares the body bytes, so retry attempts do + /// not repeat JSON serialization or compression. Request-signing auth also + /// sees the same final headers and bytes that the transport will send. + pub fn into_prepared(mut self) -> Result { + let is_json = matches!( + self.body, + Some(RequestBody::Json(_) | RequestBody::EncodedJson(_)) + ); + let trace_bytes = if self.compression != RequestCompression::None + && tracing::enabled!(target: "codex_http_client::transport", tracing::Level::TRACE) + { + match self.body.as_ref() { + Some(RequestBody::Json(body)) => Some(Bytes::from( + serde_json::to_vec(body).map_err(|err| err.to_string())?, + )), + Some(RequestBody::EncodedJson(body)) => Some(body.bytes.clone()), + Some(RequestBody::Raw(_)) | None => None, + } + } else { + None + }; + let prepared = self.prepare_body_for_send()?; + self.headers = prepared.headers; + self.body = match (is_json, prepared.body) { + (true, Some(bytes)) => Some(RequestBody::EncodedJson(EncodedJsonBody { + bytes, + trace_bytes, + prepared: true, + })), + (false, Some(body)) => Some(RequestBody::Raw(body)), + (_, None) => None, + }; + self.compression = RequestCompression::None; + Ok(self) + } + + /// Convert the request body into the exact bytes that will be sent. + /// + /// Auth schemes such as AWS SigV4 need to sign the final body bytes, including + /// compression and content headers. Calling this method does not mutate the + /// request. + pub fn prepare_body_for_send(&self) -> Result { + let headers = self.headers.clone(); + match self.body.as_ref() { + Some(RequestBody::Raw(raw_body)) => { + if self.compression != RequestCompression::None { + return Err("request compression cannot be used with raw bodies".to_string()); + } + Ok(PreparedRequestBody { + headers, + body: Some(raw_body.clone()), + }) + } + Some(RequestBody::Json(body)) => { + let body = EncodedJsonBody::encode(body).map_err(|err| err.to_string())?; + self.prepare_encoded_json(headers, &body) + } + Some(RequestBody::EncodedJson(body)) => self.prepare_encoded_json(headers, body), + None => Ok(PreparedRequestBody { + headers, + body: None, + }), + } + } + + fn prepare_encoded_json( + &self, + mut headers: HeaderMap, + body: &EncodedJsonBody, + ) -> Result { + if body.prepared { + return Ok(PreparedRequestBody { + headers, + body: Some(body.bytes.clone()), + }); + } + + let bytes = if self.compression != RequestCompression::None { + if headers.contains_key(http::header::CONTENT_ENCODING) { + return Err( + "request compression was requested but content-encoding is already set" + .to_string(), + ); + } + + let pre_compression_bytes = body.bytes.len(); + let compression_start = std::time::Instant::now(); + let (compressed, content_encoding) = match self.compression { + RequestCompression::None => unreachable!("guarded by compression != None"), + RequestCompression::Zstd => ( + zstd::stream::encode_all(std::io::Cursor::new(body.as_bytes()), 3) + .map_err(|err| err.to_string())?, + HeaderValue::from_static("zstd"), + ), + }; + let post_compression_bytes = compressed.len(); + let compression_duration = compression_start.elapsed(); + + headers.insert(http::header::CONTENT_ENCODING, content_encoding); + + tracing::debug!( + pre_compression_bytes, + post_compression_bytes, + compression_duration_ms = compression_duration.as_millis(), + "Compressed request body with zstd" + ); + + Bytes::from(compressed) + } else { + body.bytes.clone() + }; + + if !headers.contains_key(http::header::CONTENT_TYPE) { + headers.insert( + http::header::CONTENT_TYPE, + HeaderValue::from_static("application/json"), + ); + } + + Ok(PreparedRequestBody { + headers, + body: Some(bytes), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use http::HeaderValue; + use pretty_assertions::assert_eq; + use serde_json::json; + + #[test] + fn prepare_body_for_send_serializes_json_and_sets_content_type() { + let request = Request::new(Method::POST, "https://example.com/v1/responses".to_string()) + .with_json(&json!({"model": "test-model"})); + + let prepared = request + .prepare_body_for_send() + .expect("body should prepare"); + + assert_eq!( + prepared.body, + Some(Bytes::from_static(br#"{"model":"test-model"}"#)) + ); + assert_eq!( + prepared + .headers + .get(http::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()), + Some("application/json") + ); + assert_eq!( + request.body, + Some(RequestBody::Json(json!({"model": "test-model"}))) + ); + assert_eq!(request.compression, RequestCompression::None); + } + + #[test] + fn prepare_body_for_send_rejects_existing_content_encoding_when_compressing() { + let mut request = + Request::new(Method::POST, "https://example.com/v1/responses".to_string()) + .with_json(&json!({"model": "test-model"})) + .with_compression(RequestCompression::Zstd); + request.headers.insert( + http::header::CONTENT_ENCODING, + HeaderValue::from_static("gzip"), + ); + + let err = request + .prepare_body_for_send() + .expect_err("conflicting content-encoding should fail"); + + assert_eq!( + err, + "request compression was requested but content-encoding is already set" + ); + } + + #[test] + fn into_prepared_stores_compressed_body_for_reuse() { + let body = + EncodedJsonBody::encode(&json!({"model": "test-model"})).expect("JSON should encode"); + let mut request = + Request::new(Method::POST, "https://example.com/v1/responses".to_string()) + .with_compression(RequestCompression::Zstd); + request.body = Some(RequestBody::EncodedJson(body)); + let request = request.into_prepared().expect("body should prepare"); + let Some(RequestBody::EncodedJson(body)) = request.body.as_ref() else { + panic!("expected an encoded JSON body"); + }; + let decompressed = zstd::stream::decode_all(std::io::Cursor::new(body.as_bytes())) + .expect("body should decompress"); + + assert_eq!(decompressed, br#"{"model":"test-model"}"#); + assert_eq!(request.compression, RequestCompression::None); + assert_eq!( + request.headers.get(http::header::CONTENT_ENCODING), + Some(&HeaderValue::from_static("zstd")) + ); + assert_eq!( + request.headers.get(http::header::CONTENT_TYPE), + Some(&HeaderValue::from_static("application/json")) + ); + } +} + +#[derive(Debug, Clone)] +pub struct Response { + pub status: http::StatusCode, + pub headers: HeaderMap, + pub body: Bytes, +} diff --git a/codex-rs/http-client/src/route_aware_client_pool.rs b/codex-rs/http-client/src/route_aware_client_pool.rs new file mode 100644 index 0000000000000000000000000000000000000000..29705e44eb66f235adee5a8932264ee778f98ced --- /dev/null +++ b/codex-rs/http-client/src/route_aware_client_pool.rs @@ -0,0 +1,815 @@ +use std::collections::HashMap; +use std::fmt; +use std::future::Future; +use std::io; +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; + +use bytes::Bytes; +use futures::TryStream; +use http::HeaderMap; +use http::HeaderName; +use http::HeaderValue; +use http::Method; +use http::StatusCode; +use http::header::CONTENT_TYPE; +use http::header::PROXY_AUTHORIZATION; +use reqwest::IntoUrl; +use serde::Serialize; + +use crate::BuildRouteAwareHttpClientError; +use crate::ClientRouteClass; +use crate::HttpClient; +use crate::HttpClientBuilder; +use crate::HttpClientFactory; +use crate::OutboundProxyPolicy; +use crate::OutboundProxyRoute; +use crate::RouteFailureClass; +use crate::route_aware_redirect::MAX_REDIRECTS; +use crate::route_aware_redirect::insert_referer; +use crate::route_aware_redirect::is_redirect; +use crate::route_aware_redirect::redirect_request; +use crate::route_aware_redirect::redirect_url; +use crate::route_aware_redirect::remove_sensitive_headers; +use crate::tls_backend_fallback::RustlsClientCache; +use crate::tls_backend_fallback::should_retry_with_rustls; + +const MAX_CACHED_ROUTES: usize = 16; + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +enum CustomCaFallback { + #[default] + Disabled, + LegacyTransportDefault, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum SelectedTlsBackend { + TransportDefault, + RustlsFallback, +} + +/// Reuses transport clients by resolved route while selecting a route for every request URL. +/// +/// Request creation stays on the pool so the URL used for PAC or system-proxy resolution cannot +/// differ from the URL that is sent. Redirects are followed through the pool as new requests, so +/// each hop gets its own route decision while connections are still reused by route. +#[derive(Clone)] +pub struct RouteAwareClientPool { + http_client_factory: HttpClientFactory, + route_class: ClientRouteClass, + client_builder: HttpClientBuilder, + custom_ca_fallback: CustomCaFallback, + clients: Arc>>, + rustls_clients: Option, +} + +impl fmt::Debug for RouteAwareClientPool { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("RouteAwareClientPool") + .field("http_client_factory", &self.http_client_factory) + .field("route_class", &self.route_class) + .finish_non_exhaustive() + } +} + +/// Error returned when selecting a route or constructing its pooled HTTP client. +#[derive(Debug, thiserror::Error)] +pub enum RouteAwareClientPoolError { + #[error("failed to resolve the outbound proxy route: {0}")] + Resolve(#[source] io::Error), + #[error(transparent)] + Build(#[from] BuildRouteAwareHttpClientError), +} + +/// Error returned while building, routing, or sending a route-aware request. +#[derive(Debug, thiserror::Error)] +pub enum RouteAwareRequestError { + #[error(transparent)] + Request(#[from] reqwest::Error), + #[error(transparent)] + Route(#[from] RouteAwareClientPoolError), + #[error("failed to build route-aware request: {0}")] + Build(String), + #[error("redirect target uses unsupported URL scheme: {0}")] + UnsupportedRedirectScheme(String), + #[error("too many redirects")] + TooManyRedirects, + #[error("route-aware request timed out")] + Timeout, +} + +impl RouteAwareRequestError { + /// Classifies transport, proxy, and certificate failures without exposing request details. + pub fn failure_class(&self) -> Option { + if self.is_timeout() { + return Some(RouteFailureClass::ConnectTimeout); + } + if self.status() == Some(StatusCode::PROXY_AUTHENTICATION_REQUIRED) { + return Some(RouteFailureClass::ProxyAuthenticationRequired); + } + if let Self::Route(RouteAwareClientPoolError::Resolve(error)) = self + && let Some(source) = error.get_ref() + && source.is::() + { + return Some(RouteFailureClass::TlsError); + } + + let mut source: Option<&(dyn std::error::Error + 'static)> = Some(self); + while let Some(error) = source { + if error.downcast_ref::().is_some() + || error.downcast_ref::().is_some() + { + return Some(RouteFailureClass::TlsError); + } + if error.to_string() == "tunnel error: proxy authorization required" { + return Some(RouteFailureClass::ProxyAuthenticationRequired); + } + source = error.source(); + } + + match self { + Self::Route(RouteAwareClientPoolError::Build( + BuildRouteAwareHttpClientError::CustomCa(_), + )) => Some(RouteFailureClass::TlsError), + Self::Route(RouteAwareClientPoolError::Build( + BuildRouteAwareHttpClientError::InvalidProxyConfig { .. }, + )) => Some(RouteFailureClass::InvalidProxyConfig), + Self::Route(RouteAwareClientPoolError::Resolve(_)) => { + Some(RouteFailureClass::ProxyResolutionUnavailable) + } + Self::Request(_) + | Self::Build(_) + | Self::UnsupportedRedirectScheme(_) + | Self::TooManyRedirects + | Self::Timeout => None, + } + } + + pub fn status(&self) -> Option { + match self { + Self::Request(error) => error.status(), + Self::Route(_) + | Self::Build(_) + | Self::UnsupportedRedirectScheme(_) + | Self::TooManyRedirects + | Self::Timeout => None, + } + } + + pub fn is_timeout(&self) -> bool { + matches!(self, Self::Timeout) || matches!(self, Self::Request(error) if error.is_timeout()) + } + + pub fn is_connect(&self) -> bool { + matches!(self, Self::Request(error) if error.is_connect()) + } + + pub fn is_body(&self) -> bool { + matches!(self, Self::Request(error) if error.is_body()) + } + + pub fn is_request(&self) -> bool { + matches!(self, Self::Request(error) if error.is_request()) + } + + /// Removes a request URL from the underlying transport error before it is logged or returned. + /// + /// Use this for requests whose URL can contain credentials, such as signed blob uploads. + pub fn without_url(self) -> Self { + match self { + Self::Request(error) => Self::Request(error.without_url()), + other => other, + } + } +} + +#[must_use = "requests are not sent unless `send` is awaited"] +pub struct RouteAwareRequestBuilder { + pool: RouteAwareClientPool, + request: Result, +} + +impl fmt::Debug for RouteAwareRequestBuilder { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + let request = self.request.as_ref().ok(); + formatter + .debug_struct("RouteAwareRequestBuilder") + .field("pool", &self.pool) + .field("method", &request.map(reqwest::Request::method)) + .field("url", &request.map(|_| "")) + .finish_non_exhaustive() + } +} + +impl RouteAwareRequestBuilder { + fn new(pool: RouteAwareClientPool, method: Method, url: U) -> Self + where + U: IntoUrl, + { + let request = url + .into_url() + .map(|url| reqwest::Request::new(method, url)) + .map_err(RouteAwareRequestError::Request); + Self { pool, request } + } + + pub fn headers(mut self, headers: HeaderMap) -> Self { + if let Ok(request) = &mut self.request { + request.headers_mut().extend(headers); + } + self + } + + pub fn header(mut self, key: K, value: V) -> Self + where + HeaderName: TryFrom, + >::Error: Into, + HeaderValue: TryFrom, + >::Error: Into, + { + if let Ok(request) = &mut self.request { + let header = HeaderName::try_from(key) + .map_err(Into::into) + .and_then(|key| { + HeaderValue::try_from(value) + .map(|value| (key, value)) + .map_err(Into::into) + }); + match header { + Ok((key, value)) => { + request.headers_mut().append(key, value); + } + Err(error) => { + self.request = Err(RouteAwareRequestError::Build(error.to_string())); + } + } + } + self + } + + /// Sets a timeout for the request as a whole. + /// + /// The budget starts before outbound-route resolution and covers selecting or constructing a + /// pooled client, establishing a connection, sending the request, and awaiting the response. + /// Use [`HttpClientBuilder::connect_timeout`] when only connection establishment should be + /// bounded. + pub fn timeout(mut self, timeout: Duration) -> Self { + if let Ok(request) = &mut self.request { + *request.timeout_mut() = Some(timeout); + } + self + } + + pub fn json(mut self, value: &T) -> Self + where + T: ?Sized + Serialize, + { + if let Ok(request) = &mut self.request { + match serde_json::to_vec(value) { + Ok(body) => { + if !request.headers().contains_key(CONTENT_TYPE) { + request + .headers_mut() + .insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); + } + *request.body_mut() = Some(body.into()); + } + Err(error) => { + self.request = Err(RouteAwareRequestError::Build(error.to_string())); + } + } + } + self + } + + pub fn body(mut self, body: B) -> Self + where + B: Into, + { + if let Ok(request) = &mut self.request { + *request.body_mut() = Some(body.into()); + } + self + } + + /// Sets a streaming request body without exposing the underlying HTTP implementation. + pub fn body_stream(mut self, stream: S) -> Self + where + S: TryStream + Send + 'static, + S::Error: Into>, + Bytes: From, + { + if let Ok(request) = &mut self.request { + *request.body_mut() = Some(reqwest::Body::wrap_stream(stream)); + } + self + } + + pub async fn send(self) -> Result { + self.pool.send(self.request?).await + } +} + +impl RouteAwareClientPool { + pub fn outbound_proxy_policy(&self) -> OutboundProxyPolicy { + self.http_client_factory.outbound_proxy_policy() + } + + /// Creates a pool with the shared default HTTP transport settings. + pub fn new(http_client_factory: HttpClientFactory, route_class: ClientRouteClass) -> Self { + Self::with_builder(http_client_factory, route_class, HttpClientBuilder::new()) + } + + /// Creates a pool that returns redirect responses without following them. + /// + /// This applies both when reqwest owns redirect handling and when the pool follows redirects + /// manually so each hop can receive its own proxy-route decision. + pub fn new_without_redirects( + http_client_factory: HttpClientFactory, + route_class: ClientRouteClass, + ) -> Self { + Self::with_builder( + http_client_factory, + route_class, + HttpClientBuilder::new().without_redirects(), + ) + } + + /// Creates a no-redirect pool without request URL or response-header diagnostics. + pub fn new_without_redirects_or_request_logging( + http_client_factory: HttpClientFactory, + route_class: ClientRouteClass, + ) -> Self { + Self::with_builder( + http_client_factory, + route_class, + HttpClientBuilder::new() + .without_redirects() + .without_request_logging(), + ) + } + + /// Creates a pool whose clients limit only connection establishment. + /// + /// The timeout applies to every client built for a resolved route, including redirect hops. + pub fn with_connect_timeout( + http_client_factory: HttpClientFactory, + route_class: ClientRouteClass, + connect_timeout: Duration, + ) -> Self { + Self::with_builder( + http_client_factory, + route_class, + HttpClientBuilder::new().connect_timeout(connect_timeout), + ) + } + + fn with_builder( + http_client_factory: HttpClientFactory, + route_class: ClientRouteClass, + client_builder: HttpClientBuilder, + ) -> Self { + Self { + http_client_factory, + route_class, + client_builder, + custom_ca_fallback: CustomCaFallback::Disabled, + clients: Arc::new(Mutex::new(HashMap::new())), + rustls_clients: None, + } + } + + /// Retries recognized TLS protocol-negotiation failures once using rustls. + /// + /// Successful fallback is remembered for each HTTPS origin and resolved outbound route. + /// Rustls clients are reused across fallback destinations that share a route, while other + /// destinations retain the existing transport-default backend. + pub fn with_tls_backend_fallback(mut self) -> Self { + self.rustls_clients = Some(RustlsClientCache::default()); + self + } + + /// Creates a pool with the shared defaults but without URL or response-header diagnostics. + pub fn new_without_request_logging( + http_client_factory: HttpClientFactory, + route_class: ClientRouteClass, + ) -> Self { + Self::with_builder( + http_client_factory, + route_class, + HttpClientBuilder::new().without_request_logging(), + ) + } + + /// Preserves the legacy custom-CA fallback for transport-default proxy routes. + /// + /// Use this only when migrating a client that already continued with system roots after a + /// custom-CA construction failure. System-proxy routes still propagate construction errors. + pub fn with_legacy_custom_ca_fallback(mut self) -> Self { + self.custom_ca_fallback = CustomCaFallback::LegacyTransportDefault; + self + } + + /// Creates a pool that retains the Cloudflare cookies required by ChatGPT endpoints. + pub fn with_chatgpt_cloudflare_cookies( + http_client_factory: HttpClientFactory, + route_class: ClientRouteClass, + ) -> Self { + Self::with_builder( + http_client_factory, + route_class, + HttpClientBuilder::new().with_chatgpt_cloudflare_cookie_store(), + ) + } + + /// Creates a no-redirect pool that retains the Cloudflare cookies required by ChatGPT + /// endpoints. + pub fn with_chatgpt_cloudflare_cookies_without_redirects( + http_client_factory: HttpClientFactory, + route_class: ClientRouteClass, + ) -> Self { + Self::with_builder( + http_client_factory, + route_class, + HttpClientBuilder::new() + .with_chatgpt_cloudflare_cookie_store() + .without_redirects(), + ) + } + + /// Creates a no-redirect ChatGPT Cloudflare-cookie pool without request diagnostics. + pub fn with_chatgpt_cloudflare_cookies_without_redirects_or_request_logging( + http_client_factory: HttpClientFactory, + route_class: ClientRouteClass, + ) -> Self { + Self::with_builder( + http_client_factory, + route_class, + HttpClientBuilder::new() + .with_chatgpt_cloudflare_cookie_store() + .without_redirects() + .without_request_logging(), + ) + } + + /// Creates a ChatGPT Cloudflare-cookie pool without URL or response-header diagnostics. + pub fn with_chatgpt_cloudflare_cookies_without_request_logging( + http_client_factory: HttpClientFactory, + route_class: ClientRouteClass, + ) -> Self { + Self::with_builder( + http_client_factory, + route_class, + HttpClientBuilder::new() + .with_chatgpt_cloudflare_cookie_store() + .without_request_logging(), + ) + } + + pub fn get(&self, url: U) -> RouteAwareRequestBuilder + where + U: IntoUrl, + { + self.request(Method::GET, url) + } + + pub fn post(&self, url: U) -> RouteAwareRequestBuilder + where + U: IntoUrl, + { + self.request(Method::POST, url) + } + + pub fn put(&self, url: U) -> RouteAwareRequestBuilder + where + U: IntoUrl, + { + self.request(Method::PUT, url) + } + + pub fn delete(&self, url: U) -> RouteAwareRequestBuilder + where + U: IntoUrl, + { + self.request(Method::DELETE, url) + } + + pub fn request(&self, method: Method, url: U) -> RouteAwareRequestBuilder + where + U: IntoUrl, + { + RouteAwareRequestBuilder::new(self.clone(), method, url) + } + + async fn send( + &self, + request: reqwest::Request, + ) -> Result { + let http_client_factory = self.http_client_factory.clone(); + self.send_with_resolver(request, move |request_url| { + let http_client_factory = http_client_factory.clone(); + async move { + http_client_factory + .resolve_proxy_route_async(request_url) + .await + } + }) + .await + } + + async fn send_with_resolver( + &self, + mut request: reqwest::Request, + resolve_route: F, + ) -> Result + where + F: Fn(String) -> Fut, + Fut: Future>, + { + let request_method = request.method().clone(); + let request_url = request.url().to_string(); + let follows_redirects_manually = self.follows_redirects_manually(); + let timeout_deadline = request + .timeout() + .copied() + .map(|timeout| tokio::time::Instant::now() + timeout); + let mut redirects = 0; + let mut previous_route = None; + loop { + let current_url = request.url().clone(); + let (current_route, client, selected_tls_backend) = match timeout_deadline { + Some(timeout_deadline) => tokio::time::timeout_at( + timeout_deadline, + self.client_for_url_with_resolver(current_url.as_str(), &resolve_route), + ) + .await + .map_err(|_| RouteAwareRequestError::Timeout)??, + None => { + self.client_for_url_with_resolver(current_url.as_str(), &resolve_route) + .await? + } + }; + if previous_route + .as_ref() + .is_some_and(|previous_route| previous_route != ¤t_route) + { + request.headers_mut().remove(PROXY_AUTHORIZATION); + } + previous_route = Some(current_route.clone()); + if let Some(timeout_deadline) = timeout_deadline { + let remaining = timeout_deadline + .checked_duration_since(tokio::time::Instant::now()) + .ok_or(RouteAwareRequestError::Timeout)?; + if remaining.is_zero() { + return Err(RouteAwareRequestError::Timeout); + } + *request.timeout_mut() = Some(remaining); + } + let method = request.method().clone(); + let headers = request.headers().clone(); + let version = request.version(); + let timeout = request.timeout().copied(); + let replay = request.try_clone(); + let execute_request = async { + if follows_redirects_manually { + client.execute_without_request_logging(request).await + } else { + client.execute(request).await + } + }; + let response = match match timeout_deadline { + Some(timeout_deadline) => { + tokio::time::timeout_at(timeout_deadline, execute_request) + .await + .map_err(|_| RouteAwareRequestError::Timeout)? + } + None => execute_request.await, + } { + Ok(response) => response, + Err(error) => { + let result = self + .retry_with_rustls( + ¤t_url, + ¤t_route, + selected_tls_backend, + replay.as_ref(), + error, + timeout_deadline, + ) + .await; + + if follows_redirects_manually + && let Err(RouteAwareRequestError::Request(error)) = &result + { + client.log_error_summary(&request_method, &request_url, error); + } + + result? + } + }; + let status = response.status(); + if !follows_redirects_manually || !is_redirect(status) { + if follows_redirects_manually { + client.log_response(&request_method, &request_url, &response); + } + return Ok(response); + } + let Some(next_url) = redirect_url(&response) else { + if follows_redirects_manually { + client.log_response(&request_method, &request_url, &response); + } + return Ok(response); + }; + let Some(mut next_request) = + redirect_request(status, method, headers, version, timeout, replay, next_url) + else { + if follows_redirects_manually { + client.log_response(&request_method, &request_url, &response); + } + return Ok(response); + }; + let next_request_url = next_request.url().clone(); + if !matches!(next_request_url.scheme(), "http" | "https") { + return Err(RouteAwareRequestError::UnsupportedRedirectScheme( + next_request_url.scheme().to_string(), + )); + } + if redirects >= MAX_REDIRECTS { + return Err(RouteAwareRequestError::TooManyRedirects); + } + remove_sensitive_headers(next_request.headers_mut(), ¤t_url, &next_request_url); + insert_referer(next_request.headers_mut(), ¤t_url, &next_request_url); + request = next_request; + redirects += 1; + } + } + + async fn retry_with_rustls( + &self, + current_url: &reqwest::Url, + current_route: &OutboundProxyRoute, + selected_tls_backend: SelectedTlsBackend, + replay: Option<&reqwest::Request>, + error: reqwest::Error, + timeout_deadline: Option, + ) -> Result { + let Some(rustls_clients) = self.rustls_clients.as_ref() else { + return Err(error.into()); + }; + + if current_url.scheme() != "https" + || selected_tls_backend == SelectedTlsBackend::RustlsFallback + || !should_retry_with_rustls(&error) + { + return Err(error.into()); + } + + let Some(mut retry_request) = replay.and_then(reqwest::Request::try_clone) else { + return Err(error.into()); + }; + + let fallback_client = match rustls_clients.client_for_route(current_route) { + Some(client) => client, + None => self.rustls_client_for_route(current_route)?, + }; + + if let Some(timeout_deadline) = timeout_deadline { + let remaining = timeout_deadline + .checked_duration_since(tokio::time::Instant::now()) + .ok_or(RouteAwareRequestError::Timeout)?; + if remaining.is_zero() { + return Err(RouteAwareRequestError::Timeout); + } + *retry_request.timeout_mut() = Some(remaining); + } + + let execute_retry = async { + if self.follows_redirects_manually() { + fallback_client + .execute_without_request_logging(retry_request) + .await + } else { + fallback_client.execute(retry_request).await + } + }; + let response = match timeout_deadline { + Some(timeout_deadline) => tokio::time::timeout_at(timeout_deadline, execute_retry) + .await + .map_err(|_| RouteAwareRequestError::Timeout)?, + None => execute_retry.await, + }?; + + rustls_clients.remember(current_url, current_route, fallback_client); + tracing::info!( + event.name = "codex.http_client.tls_backend_fallback", + "HTTP client switched to rustls after a TLS protocol negotiation failure" + ); + + Ok(response) + } + + async fn client_for_url_with_resolver( + &self, + request_url: &str, + resolve_route: F, + ) -> Result<(OutboundProxyRoute, HttpClient, SelectedTlsBackend), RouteAwareClientPoolError> + where + F: FnOnce(String) -> Fut, + Fut: Future>, + { + let route = resolve_route(request_url.to_string()) + .await + .map_err(RouteAwareClientPoolError::Resolve)?; + if let Some(rustls_clients) = self.rustls_clients.as_ref() + && let Ok(url) = reqwest::Url::parse(request_url) + && rustls_clients.requires_rustls(&url, &route) + && let Some(client) = rustls_clients.client_for_route(&route) + { + return Ok((route, client, SelectedTlsBackend::RustlsFallback)); + } + let clients = match self.clients.lock() { + Ok(clients) => clients, + Err(error) => panic!("route-aware client cache lock should not be poisoned: {error}"), + }; + if let Some(client) = clients.get(&route) { + return Ok((route, client.clone(), SelectedTlsBackend::TransportDefault)); + } + drop(clients); + + let client_builder = if self.follows_redirects_manually() { + self.client_builder.clone().without_redirects() + } else { + self.client_builder.clone() + }; + #[expect( + deprecated, + reason = "explicitly opted-in pools preserve the legacy custom-CA fallback" + )] + let client = match ( + self.http_client_factory.outbound_proxy_policy(), + self.custom_ca_fallback, + ) { + (OutboundProxyPolicy::ReqwestDefault, CustomCaFallback::LegacyTransportDefault) => { + client_builder.build_with_transport_default_proxy_and_custom_ca_fallback() + } + (OutboundProxyPolicy::ReqwestDefault, CustomCaFallback::Disabled) + | (OutboundProxyPolicy::RespectSystemProxy, CustomCaFallback::Disabled) + | (OutboundProxyPolicy::RespectSystemProxy, CustomCaFallback::LegacyTransportDefault) => { + client_builder.build_for_resolved_route( + &self.http_client_factory, + self.route_class, + &route, + )? + } + }; + let mut clients = match self.clients.lock() { + Ok(clients) => clients, + Err(error) => panic!("route-aware client cache lock should not be poisoned: {error}"), + }; + if let Some(existing_client) = clients.get(&route) { + return Ok(( + route, + existing_client.clone(), + SelectedTlsBackend::TransportDefault, + )); + } + if clients.len() >= MAX_CACHED_ROUTES + && let Some(route_to_evict) = clients.keys().next().cloned() + { + clients.remove(&route_to_evict); + } + clients.insert(route.clone(), client.clone()); + Ok((route, client, SelectedTlsBackend::TransportDefault)) + } + + fn follows_redirects_manually(&self) -> bool { + self.client_builder.follows_redirects() + && (self.http_client_factory.outbound_proxy_policy() + == OutboundProxyPolicy::RespectSystemProxy + || self.rustls_clients.is_some()) + } + + fn rustls_client_for_route( + &self, + route: &OutboundProxyRoute, + ) -> Result { + let mut client_builder = self.client_builder.clone().with_rustls_tls(); + if self.follows_redirects_manually() { + client_builder = client_builder.without_redirects(); + } + client_builder + .build_for_resolved_route(&self.http_client_factory, self.route_class, route) + .map_err(Into::into) + } +} + +#[cfg(test)] +#[path = "route_aware_client_pool_tests.rs"] +mod tests; + +#[cfg(test)] +#[path = "route_aware_tls_fallback_tests.rs"] +mod tls_fallback_tests; diff --git a/codex-rs/http-client/src/route_aware_client_pool_tests.rs b/codex-rs/http-client/src/route_aware_client_pool_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..1be6042edb0b7c130661ced5337fae4448d0db16 --- /dev/null +++ b/codex-rs/http-client/src/route_aware_client_pool_tests.rs @@ -0,0 +1,899 @@ +use std::collections::HashMap; +use std::io; +use std::io::Read; +use std::io::Write; +use std::net::TcpListener; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; +use std::time::Instant; + +use bytes::Bytes; +use futures::stream; +use pretty_assertions::assert_eq; +use tracing_subscriber::Layer; +use tracing_subscriber::layer::SubscriberExt; + +use super::*; +use crate::OutboundProxyPolicy; + +#[tokio::test] +async fn request_failures_classify_real_untrusted_certificate_handshakes() { + codex_utils_rustls_provider::ensure_rustls_crypto_provider(); + let certificate = rcgen::generate_simple_self_signed(vec!["localhost".to_string()]) + .expect("self-signed certificate should generate"); + let private_key = rustls::pki_types::PrivateKeyDer::Pkcs8( + rustls::pki_types::PrivatePkcs8KeyDer::from(certificate.signing_key.serialize_der()), + ); + let configuration = Arc::new( + rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(vec![certificate.cert.der().clone()], private_key) + .expect("TLS server should be configured"), + ); + let listener = TcpListener::bind("127.0.0.1:0").expect("TLS server should bind"); + let address = listener + .local_addr() + .expect("TLS server should have an address"); + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("TLS server should accept"); + stream + .set_read_timeout(Some(Duration::from_secs(3))) + .expect("TLS handshake timeout"); + let mut connection = rustls::ServerConnection::new(configuration) + .expect("TLS server connection should be created"); + let _ = connection.complete_io(&mut stream); + }); + let pool = RouteAwareClientPool::new_without_request_logging( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Api, + ); + + let request = pool + .get(format!("https://localhost:{}/", address.port())) + .timeout(Duration::from_secs(3)) + .request + .expect("TLS request should build"); + let error = pool + .send_with_resolver(request, |_| async { Ok(OutboundProxyRoute::Direct) }) + .await + .expect_err("self-signed server certificate must not be trusted"); + drop(std::net::TcpStream::connect(address)); + server.join().expect("TLS server should finish"); + + assert_eq!( + error.failure_class(), + Some(RouteFailureClass::TlsError), + "unexpected certificate error: {error:?}" + ); +} + +#[tokio::test] +async fn request_failures_classify_https_proxy_authentication_challenges() { + let (address, proxy) = spawn_response_server(vec![ + "HTTP/1.1 407 Proxy Authentication Required\r\n\ + Proxy-Authenticate: Basic realm=\"codex\"\r\n\ + Content-Length: 0\r\n\ + Connection: close\r\n\r\n" + .to_string(), + ]); + let pool = RouteAwareClientPool::new_without_request_logging( + HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy), + ClientRouteClass::Api, + ); + let mut request = reqwest::Request::new( + Method::GET, + reqwest::Url::parse("https://example.com/").expect("request URL should parse"), + ); + *request.timeout_mut() = Some(Duration::from_secs(3)); + + let error = pool + .send_with_resolver(request, move |_| async move { + Ok(OutboundProxyRoute::Proxy { + url: format!("http://{address}"), + no_proxy: None, + }) + }) + .await + .expect_err("HTTPS proxy challenge should reject the CONNECT request"); + let requests = proxy.join().expect("proxy fixture should finish"); + + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("CONNECT example.com:443 HTTP/1.1\r\n")); + assert_eq!( + error.failure_class(), + Some(RouteFailureClass::ProxyAuthenticationRequired), + "unexpected HTTPS proxy error: {error:?}" + ); +} + +#[test] +fn request_builder_debug_redacts_url_secrets() { + let pool = RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Api, + ); + let request = pool.get( + "https://username:password@private.example/secret-path?sig=query-secret#fragment-secret", + ); + + assert_eq!( + format!("{request:?}"), + concat!( + "RouteAwareRequestBuilder { pool: RouteAwareClientPool { ", + "http_client_factory: HttpClientFactory { outbound_proxy_policy: ReqwestDefault, ", + "network_policy: NetworkPolicy { managed: false, .. } }, ", + "route_class: Api, .. }, method: Some(GET), ", + "url: Some(\"\"), .. }" + ) + ); +} + +#[tokio::test] +async fn streams_request_bodies_without_exposing_reqwest_body() { + let (address, server) = spawn_response_server(vec![ + "HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n".to_string(), + ]); + let pool = RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Api, + ); + let response = pool + .put(format!("http://{address}/upload")) + .header(http::header::CONTENT_LENGTH, /*value*/ 5) + .body_stream(stream::iter(vec![Ok::<_, io::Error>(Bytes::from_static( + b"hello", + ))])) + .send() + .await + .expect("streaming request should succeed"); + + assert_eq!(response.status(), StatusCode::OK); + let requests = server.join().expect("response server should finish"); + assert_eq!(requests.len(), 1); + assert!(requests[0].ends_with("\r\n\r\nhello")); +} + +#[tokio::test] +async fn legacy_custom_ca_fallback_is_limited_to_reqwest_default() { + const CHILD_POLICY_ENV: &str = "CODEX_HTTP_CLIENT_POOL_INVALID_CA_TEST_POLICY"; + + let Ok(policy_name) = std::env::var(CHILD_POLICY_ENV) else { + let temp_dir = tempfile::tempdir().expect("temporary directory should be created"); + let invalid_ca_path = temp_dir.path().join("invalid-ca.pem"); + std::fs::write(&invalid_ca_path, "not a PEM certificate") + .expect("invalid CA fixture should be written"); + + for ca_env in ["CODEX_CA_CERTIFICATE", "SSL_CERT_FILE"] { + for policy_name in ["reqwest-default", "respect-system-proxy"] { + let output = std::process::Command::new( + std::env::current_exe().expect("test executable should be available"), + ) + .arg("--exact") + .arg("route_aware_client_pool::tests::legacy_custom_ca_fallback_is_limited_to_reqwest_default") + .arg("--nocapture") + .env_remove("CODEX_CA_CERTIFICATE") + .env_remove("SSL_CERT_FILE") + .env(ca_env, &invalid_ca_path) + .env(CHILD_POLICY_ENV, policy_name) + .output() + .expect("isolated CA subprocess should run"); + + assert!( + output.status.success(), + "{policy_name} failed with invalid {ca_env}\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr), + ); + } + } + return; + }; + + let outbound_proxy_policy = match policy_name.as_str() { + "reqwest-default" => OutboundProxyPolicy::ReqwestDefault, + "respect-system-proxy" => OutboundProxyPolicy::RespectSystemProxy, + _ => panic!("unexpected test proxy policy: {policy_name}"), + }; + let pool = RouteAwareClientPool::with_chatgpt_cloudflare_cookies( + HttpClientFactory::new(outbound_proxy_policy), + ClientRouteClass::Other, + ) + .with_legacy_custom_ca_fallback(); + + match outbound_proxy_policy { + OutboundProxyPolicy::ReqwestDefault => { + let (address, server) = spawn_response_server(vec![ + "HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n".to_string(), + ]); + let response = pool + .get(format!("http://{address}/update")) + .send() + .await + .expect("default-routed request should fall back to system roots"); + + assert_eq!(response.status(), StatusCode::OK); + let requests = server.join().expect("response server should finish"); + assert_eq!(requests.len(), 1); + } + OutboundProxyPolicy::RespectSystemProxy => { + let error = pool + .client_for_url_with_resolver("http://127.0.0.1/update", |_| async { + Ok(OutboundProxyRoute::Direct) + }) + .await + .expect_err("system-proxy routes should reject invalid custom CAs"); + + assert!(matches!(error, RouteAwareClientPoolError::Build(_))); + } + } +} + +#[tokio::test] +async fn without_url_redacts_transport_error_urls() { + let listener = TcpListener::bind(("127.0.0.1", 0)).expect("listener should bind"); + let address = listener.local_addr().expect("listener should have address"); + drop(listener); + let secret = "signed-secret"; + let error = reqwest::Client::new() + .get(format!("http://{address}/upload?sig={secret}")) + .send() + .await + .expect_err("closed listener should reject request"); + let error = RouteAwareRequestError::from(error).without_url(); + + assert!(!error.to_string().contains(secret)); +} + +#[tokio::test] +async fn forwards_exact_urls_and_caches_clients_by_resolved_route() { + let pool = RouteAwareClientPool::with_builder( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Api, + HttpClientBuilder::new(), + ); + + let direct_url = "https://example.com/first?target=direct"; + let same_route_url = "https://example.com/second?target=direct%202"; + let proxy_url = "https://example.com/third?target=proxy"; + let resolver = FakeRouteResolver::new(HashMap::from([ + (direct_url.to_string(), OutboundProxyRoute::Direct), + (same_route_url.to_string(), OutboundProxyRoute::Direct), + ( + proxy_url.to_string(), + OutboundProxyRoute::Proxy { + url: "http://proxy.example".to_string(), + no_proxy: None, + }, + ), + ])); + + resolve_with(&pool, &resolver, direct_url) + .await + .expect("first client should build"); + resolve_with(&pool, &resolver, same_route_url) + .await + .expect("second client should reuse the route"); + resolve_with(&pool, &resolver, proxy_url) + .await + .expect("proxy client should build separately"); + + assert_eq!(pool.clients.lock().expect("client cache lock").len(), 2); + assert_eq!( + resolver.observed_urls(), + vec![ + direct_url.to_string(), + same_route_url.to_string(), + proxy_url.to_string(), + ] + ); +} + +#[tokio::test] +async fn cached_tls_backend_only_changes_its_destination_and_route() { + let pool = RouteAwareClientPool::with_builder( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Other, + HttpClientBuilder::new(), + ) + .with_tls_backend_fallback(); + let rustls_url = "https://mcp.example.com/first"; + let native_url = "https://another.example.com/first"; + let proxied_url = "https://mcp.example.com/proxied"; + let proxy = OutboundProxyRoute::Proxy { + url: "http://proxy.example.com".to_string(), + no_proxy: None, + }; + let resolver = FakeRouteResolver::new(HashMap::from([ + (rustls_url.to_string(), OutboundProxyRoute::Direct), + (native_url.to_string(), OutboundProxyRoute::Direct), + (proxied_url.to_string(), proxy), + ])); + let fallback_client = HttpClientBuilder::new() + .with_rustls_tls() + .build_direct() + .expect("rustls client should build without proxy autodiscovery"); + pool.rustls_clients + .as_ref() + .expect("TLS fallback cache") + .remember( + &reqwest::Url::parse(rustls_url).expect("valid rustls URL"), + &OutboundProxyRoute::Direct, + fallback_client, + ); + + resolve_with(&pool, &resolver, rustls_url) + .await + .expect("remembered destination should build the rustls client"); + assert_eq!(pool.clients.lock().expect("client cache lock").len(), 0); + + resolve_with(&pool, &resolver, native_url) + .await + .expect("another destination should retain its native TLS client"); + resolve_with(&pool, &resolver, proxied_url) + .await + .expect("another proxy route should retain its native TLS client"); + assert_eq!(pool.clients.lock().expect("client cache lock").len(), 2); +} + +#[tokio::test] +async fn tls_fallback_pool_reselects_routes_for_each_redirect_hop() { + let (address, server) = spawn_response_server(vec![ + "HTTP/1.1 302 Found\r\nLocation: /final\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + .to_string(), + "HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok".to_string(), + ]); + let initial_url = format!("http://{address}/start"); + let final_url = format!("http://{address}/final"); + let resolver = FakeRouteResolver::new(HashMap::from([ + (initial_url.clone(), OutboundProxyRoute::Direct), + (final_url.clone(), OutboundProxyRoute::Direct), + ])); + let pool = RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Other, + ) + .with_tls_backend_fallback(); + let request = reqwest::Request::new( + Method::GET, + reqwest::Url::parse(&initial_url).expect("valid initial URL"), + ); + + let response = pool + .send_with_resolver(request, |url| resolver.resolve(url)) + .await + .expect("fallback-enabled client should follow redirects"); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(resolver.observed_urls(), vec![initial_url, final_url]); + assert_eq!( + server.join().expect("redirect server should finish").len(), + 2 + ); +} + +#[tokio::test] +async fn reqwest_default_route_preserves_transport_redirects() { + let listener = TcpListener::bind(("127.0.0.1", 0)).expect("redirect listener should bind"); + let address = listener + .local_addr() + .expect("redirect listener should have an address"); + listener + .set_nonblocking(true) + .expect("redirect listener should become nonblocking"); + let server = std::thread::spawn(move || { + let mut request_lines = Vec::new(); + for response in [ + "HTTP/1.1 302 Found\r\nLocation: /final\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + "HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok", + ] { + let deadline = Instant::now() + Duration::from_secs(2); + let (mut stream, _) = loop { + match listener.accept() { + Ok(connection) => break connection, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => { + assert!( + Instant::now() < deadline, + "redirect server should receive the next request" + ); + std::thread::sleep(Duration::from_millis(10)); + } + Err(error) => panic!("redirect server should accept: {error}"), + } + }; + // Accepted sockets inherit the listener's nonblocking mode on macOS. + stream + .set_nonblocking(false) + .expect("redirect stream should become blocking"); + stream + .set_read_timeout(Some(Duration::from_secs(2))) + .expect("redirect stream should get a read timeout"); + let mut buffer = [0_u8; 1024]; + let size = stream + .read(&mut buffer) + .expect("redirect server should read request"); + let request = String::from_utf8_lossy(&buffer[..size]); + request_lines.push( + request + .lines() + .next() + .expect("request should have a request line") + .to_string(), + ); + stream + .write_all(response.as_bytes()) + .expect("redirect server should write response"); + } + request_lines + }); + let pool = RouteAwareClientPool::with_builder( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Api, + HttpClientBuilder::new(), + ); + let initial_url = format!("http://{address}/start"); + let request = reqwest::Request::new( + Method::GET, + reqwest::Url::parse(&initial_url).expect("request URL should parse"), + ); + + let response = pool + .send_with_resolver(request, |_| async { Ok(OutboundProxyRoute::Direct) }) + .await + .expect("default-routed request should follow redirect"); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response.url().as_str(), format!("http://{address}/final")); + assert_eq!( + server.join().expect("redirect server should finish"), + vec![ + "GET /start HTTP/1.1".to_string(), + "GET /final HTTP/1.1".to_string(), + ] + ); +} + +#[tokio::test] +async fn no_redirect_pool_returns_redirect_response() { + for outbound_proxy_policy in [ + OutboundProxyPolicy::ReqwestDefault, + OutboundProxyPolicy::RespectSystemProxy, + ] { + let (address, server) = spawn_response_server(vec![ + "HTTP/1.1 302 Found\r\nLocation: /final\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + .to_string(), + ]); + let pool = RouteAwareClientPool::new_without_redirects( + HttpClientFactory::new(outbound_proxy_policy), + ClientRouteClass::Api, + ); + let initial_url = format!("http://{address}/start"); + let request = reqwest::Request::new( + Method::GET, + reqwest::Url::parse(&initial_url).expect("request URL should parse"), + ); + + let response = pool + .send_with_resolver(request, |_| async { Ok(OutboundProxyRoute::Direct) }) + .await + .expect("no-redirect request should finish"); + + assert_eq!(response.status(), StatusCode::FOUND); + let requests = server.join().expect("redirect server should finish"); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("GET /start HTTP/1.1\r\n")); + } +} + +#[tokio::test] +async fn bounds_cached_routes_and_rebuilds_an_evicted_route() { + let pool = RouteAwareClientPool::with_builder( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Api, + HttpClientBuilder::new(), + ); + let routes = (0..=MAX_CACHED_ROUTES) + .map(|index| { + ( + format!("https://target-{index}.example"), + OutboundProxyRoute::Proxy { + url: format!("http://proxy-{index}.example"), + no_proxy: None, + }, + ) + }) + .collect::>(); + let resolver = FakeRouteResolver::new(routes.clone()); + + for request_url in routes.keys() { + resolve_with(&pool, &resolver, request_url) + .await + .expect("client should build"); + } + let evicted_route = { + let clients = pool.clients.lock().expect("client cache lock"); + assert_eq!(clients.len(), MAX_CACHED_ROUTES); + routes + .iter() + .find(|(_, route)| !clients.contains_key(*route)) + .map(|(request_url, _)| request_url.clone()) + .expect("one route should have been evicted") + }; + + resolve_with(&pool, &resolver, &evicted_route) + .await + .expect("evicted client should rebuild"); + + let clients = pool.clients.lock().expect("client cache lock"); + assert_eq!(clients.len(), MAX_CACHED_ROUTES); + assert!(clients.contains_key(&routes[&evicted_route])); +} + +#[tokio::test] +async fn request_timeout_covers_route_selection() { + let pool = manual_redirect_pool(); + let mut request = reqwest::Request::new( + Method::GET, + reqwest::Url::parse("http://route-selection-timeout.test/start") + .expect("request URL should parse"), + ); + *request.timeout_mut() = Some(Duration::from_millis(10)); + let resolver_calls = Arc::new(AtomicUsize::new(0)); + let observed_resolver_calls = Arc::clone(&resolver_calls); + + let error = pool + .send_with_resolver(request, move |_| { + observed_resolver_calls.fetch_add(1, Ordering::SeqCst); + async { + tokio::time::sleep(Duration::from_millis(100)).await; + Ok(OutboundProxyRoute::Direct) + } + }) + .await + .expect_err("request should time out during route selection"); + + assert!(matches!(error, RouteAwareRequestError::Timeout)); + assert_eq!(resolver_calls.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn request_timeout_is_shared_across_redirect_hops() { + let (address, server) = spawn_response_server(vec![ + "HTTP/1.1 302 Found\r\nLocation: /final\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + .to_string(), + ]); + let pool = manual_redirect_pool(); + let mut request = reqwest::Request::new( + Method::GET, + reqwest::Url::parse(&format!("http://{address}/start")).expect("request URL should parse"), + ); + *request.timeout_mut() = Some(Duration::from_secs(2)); + let resolver_calls = Arc::new(AtomicUsize::new(0)); + let observed_resolver_calls = Arc::clone(&resolver_calls); + + let error = pool + .send_with_resolver(request, move |_| { + let resolver_call = observed_resolver_calls.fetch_add(1, Ordering::SeqCst); + async move { + let delay = if resolver_call == 0 { + Duration::from_millis(500) + } else { + Duration::from_millis(1_750) + }; + tokio::time::sleep(delay).await; + Ok(OutboundProxyRoute::Direct) + } + }) + .await + .expect_err("redirect chain should exceed its shared timeout"); + + assert!(matches!(error, RouteAwareRequestError::Timeout)); + assert_eq!(resolver_calls.load(Ordering::SeqCst), 2); + assert_eq!( + server.join().expect("redirect server should finish").len(), + 1 + ); +} + +#[tokio::test] +async fn rejects_replayable_redirect_to_unsupported_scheme() { + let (address, server) = spawn_response_server(vec![ + "HTTP/1.1 307 Temporary Redirect\r\nLocation: ftp://example.com/final\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + .to_string(), + ]); + let pool = manual_redirect_pool(); + let request = reqwest::Request::new( + Method::GET, + reqwest::Url::parse(&format!("http://{address}/start")).expect("request URL should parse"), + ); + + let error = pool + .send_with_resolver(request, |_| async { Ok(OutboundProxyRoute::Direct) }) + .await + .expect_err("unsupported redirect scheme should fail"); + + assert!(matches!( + error, + RouteAwareRequestError::UnsupportedRedirectScheme(scheme) if scheme == "ftp" + )); + assert_eq!( + server.join().expect("redirect server should finish").len(), + 1 + ); +} + +#[tokio::test] +async fn rejects_redirects_beyond_the_limit() { + let responses = (0..=MAX_REDIRECTS) + .map(|redirect| { + format!( + "HTTP/1.1 302 Found\r\nLocation: /hop/{}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + redirect + 1 + ) + }) + .collect(); + let (address, server) = spawn_response_server(responses); + let pool = manual_redirect_pool(); + let request = reqwest::Request::new( + Method::GET, + reqwest::Url::parse(&format!("http://{address}/start")).expect("request URL should parse"), + ); + + let error = pool + .send_with_resolver(request, |_| async { Ok(OutboundProxyRoute::Direct) }) + .await + .expect_err("redirect chain should stop at the limit"); + let requests = server.join().expect("redirect server should finish"); + + assert!(matches!(error, RouteAwareRequestError::TooManyRedirects)); + assert_eq!(requests.len(), MAX_REDIRECTS + 1); + assert_eq!( + requests.last().and_then(|request| request.lines().next()), + Some("GET /hop/10 HTTP/1.1") + ); +} + +#[tokio::test] +async fn disabled_pool_logging_does_not_expose_request_or_response_data() { + let (address, server) = spawn_response_server(vec![ + "HTTP/1.1 200 OK\r\nx-sensitive-response: response-secret-value\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok" + .to_string(), + ]); + let pool = RouteAwareClientPool::with_builder( + HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy), + ClientRouteClass::Api, + HttpClientBuilder::new().without_request_logging(), + ); + let buffer = Arc::new(Mutex::new(Vec::new())); + let subscriber = tracing_subscriber::registry().with( + tracing_subscriber::fmt::layer() + .with_ansi(false) + .with_writer(TestLogWriter { + buffer: Arc::clone(&buffer), + }) + .with_filter( + tracing_subscriber::filter::Targets::new() + .with_target("codex_http_client", tracing::Level::TRACE), + ), + ); + let _guard = tracing::subscriber::set_default(subscriber); + tracing::debug!(target: "codex_http_client", "log capture sentinel"); + let request_url = format!( + "http://auth-user:password-secret-value@{address}/token?client_secret=query-secret-value" + ); + let mut request = reqwest::Request::new( + Method::POST, + reqwest::Url::parse(&request_url).expect("request URL should parse"), + ); + request.headers_mut().insert( + "x-sensitive-request", + HeaderValue::from_static("request-header-secret-value"), + ); + *request.body_mut() = Some("request-body-secret-value".into()); + *request.timeout_mut() = Some(Duration::from_secs(2)); + + let response = pool + .send_with_resolver(request, |_| async { Ok(OutboundProxyRoute::Direct) }) + .await + .expect("route-aware request should succeed"); + assert_eq!(response.status(), StatusCode::OK); + server.join().expect("server thread should finish"); + + let unresponsive_listener = + TcpListener::bind(("127.0.0.1", 0)).expect("unresponsive listener should bind"); + let unresponsive_address = unresponsive_listener + .local_addr() + .expect("unresponsive listener should have an address"); + let failure_url = format!( + "http://auth-user:failure-password-secret-value@{unresponsive_address}/token?client_secret=failure-query-secret-value" + ); + let mut request = reqwest::Request::new( + Method::POST, + reqwest::Url::parse(&failure_url).expect("failure URL should parse"), + ); + *request.timeout_mut() = Some(Duration::from_millis(100)); + + let error = pool + .send_with_resolver(request, |_| async { Ok(OutboundProxyRoute::Direct) }) + .await + .expect_err("request to an unresponsive listener should time out"); + assert!(error.is_timeout()); + + let logs = String::from_utf8(buffer.lock().expect("log buffer lock").clone()) + .expect("logs should be UTF-8"); + assert!(logs.contains("log capture sentinel")); + for secret in [ + "password-secret-value", + "query-secret-value", + "request-header-secret-value", + "request-body-secret-value", + "response-secret-value", + "failure-password-secret-value", + "failure-query-secret-value", + ] { + assert!(!logs.contains(secret), "logs exposed {secret}:\n{logs}"); + } +} + +#[derive(Clone)] +struct FakeRouteResolver { + routes: Arc>, + observed_urls: Arc>>, +} + +impl FakeRouteResolver { + fn new(routes: HashMap) -> Self { + Self { + routes: Arc::new(routes), + observed_urls: Arc::new(Mutex::new(Vec::new())), + } + } + + async fn resolve(&self, request_url: String) -> io::Result { + self.observed_urls + .lock() + .expect("observed URL lock") + .push(request_url.clone()); + self.routes + .get(&request_url) + .cloned() + .ok_or_else(|| io::Error::other(format!("no route for {request_url}"))) + } + + fn observed_urls(&self) -> Vec { + self.observed_urls + .lock() + .expect("observed URL lock") + .clone() + } +} + +async fn resolve_with( + pool: &RouteAwareClientPool, + resolver: &FakeRouteResolver, + request_url: &str, +) -> Result { + let resolver = resolver.clone(); + let (_, client, _) = pool + .client_for_url_with_resolver(request_url, move |request_url| async move { + resolver.resolve(request_url).await + }) + .await?; + Ok(client) +} + +fn manual_redirect_pool() -> RouteAwareClientPool { + RouteAwareClientPool::with_builder( + HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy), + ClientRouteClass::Api, + HttpClientBuilder::new(), + ) +} + +fn spawn_response_server( + responses: Vec, +) -> (std::net::SocketAddr, std::thread::JoinHandle>) { + let listener = TcpListener::bind(("127.0.0.1", 0)).expect("response listener should bind"); + let address = listener + .local_addr() + .expect("response listener should have an address"); + listener + .set_nonblocking(true) + .expect("response listener should become nonblocking"); + let server = std::thread::spawn(move || { + let mut requests = Vec::new(); + for response in responses { + let deadline = Instant::now() + Duration::from_secs(2); + let (mut stream, _) = loop { + match listener.accept() { + Ok(connection) => break connection, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => { + assert!( + Instant::now() < deadline, + "response server should receive the next request" + ); + std::thread::sleep(Duration::from_millis(10)); + } + Err(error) => panic!("response server should accept: {error}"), + } + }; + stream + .set_nonblocking(false) + .expect("response stream should become blocking"); + stream + .set_read_timeout(Some(Duration::from_secs(2))) + .expect("response stream should get a read timeout"); + requests.push(read_http_message(&mut stream)); + stream + .write_all(response.as_bytes()) + .expect("response server should write response"); + } + requests + }); + (address, server) +} + +fn read_http_message(stream: &mut impl Read) -> String { + let mut buffer = Vec::new(); + let mut chunk = [0_u8; 1024]; + loop { + let bytes_read = stream.read(&mut chunk).expect("HTTP message should read"); + if bytes_read == 0 { + break; + } + buffer.extend_from_slice(&chunk[..bytes_read]); + if let Some(header_end) = buffer.windows(4).position(|window| window == b"\r\n\r\n") { + let body_start = header_end + 4; + let headers = String::from_utf8_lossy(&buffer[..body_start]); + let content_length = headers + .lines() + .filter_map(|line| line.split_once(':')) + .find_map(|(name, value)| { + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + if buffer.len() >= body_start + content_length { + break; + } + } + } + String::from_utf8_lossy(&buffer).into_owned() +} + +#[derive(Clone)] +struct TestLogWriter { + buffer: Arc>>, +} + +struct TestLogSink { + buffer: Arc>>, +} + +impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for TestLogWriter { + type Writer = TestLogSink; + + fn make_writer(&'a self) -> Self::Writer { + TestLogSink { + buffer: Arc::clone(&self.buffer), + } + } +} + +impl Write for TestLogSink { + fn write(&mut self, buffer: &[u8]) -> io::Result { + let mut log_buffer = self + .buffer + .lock() + .map_err(|_| io::Error::other("log buffer lock was poisoned"))?; + log_buffer.extend(buffer); + Ok(buffer.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} diff --git a/codex-rs/http-client/src/route_aware_redirect.rs b/codex-rs/http-client/src/route_aware_redirect.rs new file mode 100644 index 0000000000000000000000000000000000000000..77af745725e66dee0112e240c95030065318978d --- /dev/null +++ b/codex-rs/http-client/src/route_aware_redirect.rs @@ -0,0 +1,142 @@ +//! Redirect semantics for requests whose proxy route is selected from the complete URL. +//! +//! Reqwest normally follows redirects inside one `Client::execute` call. That is safe for +//! `ReqwestDefault`, where reqwest owns proxy selection, but not for `RespectSystemProxy`: Codex +//! resolves PAC and operating-system proxy settings before building a client pinned to the +//! resulting direct or concrete-proxy route. An internal reqwest redirect would reuse that client +//! for the new URL without giving [`crate::RouteAwareClientPool`] an opportunity to resolve the +//! redirect destination. +//! +//! The pool therefore disables reqwest redirects for `RespectSystemProxy`, observes each redirect +//! response, and uses these helpers to construct the next request before resolving its route. The +//! helpers preserve the relevant reqwest method/body behavior while applying stricter +//! origin-sensitive credential and `Referer` rules at the boundary between requests. + +use std::time::Duration; + +use http::HeaderMap; +use http::Method; +use http::StatusCode; +use http::header::AUTHORIZATION; +use http::header::CONTENT_ENCODING; +use http::header::CONTENT_LENGTH; +use http::header::CONTENT_TYPE; +use http::header::COOKIE; +use http::header::LOCATION; +use http::header::PROXY_AUTHORIZATION; +use http::header::REFERER; +use http::header::TRANSFER_ENCODING; +use http::header::WWW_AUTHENTICATE; + +pub(super) const MAX_REDIRECTS: usize = 10; + +pub(super) fn is_redirect(status: StatusCode) -> bool { + matches!( + status, + StatusCode::MOVED_PERMANENTLY + | StatusCode::FOUND + | StatusCode::SEE_OTHER + | StatusCode::TEMPORARY_REDIRECT + | StatusCode::PERMANENT_REDIRECT + ) +} + +pub(super) fn redirect_url(response: &reqwest::Response) -> Option { + let location = response.headers().get(LOCATION)?.to_str().ok()?; + response.url().join(location).ok() +} + +pub(super) fn redirect_request( + status: StatusCode, + mut method: Method, + mut headers: HeaderMap, + version: http::Version, + timeout: Option, + replay: Option, + next_url: reqwest::Url, +) -> Option { + let drop_body = match status { + StatusCode::MOVED_PERMANENTLY | StatusCode::FOUND if method == Method::POST => { + method = Method::GET; + true + } + StatusCode::SEE_OTHER => { + if method != Method::HEAD { + method = Method::GET; + } + true + } + StatusCode::MOVED_PERMANENTLY + | StatusCode::FOUND + | StatusCode::TEMPORARY_REDIRECT + | StatusCode::PERMANENT_REDIRECT => false, + _ => return None, + }; + + if drop_body { + for header in [ + CONTENT_TYPE, + CONTENT_LENGTH, + CONTENT_ENCODING, + TRANSFER_ENCODING, + ] { + headers.remove(header); + } + let mut request = reqwest::Request::new(method, next_url); + *request.headers_mut() = headers; + *request.version_mut() = version; + *request.timeout_mut() = timeout; + Some(request) + } else { + replay.map(|mut request| { + *request.url_mut() = next_url; + request + }) + } +} + +pub(super) fn remove_sensitive_headers( + headers: &mut HeaderMap, + previous: &reqwest::Url, + next: &reqwest::Url, +) { + if !same_origin(previous, next) { + for header in [AUTHORIZATION, COOKIE, PROXY_AUTHORIZATION, WWW_AUTHENTICATE] { + headers.remove(header); + } + headers.remove("cookie2"); + } +} + +pub(super) fn insert_referer( + headers: &mut HeaderMap, + previous: &reqwest::Url, + next: &reqwest::Url, +) { + headers.remove(REFERER); + if next.scheme() == "http" && previous.scheme() == "https" { + return; + } + + let mut referer = previous.clone(); + let _ = referer.set_username(""); + let _ = referer.set_password(None); + referer.set_fragment(None); + if !same_origin(previous, next) { + referer.set_path("/"); + referer.set_query(None); + } + if let Ok(value) = referer.as_str().parse() { + headers.insert(REFERER, value); + } +} + +fn same_origin(previous: &reqwest::Url, next: &reqwest::Url) -> bool { + previous.scheme() == next.scheme() + && previous.host_str() == next.host_str() + && previous.port_or_known_default() == next.port_or_known_default() +} + +#[cfg(test)] +#[path = "route_aware_redirect_tests.rs"] +mod tests; diff --git a/codex-rs/http-client/src/route_aware_redirect_integration_tests.rs b/codex-rs/http-client/src/route_aware_redirect_integration_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..344e8a7f6b010c3e47d2b21d20c77dafb5adc932 --- /dev/null +++ b/codex-rs/http-client/src/route_aware_redirect_integration_tests.rs @@ -0,0 +1,172 @@ +use std::io; +use std::io::Read; +use std::io::Write; +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; +use std::time::Instant; + +use http::header::AUTHORIZATION; +use http::header::COOKIE; +use http::header::PROXY_AUTHORIZATION; +use pretty_assertions::assert_eq; +use tracing_subscriber::Layer; +use tracing_subscriber::layer::SubscriberExt; + +use super::*; +use crate::RouteAwareClientPool; + +#[tokio::test] +async fn route_aware_pool_re_resolves_redirects_and_logs_only_final_outcome() { + let (proxy_addr, proxy_thread) = + spawn_response("HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok"); + let (redirect_addr, redirect_thread) = spawn_response( + "HTTP/1.1 302 Found\r\nLocation: /final?token=redirect-target-secret\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ); + let initial_url = format!("http://{redirect_addr}/start"); + let redirected_url = format!("http://{redirect_addr}/final?token=redirect-target-secret"); + cache_system_proxy_decision(&initial_url, SystemProxyDecision::Direct); + cache_system_proxy_decision( + &redirected_url, + SystemProxyDecision::Proxy { + url: format!("http://{proxy_addr}"), + }, + ); + let pool = RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy), + ClientRouteClass::Api, + ); + let log_buffer = Arc::new(Mutex::new(Vec::new())); + let subscriber = tracing_subscriber::registry().with( + tracing_subscriber::fmt::layer() + .with_ansi(false) + .with_writer(TestLogWriter { + buffer: Arc::clone(&log_buffer), + }) + .with_filter( + tracing_subscriber::filter::Targets::new() + .with_target("codex_http_client", tracing::Level::TRACE), + ), + ); + let _guard = tracing::subscriber::set_default(subscriber); + + tokio::time::timeout( + Duration::from_secs(2), + pool.get(&initial_url) + .header(AUTHORIZATION, "Bearer origin-secret") + .header(COOKIE, "session=origin-secret") + .header(PROXY_AUTHORIZATION, "Basic proxy-secret") + .send(), + ) + .await + .expect("redirected request should finish") + .expect("redirected request should use both selected routes"); + redirect_thread + .join() + .expect("redirect listener should finish"); + let redirected_request = proxy_thread.join().expect("proxy listener should finish"); + assert_eq!( + redirected_request.lines().next(), + Some(format!("GET {redirected_url} HTTP/1.1").as_str()) + ); + assert!(has_header(&redirected_request, "authorization")); + assert!(has_header(&redirected_request, "cookie")); + assert!(!has_header(&redirected_request, "proxy-authorization")); + + let logs = String::from_utf8(log_buffer.lock().expect("log buffer lock").clone()) + .expect("logs should be UTF-8"); + assert!(logs.contains(&initial_url)); + assert_eq!(logs.matches("Request completed").count(), 1); + assert!(!logs.contains("redirect-target-secret")); +} + +fn spawn_response( + response: &'static str, +) -> (std::net::SocketAddr, std::thread::JoinHandle) { + let listener = std::net::TcpListener::bind(("127.0.0.1", 0)).expect("bind listener"); + let address = listener.local_addr().expect("HTTP listener address"); + listener.set_nonblocking(true).expect("set nonblocking"); + let thread = std::thread::spawn(move || { + let deadline = Instant::now() + Duration::from_secs(2); + let (mut stream, _) = loop { + match listener.accept() { + Ok(connection) => break connection, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => { + assert!(Instant::now() < deadline, "request timed out"); + std::thread::sleep(Duration::from_millis(10)); + } + Err(error) => panic!("HTTP listener should accept: {error}"), + } + }; + // Accepted sockets inherit the listener's nonblocking mode on macOS. + stream + .set_nonblocking(false) + .expect("set accepted stream blocking"); + stream + .set_read_timeout(Some(Duration::from_secs(2))) + .expect("read timeout"); + let request = read_http_headers(&mut stream); + stream + .write_all(response.as_bytes()) + .expect("write response"); + request + }); + (address, thread) +} + +fn read_http_headers(stream: &mut impl Read) -> String { + let mut buffer = Vec::new(); + let mut chunk = [0_u8; 1024]; + loop { + let bytes_read = stream.read(&mut chunk).expect("HTTP headers should read"); + if bytes_read == 0 { + break; + } + buffer.extend_from_slice(&chunk[..bytes_read]); + if buffer.windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + } + String::from_utf8_lossy(&buffer).into_owned() +} + +fn has_header(request: &str, expected_name: &str) -> bool { + request.lines().any(|line| { + line.split_once(':') + .is_some_and(|(name, _)| name.eq_ignore_ascii_case(expected_name)) + }) +} + +#[derive(Clone)] +struct TestLogWriter { + buffer: Arc>>, +} + +struct TestLogSink { + buffer: Arc>>, +} + +impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for TestLogWriter { + type Writer = TestLogSink; + + fn make_writer(&'a self) -> Self::Writer { + TestLogSink { + buffer: Arc::clone(&self.buffer), + } + } +} + +impl Write for TestLogSink { + fn write(&mut self, buffer: &[u8]) -> io::Result { + let mut log_buffer = self + .buffer + .lock() + .map_err(|_| io::Error::other("log buffer lock was poisoned"))?; + log_buffer.extend(buffer); + Ok(buffer.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} diff --git a/codex-rs/http-client/src/route_aware_redirect_tests.rs b/codex-rs/http-client/src/route_aware_redirect_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..d9e8f6053f3e61a1bba7db149cd295bb42de1df5 --- /dev/null +++ b/codex-rs/http-client/src/route_aware_redirect_tests.rs @@ -0,0 +1,233 @@ +use http::HeaderValue; +use http::header::CONTENT_ENCODING; +use http::header::CONTENT_LENGTH; +use http::header::CONTENT_TYPE; +use http::header::COOKIE; +use http::header::REFERER; +use http::header::TRANSFER_ENCODING; +use pretty_assertions::assert_eq; + +use super::*; + +#[derive(Clone, Copy, Debug)] +enum RedirectBodyBehavior { + Drop, + Preserve, +} + +#[test] +fn redirects_match_reqwest_method_and_body_rules() { + let original_url = + reqwest::Url::parse("https://example.com/original").expect("original URL should parse"); + let redirected_url = + reqwest::Url::parse("https://example.com/next").expect("redirect URL should parse"); + + for (status, method, expected_method, body_behavior) in [ + ( + StatusCode::MOVED_PERMANENTLY, + Method::POST, + Method::GET, + RedirectBodyBehavior::Drop, + ), + ( + StatusCode::FOUND, + Method::POST, + Method::GET, + RedirectBodyBehavior::Drop, + ), + ( + StatusCode::SEE_OTHER, + Method::POST, + Method::GET, + RedirectBodyBehavior::Drop, + ), + ( + StatusCode::SEE_OTHER, + Method::HEAD, + Method::HEAD, + RedirectBodyBehavior::Drop, + ), + ( + StatusCode::MOVED_PERMANENTLY, + Method::PUT, + Method::PUT, + RedirectBodyBehavior::Preserve, + ), + ( + StatusCode::FOUND, + Method::PUT, + Method::PUT, + RedirectBodyBehavior::Preserve, + ), + ( + StatusCode::TEMPORARY_REDIRECT, + Method::POST, + Method::POST, + RedirectBodyBehavior::Preserve, + ), + ( + StatusCode::PERMANENT_REDIRECT, + Method::POST, + Method::POST, + RedirectBodyBehavior::Preserve, + ), + ] { + let mut original = reqwest::Request::new(method.clone(), original_url.clone()); + *original.headers_mut() = HeaderMap::from_iter([ + (CONTENT_TYPE, HeaderValue::from_static("application/json")), + (CONTENT_LENGTH, HeaderValue::from_static("2")), + (CONTENT_ENCODING, HeaderValue::from_static("identity")), + (TRANSFER_ENCODING, HeaderValue::from_static("chunked")), + ( + http::header::ACCEPT, + HeaderValue::from_static("application/json"), + ), + ]); + *original.version_mut() = http::Version::HTTP_2; + *original.timeout_mut() = Some(Duration::from_secs(7)); + *original.body_mut() = Some("{}".into()); + let mut expected_headers = original.headers().clone(); + let expected_body = match body_behavior { + RedirectBodyBehavior::Drop => { + for header in [ + CONTENT_TYPE, + CONTENT_LENGTH, + CONTENT_ENCODING, + TRANSFER_ENCODING, + ] { + expected_headers.remove(header); + } + None + } + RedirectBodyBehavior::Preserve => Some(b"{}".to_vec()), + }; + + let redirected = redirect_request( + status, + original.method().clone(), + original.headers().clone(), + original.version(), + original.timeout().copied(), + original.try_clone(), + redirected_url.clone(), + ) + .expect("supported redirect should be followed"); + + assert_eq!( + ( + redirected.method().clone(), + redirected.url().clone(), + redirected.headers().clone(), + redirected.version(), + redirected.timeout().copied(), + redirected + .body() + .and_then(reqwest::Body::as_bytes) + .map(<[u8]>::to_vec), + ), + ( + expected_method, + redirected_url.clone(), + expected_headers, + http::Version::HTTP_2, + Some(Duration::from_secs(7)), + expected_body, + ), + "redirect behavior for {status} {method}", + ); + } +} + +#[test] +fn redirect_referer_follows_strict_origin_when_cross_origin() { + let previous = + reqwest::Url::parse("https://user:password@example.com/start?access_token=secret#fragment") + .expect("valid URL"); + + let mut same_origin_headers = HeaderMap::new(); + let same_origin = reqwest::Url::parse("https://example.com/next").expect("valid URL"); + insert_referer(&mut same_origin_headers, &previous, &same_origin); + assert_eq!( + same_origin_headers.get(REFERER), + Some(&HeaderValue::from_static( + "https://example.com/start?access_token=secret" + )) + ); + + let mut cross_origin_headers = HeaderMap::new(); + let cross_origin = reqwest::Url::parse("https://other.example/next").expect("valid URL"); + insert_referer(&mut cross_origin_headers, &previous, &cross_origin); + assert_eq!( + cross_origin_headers.get(REFERER), + Some(&HeaderValue::from_static("https://example.com/")) + ); + + let mut downgrade_headers = HeaderMap::from_iter([( + REFERER, + HeaderValue::from_static("https://stale.example/path?secret=value"), + )]); + let downgrade = reqwest::Url::parse("http://other.example/next").expect("valid URL"); + insert_referer(&mut downgrade_headers, &previous, &downgrade); + assert_eq!(downgrade_headers.get(REFERER), None); +} + +#[test] +fn non_replayable_redirect_is_not_followed_regardless_of_target_scheme() { + let next_url = reqwest::Url::parse("ftp://example.com/next").expect("valid URL"); + + let redirected = redirect_request( + StatusCode::TEMPORARY_REDIRECT, + Method::POST, + HeaderMap::new(), + http::Version::HTTP_11, + /*timeout*/ None, + /*replay*/ None, + next_url, + ); + + assert!(redirected.is_none()); +} + +#[test] +fn redirect_credentials_are_retained_only_for_the_same_origin() { + for (previous, next, retain_credentials) in [ + ( + "https://example.com:8080/start", + "https://example.com:8080/next", + true, + ), + ( + "https://example.com:8080/start", + "http://example.com:8080/next", + false, + ), + ( + "https://example.com:8080/start", + "https://other.example:8080/next", + false, + ), + ( + "https://example.com:8080/start", + "https://example.com:8081/next", + false, + ), + ] { + let previous = reqwest::Url::parse(previous).expect("previous URL should parse"); + let next = reqwest::Url::parse(next).expect("next URL should parse"); + let mut headers = HeaderMap::from_iter([ + (AUTHORIZATION, HeaderValue::from_static("Bearer secret")), + (COOKIE, HeaderValue::from_static("session=secret")), + ]); + + remove_sensitive_headers(&mut headers, &previous, &next); + + assert_eq!( + ( + headers.contains_key(AUTHORIZATION), + headers.contains_key(COOKIE), + ), + (retain_credentials, retain_credentials), + "credential handling for {previous} -> {next}" + ); + } +} diff --git a/codex-rs/http-client/src/route_aware_tls_fallback_tests.rs b/codex-rs/http-client/src/route_aware_tls_fallback_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..e5968030907599d7cdcdc608824aa27c4815c54a --- /dev/null +++ b/codex-rs/http-client/src/route_aware_tls_fallback_tests.rs @@ -0,0 +1,510 @@ +use std::io; +use std::io::Read; +use std::io::Write; +use std::net::TcpListener; +use std::sync::Arc; +use std::sync::mpsc; +use std::thread; +use std::time::Duration; +use std::time::Instant; + +use bytes::Bytes; +use futures::stream; +use pretty_assertions::assert_eq; +use rcgen::CertifiedKey; +use rcgen::generate_simple_self_signed; +use rustls_pki_types::PrivateKeyDer; + +use super::ClientRouteClass; +use super::HttpClient; +use super::HttpClientFactory; +use super::Method; +use super::OutboundProxyPolicy; +use super::OutboundProxyRoute; +use super::RouteAwareClientPool; +use super::RouteAwareRequestError; +use super::SelectedTlsBackend; +use crate::tls_backend_fallback::should_retry_with_rustls; + +const PROTOCOL_VERSION_TLS_ALERT: &[u8] = &[21, 3, 3, 0, 2, 2, 70]; + +type SuccessfulTlsFallbackServer = (String, HttpClient, mpsc::Receiver>>); + +#[tokio::test] +async fn default_pool_does_not_retry_a_native_tls_protocol_failure() { + let (url, attempts, stop_server) = + spawn_protocol_version_rejection_server(/*maximum_attempts*/ 2) + .expect("TLS rejection server should start"); + let pool = RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Other, + ); + let request = reqwest::Request::new( + Method::POST, + reqwest::Url::parse(&url).expect("valid HTTPS URL"), + ); + + let error = pool + .send_with_resolver(request, |_| async { Ok(OutboundProxyRoute::Direct) }) + .await + .expect_err("ordinary traffic should preserve the native TLS failure"); + + let _ = stop_server.send(()); + assert!(error.is_connect()); + assert_eq!( + attempts + .recv() + .expect("TLS server should finish") + .expect("TLS server should reject one connection"), + 1 + ); +} + +#[tokio::test] +async fn retries_a_native_tls_protocol_failure_once_with_rustls() { + let (url, attempts, stop_server) = + spawn_protocol_version_rejection_server(/*maximum_attempts*/ 2) + .expect("TLS rejection server should start"); + let pool = RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Other, + ) + .with_tls_backend_fallback(); + let destination = reqwest::Url::parse(&url).expect("valid HTTPS URL"); + let mut request = reqwest::Request::new(Method::POST, destination.clone()); + *request.body_mut() = Some(Bytes::from_static(b"mcp-initialize").into()); + + let error = pool + .send_with_resolver(request, |_| async { Ok(OutboundProxyRoute::Direct) }) + .await + .expect_err("both TLS handshakes should be rejected"); + + let _ = stop_server.send(()); + assert!(error.is_connect()); + assert_eq!( + attempts + .recv() + .expect("TLS server should finish") + .expect("TLS server should reject both connections"), + 2, + "native TLS protocol failure should retry with rustls: {error:?}" + ); + let rustls_clients = pool.rustls_clients.as_ref().expect("TLS fallback cache"); + assert_eq!( + ( + rustls_clients.requires_rustls(&destination, &OutboundProxyRoute::Direct), + rustls_clients + .client_for_route(&OutboundProxyRoute::Direct) + .is_some(), + ), + (false, false) + ); +} + +#[tokio::test] +async fn retries_a_native_tls_failure_after_another_request_caches_rustls() { + let (url, attempts, stop_server) = + spawn_protocol_version_rejection_server(/*maximum_attempts*/ 2) + .expect("TLS rejection server should start"); + let pool = RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Other, + ) + .with_tls_backend_fallback(); + let destination = reqwest::Url::parse(&url).expect("valid HTTPS URL"); + let (route, native_client, selected_tls_backend) = pool + .client_for_url_with_resolver(destination.as_str(), |_| async { + Ok(OutboundProxyRoute::Direct) + }) + .await + .expect("native TLS client should resolve"); + assert_eq!(selected_tls_backend, SelectedTlsBackend::TransportDefault); + + let request = reqwest::Request::new(Method::POST, destination.clone()); + let replay = request.try_clone().expect("request should be replayable"); + let error = native_client + .execute_without_request_logging(request) + .await + .expect_err("native TLS handshake should fail"); + let rustls_client = pool + .rustls_client_for_route(&route) + .expect("rustls fallback client should build"); + pool.rustls_clients + .as_ref() + .expect("TLS fallback cache") + .remember(&destination, &route, rustls_client); + + let error = pool + .retry_with_rustls( + &destination, + &route, + selected_tls_backend, + Some(&replay), + error, + /*timeout_deadline*/ None, + ) + .await + .expect_err("both TLS handshakes should be rejected"); + + let _ = stop_server.send(()); + assert!(error.is_connect()); + assert_eq!( + attempts + .recv() + .expect("TLS server should finish") + .expect("TLS server should reject both connections"), + 2, + "an in-flight native TLS failure should retry despite a concurrent cache update" + ); +} + +#[tokio::test] +async fn does_not_retry_a_cached_rustls_tls_protocol_failure() { + let (url, attempts, stop_server) = + spawn_protocol_version_rejection_server(/*maximum_attempts*/ 2) + .expect("TLS rejection server should start"); + let pool = RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Other, + ) + .with_tls_backend_fallback(); + let destination = reqwest::Url::parse(&url).expect("valid HTTPS URL"); + let rustls_client = pool + .rustls_client_for_route(&OutboundProxyRoute::Direct) + .expect("rustls fallback client should build"); + pool.rustls_clients + .as_ref() + .expect("TLS fallback cache") + .remember(&destination, &OutboundProxyRoute::Direct, rustls_client); + let request = reqwest::Request::new(Method::POST, destination); + + let error = pool + .send_with_resolver(request, |_| async { Ok(OutboundProxyRoute::Direct) }) + .await + .expect_err("cached rustls TLS handshake should fail"); + + let _ = stop_server.send(()); + assert!(error.is_connect()); + assert_eq!( + attempts + .recv() + .expect("TLS server should finish") + .expect("TLS server should reject one connection"), + 1, + "a cached rustls client must not retry against itself" + ); +} + +#[tokio::test] +async fn successful_rustls_fallback_replays_the_request_and_reuses_the_destination() { + let (url, trusted_rustls_client, observed_requests) = + spawn_successful_tls_fallback_server().expect("TLS fallback server should start"); + let pool = RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Other, + ) + .with_tls_backend_fallback(); + let destination = reqwest::Url::parse(&url).expect("valid HTTPS URL"); + let existing_destination = + reqwest::Url::parse("https://another-mcp.example.com/mcp").expect("valid HTTPS URL"); + let rustls_clients = pool.rustls_clients.as_ref().expect("TLS fallback cache"); + + // Another destination may already have established a trusted rustls client on this route. + // The test destination must still attempt native TLS before reusing that route-level client. + rustls_clients.remember( + &existing_destination, + &OutboundProxyRoute::Direct, + trusted_rustls_client, + ); + assert!(!rustls_clients.requires_rustls(&destination, &OutboundProxyRoute::Direct)); + + let mut initialize_request = reqwest::Request::new(Method::POST, destination.clone()); + *initialize_request.body_mut() = Some(Bytes::from_static(b"mcp-initialize").into()); + let initialize_response = pool + .send_with_resolver(initialize_request, |_| async { + Ok(OutboundProxyRoute::Direct) + }) + .await + .expect("native TLS protocol failure should recover with rustls"); + + assert_eq!( + initialize_response + .text() + .await + .expect("fallback response body should be readable"), + "ok" + ); + assert!(rustls_clients.requires_rustls(&destination, &OutboundProxyRoute::Direct)); + + let mut tool_request = reqwest::Request::new(Method::POST, destination); + *tool_request.body_mut() = Some(Bytes::from_static(b"mcp-tool-call").into()); + let tool_response = pool + .send_with_resolver(tool_request, |_| async { Ok(OutboundProxyRoute::Direct) }) + .await + .expect("remembered destination should reuse rustls without another native TLS attempt"); + + assert_eq!( + tool_response + .text() + .await + .expect("cached rustls response body should be readable"), + "ok" + ); + assert_eq!( + observed_requests + .recv_timeout(Duration::from_secs(5)) + .expect("TLS fallback server should finish") + .expect("TLS fallback server should capture both requests"), + vec!["mcp-initialize".to_string(), "mcp-tool-call".to_string()] + ); +} + +#[tokio::test] +async fn retries_a_tls_protocol_failure_when_request_url_contains_certificate_markers() { + let (url, attempts, stop_server) = + spawn_protocol_version_rejection_server(/*maximum_attempts*/ 2) + .expect("TLS rejection server should start"); + let pool = RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Other, + ) + .with_tls_backend_fallback(); + let mut destination = reqwest::Url::parse(&url).expect("valid HTTPS URL"); + destination.set_path("/certificate/hostname/expired/revoked/mcp"); + let request = reqwest::Request::new(Method::POST, destination); + + let error = pool + .send_with_resolver(request, |_| async { Ok(OutboundProxyRoute::Direct) }) + .await + .expect_err("both TLS handshakes should be rejected"); + + let _ = stop_server.send(()); + assert!(error.is_connect()); + assert_eq!( + attempts + .recv() + .expect("TLS server should finish") + .expect("TLS server should reject both connections"), + 2, + "request URL text should not prevent TLS backend fallback: {error:?}" + ); +} + +#[tokio::test] +async fn does_not_retry_a_non_replayable_streaming_request() { + let (url, attempts, stop_server) = + spawn_protocol_version_rejection_server(/*maximum_attempts*/ 2) + .expect("TLS rejection server should start"); + let pool = RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Other, + ) + .with_tls_backend_fallback(); + let mut request = reqwest::Request::new( + Method::POST, + reqwest::Url::parse(&url).expect("valid HTTPS URL"), + ); + *request.body_mut() = Some(reqwest::Body::wrap_stream(stream::iter(vec![Ok::< + _, + io::Error, + >( + Bytes::from_static(b"streaming body"), + )]))); + + let error = pool + .send_with_resolver(request, |_| async { Ok(OutboundProxyRoute::Direct) }) + .await + .expect_err("non-replayable request should preserve the native TLS failure"); + + let _ = stop_server.send(()); + assert!(error.is_connect()); + let RouteAwareRequestError::Request(ref request_error) = error else { + panic!("expected native TLS request error, got {error:?}"); + }; + assert!( + should_retry_with_rustls(request_error), + "native TLS protocol failure should be retryable: {request_error:?}" + ); + assert_eq!( + attempts + .recv() + .expect("TLS server should finish") + .expect("TLS server should reject one connection"), + 1 + ); +} + +fn spawn_successful_tls_fallback_server() -> io::Result { + codex_utils_rustls_provider::ensure_rustls_crypto_provider(); + let CertifiedKey { cert, signing_key } = + generate_simple_self_signed(vec!["127.0.0.1".to_string()]).map_err(io::Error::other)?; + let certificate = cert.der().clone(); + let private_key = PrivateKeyDer::from(signing_key); + let tls_config = Arc::new( + rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(vec![certificate.clone()], private_key) + .map_err(io::Error::other)?, + ); + let trusted_rustls_client = HttpClient::new( + reqwest::Client::builder() + .use_rustls_tls() + .add_root_certificate( + reqwest::Certificate::from_der(certificate.as_ref()).map_err(io::Error::other)?, + ) + .no_proxy() + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(io::Error::other)?, + ); + + let listener = TcpListener::bind(("127.0.0.1", 0))?; + let address = listener.local_addr()?; + listener.set_nonblocking(true)?; + let (requests_tx, requests_rx) = mpsc::channel(); + + thread::spawn(move || { + let result = (|| -> io::Result> { + let mut observed_requests = Vec::new(); + for connection_index in 0..3 { + let deadline = Instant::now() + Duration::from_secs(5); + let mut stream = loop { + match listener.accept() { + Ok((stream, _)) => break stream, + Err(error) if error.kind() == io::ErrorKind::WouldBlock => { + if Instant::now() >= deadline { + return Err(io::Error::new( + io::ErrorKind::TimedOut, + "timed out waiting for the TLS fallback client", + )); + } + thread::sleep(Duration::from_millis(10)); + } + Err(error) => return Err(error), + } + }; + stream.set_nonblocking(false)?; + stream.set_read_timeout(Some(Duration::from_secs(5)))?; + stream.set_write_timeout(Some(Duration::from_secs(5)))?; + + if connection_index == 0 { + let mut client_hello = [0_u8; 2_048]; + if stream.read(&mut client_hello)? == 0 { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "native TLS peer closed before sending a ClientHello", + )); + } + stream.write_all(PROTOCOL_VERSION_TLS_ALERT)?; + stream.flush()?; + continue; + } + + let connection = rustls::ServerConnection::new(Arc::clone(&tls_config)) + .map_err(io::Error::other)?; + let mut tls = rustls::StreamOwned::new(connection, stream); + let mut request = Vec::new(); + let mut chunk = [0_u8; 1_024]; + let header_end = loop { + let bytes_read = tls.read(&mut chunk)?; + if bytes_read == 0 { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "rustls peer closed before sending HTTP headers", + )); + } + request.extend_from_slice(&chunk[..bytes_read]); + if let Some(header_end) = + request.windows(4).position(|window| window == b"\r\n\r\n") + { + break header_end + 4; + } + }; + let headers = + std::str::from_utf8(&request[..header_end]).map_err(io::Error::other)?; + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then_some(value.trim()) + }) + .ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "missing Content-Length") + })? + .parse::() + .map_err(io::Error::other)?; + while request.len() < header_end + content_length { + let bytes_read = tls.read(&mut chunk)?; + if bytes_read == 0 { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "rustls peer closed before sending the HTTP body", + )); + } + request.extend_from_slice(&chunk[..bytes_read]); + } + let body = std::str::from_utf8(&request[header_end..header_end + content_length]) + .map_err(io::Error::other)?; + observed_requests.push(body.to_string()); + tls.write_all( + b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok", + )?; + tls.flush()?; + } + Ok(observed_requests) + })(); + let _ = requests_tx.send(result); + }); + + Ok(( + format!("https://{address}/mcp"), + trusted_rustls_client, + requests_rx, + )) +} + +fn spawn_protocol_version_rejection_server( + maximum_attempts: usize, +) -> io::Result<(String, mpsc::Receiver>, mpsc::Sender<()>)> { + let listener = TcpListener::bind(("127.0.0.1", 0))?; + let address = listener.local_addr()?; + listener.set_nonblocking(true)?; + let (attempts_tx, attempts_rx) = mpsc::channel(); + let (stop_tx, stop_rx) = mpsc::channel(); + + thread::spawn(move || { + let result = (|| -> io::Result { + let mut attempts = 0; + let deadline = Instant::now() + Duration::from_secs(5); + while attempts < maximum_attempts && Instant::now() < deadline { + if stop_rx.try_recv().is_ok() { + break; + } + match listener.accept() { + Ok((mut stream, _)) => { + stream.set_nonblocking(false)?; + stream.set_read_timeout(Some(Duration::from_secs(2)))?; + let mut client_hello = [0_u8; 2_048]; + if stream.read(&mut client_hello)? == 0 { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "TLS peer closed before sending a ClientHello", + )); + } + stream.write_all(PROTOCOL_VERSION_TLS_ALERT)?; + attempts += 1; + } + Err(error) if error.kind() == io::ErrorKind::WouldBlock => { + thread::sleep(Duration::from_millis(10)); + } + Err(error) => return Err(error), + } + } + Ok(attempts) + })(); + let _ = attempts_tx.send(result); + }); + + Ok((format!("https://{address}/mcp"), attempts_rx, stop_tx)) +} diff --git a/codex-rs/http-client/src/tls_backend_fallback.rs b/codex-rs/http-client/src/tls_backend_fallback.rs new file mode 100644 index 0000000000000000000000000000000000000000..760df1824ad752fe8d4b5f3f2b0d0e3e848c0144 --- /dev/null +++ b/codex-rs/http-client/src/tls_backend_fallback.rs @@ -0,0 +1,156 @@ +//! Narrow TLS backend fallback for delegated requests that select their route per destination. +//! +//! Native TLS remains the default. A recognized connection-time protocol negotiation failure can +//! select rustls for one HTTPS origin and outbound route without changing other destinations. + +use std::collections::HashMap; +use std::collections::HashSet; +use std::error::Error; +use std::sync::Arc; +use std::sync::Mutex; + +use crate::HttpClient; +use crate::OutboundProxyRoute; + +const MAX_CACHED_RUSTLS_DESTINATIONS: usize = 16; +// Schannel maps TLS alert 70 (protocol_version) to SEC_E_UNSUPPORTED_FUNCTION. +const SCHANNEL_PROTOCOL_VERSION_ERROR: i32 = 0x8009_0302_u32 as i32; + +#[derive(Clone, Default)] +pub(crate) struct RustlsClientCache { + state: Arc>, +} + +#[derive(Default)] +struct RustlsClientCacheState { + destinations: HashSet, + clients: HashMap, +} + +#[derive(Clone, Hash, PartialEq, Eq)] +struct DestinationRoute { + host: String, + port: u16, + route: OutboundProxyRoute, +} + +impl RustlsClientCache { + pub(crate) fn requires_rustls(&self, url: &reqwest::Url, route: &OutboundProxyRoute) -> bool { + let Some(destination) = DestinationRoute::new(url, route) else { + return false; + }; + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .destinations + .contains(&destination) + } + + pub(crate) fn client_for_route(&self, route: &OutboundProxyRoute) -> Option { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clients + .get(route) + .cloned() + } + + pub(crate) fn remember( + &self, + url: &reqwest::Url, + route: &OutboundProxyRoute, + client: HttpClient, + ) { + let Some(destination) = DestinationRoute::new(url, route) else { + return; + }; + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.destinations.contains(&destination) { + return; + } + if state.destinations.len() >= MAX_CACHED_RUSTLS_DESTINATIONS + && let Some(destination_to_evict) = state.destinations.iter().next().cloned() + { + state.destinations.remove(&destination_to_evict); + if !state + .destinations + .iter() + .any(|destination| destination.route == destination_to_evict.route) + { + state.clients.remove(&destination_to_evict.route); + } + } + state.clients.entry(route.clone()).or_insert(client); + state.destinations.insert(destination); + } +} + +impl DestinationRoute { + fn new(url: &reqwest::Url, route: &OutboundProxyRoute) -> Option { + if url.scheme() != "https" { + return None; + } + Some(Self { + host: url.host_str()?.to_ascii_lowercase(), + port: url.port_or_known_default()?, + route: route.clone(), + }) + } +} + +pub(crate) fn should_retry_with_rustls(error: &reqwest::Error) -> bool { + error.is_connect() && !error.is_timeout() && error.source().is_some_and(has_retryable_tls_error) +} + +fn has_retryable_tls_error(error: &(dyn Error + 'static)) -> bool { + let mut source = Some(error); + let mut recognized_negotiation_failure = false; + + while let Some(error) = source { + let message = error.to_string().to_ascii_lowercase(); + if [ + "certificate", + "unknown issuer", + "unknown ca", + "untrusted", + "self signed", + "self-signed", + "hostname", + "expired", + "revoked", + ] + .iter() + .any(|marker| message.contains(marker)) + { + return false; + } + + // macOS Secure Transport reports the protocol alert as "bad protocol version". + let is_macos_protocol_version_error = message.contains("bad protocol version"); + // Linux OpenSSL reports the peer's "tlsv1 alert protocol version". + let is_linux_protocol_version_error = message.contains("tlsv1 alert protocol version"); + // Windows Schannel may expose the protocol alert as a raw or formatted OS error. + let is_schannel_protocol_version_error = error + .downcast_ref::() + .and_then(std::io::Error::raw_os_error) + == Some(SCHANNEL_PROTOCOL_VERSION_ERROR) + || message.contains("(os error -2146893054)") + || message.contains("0x80090302"); + if is_macos_protocol_version_error + || is_linux_protocol_version_error + || is_schannel_protocol_version_error + { + recognized_negotiation_failure = true; + } + source = error.source(); + } + + recognized_negotiation_failure +} + +#[cfg(test)] +#[path = "tls_backend_fallback_tests.rs"] +mod tests; diff --git a/codex-rs/http-client/src/tls_backend_fallback_tests.rs b/codex-rs/http-client/src/tls_backend_fallback_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..c7ca0adceb94a6729d8657c7a7feb4e066044f68 --- /dev/null +++ b/codex-rs/http-client/src/tls_backend_fallback_tests.rs @@ -0,0 +1,216 @@ +use std::io; + +use pretty_assertions::assert_eq; + +use super::MAX_CACHED_RUSTLS_DESTINATIONS; +use super::RustlsClientCache; +use super::SCHANNEL_PROTOCOL_VERSION_ERROR; +use super::has_retryable_tls_error; +use crate::HttpClientBuilder; +use crate::OutboundProxyRoute; + +#[test] +fn recognizes_platform_specific_tls_protocol_negotiation_failures() { + let errors = [ + ("client error (Connect): bad protocol version", true), + ("BAD PROTOCOL VERSION", true), + ( + "error:0A00042E:SSL routines:ssl3_read_bytes:tlsv1 alert protocol version", + true, + ), + ("TLSV1 ALERT PROTOCOL VERSION", true), + ( + "The function requested is not supported. (os error -2146893054)", + true, + ), + ("Schannel protocol error 0x80090302", true), + ("SCHANNEL PROTOCOL ERROR 0X80090302", true), + ("certificate validation failed: bad protocol version", false), + ("bad protocol version: certificate has expired", false), + ( + "certificate validation failed: tlsv1 alert protocol version", + false, + ), + ( + "tlsv1 alert protocol version: certificate has expired", + false, + ), + ( + "certificate validation failed (os error -2146893054)", + false, + ), + ("certificate validation failed: 0x80090302", false), + ("unknown issuer", false), + ("self-signed certificate", false), + ("hostname mismatch", false), + ("connection refused", false), + ("connection reset", false), + ("dns lookup failed", false), + ("tls handshake failed", false), + ("unsupported protocol", false), + ("wrong version number", false), + ("The function requested is not supported.", false), + ( + "The client and server cannot communicate. (os error -2146893007)", + false, + ), + ("operation timed out", false), + ("407 Proxy Authentication Required", false), + ]; + + for (message, expected) in errors { + let error = io::Error::other(message); + assert_eq!(has_retryable_tls_error(&error), expected, "{message}"); + } +} + +#[test] +fn recognizes_schannel_protocol_version_error_codes() { + let protocol_error = io::Error::from_raw_os_error(SCHANNEL_PROTOCOL_VERSION_ERROR); + let another_schannel_error = io::Error::from_raw_os_error(/*code*/ -2_146_893_007); + + assert_eq!( + ( + has_retryable_tls_error(&protocol_error), + has_retryable_tls_error(&another_schannel_error), + ), + (true, false) + ); +} + +#[test] +fn certificate_errors_in_an_error_source_never_enable_fallback() { + for message in [ + "certificate verification failed: bad protocol version", + "certificate verification failed: tlsv1 alert protocol version", + "certificate verification failed (os error -2146893054)", + "certificate verification failed: 0x80090302", + ] { + let error = io::Error::other(io::Error::other(message)); + + assert!(!has_retryable_tls_error(&error), "{message}"); + } +} + +#[test] +fn rustls_fallback_decisions_are_scoped_to_origin_and_outbound_route() { + let cache = RustlsClientCache::default(); + let destination = reqwest::Url::parse("https://mcp.example.com/first").expect("valid URL"); + let same_origin = reqwest::Url::parse("https://mcp.example.com/second").expect("valid URL"); + let another_host = reqwest::Url::parse("https://another.example.com/first").expect("valid URL"); + let another_port = + reqwest::Url::parse("https://mcp.example.com:8443/first").expect("valid URL"); + let insecure = reqwest::Url::parse("http://mcp.example.com/first").expect("valid URL"); + let direct = OutboundProxyRoute::Direct; + let proxy = OutboundProxyRoute::Proxy { + url: "http://proxy.example.com".to_string(), + no_proxy: None, + }; + + let client = HttpClientBuilder::new() + .with_rustls_tls() + .build_direct() + .expect("rustls client should build without proxy autodiscovery"); + cache.remember(&destination, &direct, client); + + assert_eq!( + [ + cache.requires_rustls(&destination, &direct), + cache.requires_rustls(&same_origin, &direct), + cache.requires_rustls(&another_host, &direct), + cache.requires_rustls(&another_port, &direct), + cache.requires_rustls(&destination, &proxy), + cache.requires_rustls(&insecure, &direct), + ], + [true, true, false, false, false, false] + ); +} + +#[test] +fn cached_rustls_clients_are_reused_for_the_same_outbound_route() { + let cache = RustlsClientCache::default(); + let first_destination = + reqwest::Url::parse("https://first.example.com").expect("valid first URL"); + let second_destination = + reqwest::Url::parse("https://second.example.com").expect("valid second URL"); + let direct = OutboundProxyRoute::Direct; + let client = HttpClientBuilder::new() + .with_rustls_tls() + .build_direct() + .expect("rustls client should build without proxy autodiscovery"); + + cache.remember(&first_destination, &direct, client.clone()); + cache.remember(&second_destination, &direct, client); + + assert_eq!( + ( + cache.requires_rustls(&first_destination, &direct), + cache.requires_rustls(&second_destination, &direct), + cache.client_for_route(&direct).is_some(), + cache + .state + .lock() + .expect("rustls client cache lock") + .clients + .len(), + ), + (true, true, true, 1) + ); +} + +#[test] +fn cached_rustls_destinations_remain_bounded_while_sharing_a_route_client() { + let cache = RustlsClientCache::default(); + let client = HttpClientBuilder::new() + .with_rustls_tls() + .build_direct() + .expect("rustls client should build without proxy autodiscovery"); + + for index in 0..=MAX_CACHED_RUSTLS_DESTINATIONS { + let destination = + reqwest::Url::parse(&format!("https://mcp-{index}.example.com")).expect("valid URL"); + cache.remember(&destination, &OutboundProxyRoute::Direct, client.clone()); + } + + let state = cache.state.lock().expect("rustls client cache lock"); + assert_eq!( + (state.destinations.len(), state.clients.len()), + (MAX_CACHED_RUSTLS_DESTINATIONS, 1) + ); +} + +#[test] +fn evicting_a_destination_removes_its_unshared_route_client() { + let cache = RustlsClientCache::default(); + let client = HttpClientBuilder::new() + .with_rustls_tls() + .build_direct() + .expect("rustls client should build without proxy autodiscovery"); + + for index in 0..=MAX_CACHED_RUSTLS_DESTINATIONS { + let destination = + reqwest::Url::parse(&format!("https://mcp-{index}.example.com")).expect("valid URL"); + let route = OutboundProxyRoute::Proxy { + url: format!("http://proxy-{index}.example.com"), + no_proxy: None, + }; + cache.remember(&destination, &route, client.clone()); + } + + let state = cache.state.lock().expect("rustls client cache lock"); + assert_eq!( + ( + state.destinations.len(), + state.clients.len(), + state + .destinations + .iter() + .all(|destination| state.clients.contains_key(&destination.route)), + ), + ( + MAX_CACHED_RUSTLS_DESTINATIONS, + MAX_CACHED_RUSTLS_DESTINATIONS, + true, + ) + ); +} diff --git a/codex-rs/http-client/src/transport.rs b/codex-rs/http-client/src/transport.rs new file mode 100644 index 0000000000000000000000000000000000000000..ca351c1f0030cbe384fcb2a6eebce5819f3a51c2 --- /dev/null +++ b/codex-rs/http-client/src/transport.rs @@ -0,0 +1,252 @@ +//! Shared HTTP transport with optional limits on complete response bodies. +//! +//! Each request's limit applies to observed body bytes, including unsuccessful responses; +//! declared content lengths are only an additional early-rejection check. + +use crate::client::HttpClient; +use crate::client::RequestBuilder; +use crate::error::TransportError; +use crate::request::Request; +use crate::request::RequestBody; +use crate::request::Response; +use bytes::Bytes; +use bytes::BytesMut; +use futures::StreamExt; +use futures::stream::BoxStream; +use http::HeaderMap; +use http::Method; +use http::StatusCode; +use tracing::Level; +use tracing::enabled; +use tracing::trace; + +pub type ByteStream = BoxStream<'static, Result>; + +pub struct StreamResponse { + pub status: StatusCode, + pub headers: HeaderMap, + pub bytes: ByteStream, +} + +pub trait HttpTransport: Send + Sync { + fn execute( + &self, + req: Request, + ) -> impl std::future::Future> + Send; + fn stream( + &self, + req: Request, + ) -> impl std::future::Future> + Send; +} + +#[derive(Clone, Debug)] +pub struct ReqwestTransport { + client: HttpClient, +} + +impl ReqwestTransport { + pub fn new(client: reqwest::Client) -> Self { + Self { + client: HttpClient::new(client), + } + } + + pub fn from_http_client(client: HttpClient) -> Self { + Self { client } + } + + fn build(&self, req: Request) -> Result { + let prepared = req.prepare_body_for_send().map_err(TransportError::Build)?; + + let Request { + method, + url, + headers: _, + body: _, + compression: _, + timeout, + response_body_limit_bytes: _, + } = req; + + let mut builder = self.client.request( + Method::from_bytes(method.as_str().as_bytes()).unwrap_or(Method::GET), + &url, + ); + + if let Some(timeout) = timeout { + builder = builder.timeout(timeout); + } + + builder = builder.headers(prepared.headers); + if let Some(body) = prepared.body { + builder = builder.body(body); + } + Ok(builder) + } + + fn map_error(err: reqwest::Error) -> TransportError { + if err.is_connect() { + TransportError::Connection(err.without_url()) + } else if err.is_timeout() { + TransportError::Timeout + } else { + TransportError::Network(err.to_string()) + } + } + + fn trace_request(&self, req: &Request) { + if self.client.request_logging_enabled() && enabled!(Level::TRACE) { + trace!( + "{} to {}: {}", + req.method, + req.url, + request_body_for_trace(req) + ); + } + } +} + +fn request_body_for_trace(req: &Request) -> String { + match req.body.as_ref() { + Some(RequestBody::Json(body)) => body.to_string(), + Some(RequestBody::EncodedJson(body)) => { + String::from_utf8_lossy(body.trace_bytes()).into_owned() + } + Some(RequestBody::Raw(body)) => format!("", body.len()), + None => String::new(), + } +} + +impl HttpTransport for ReqwestTransport { + async fn execute(&self, req: Request) -> Result { + self.trace_request(&req); + + let url = req.url.clone(); + let response_body_limit_bytes = req.response_body_limit_bytes; + let builder = self.build(req)?; + let resp = builder.send().await.map_err(Self::map_error)?; + let status = resp.status(); + let headers = resp.headers().clone(); + let bytes = match response_body_limit_bytes { + Some(max_bytes) => bounded_response_bytes(resp, max_bytes).await, + None => resp.bytes().await.map_err(Self::map_error), + }; + if !status.is_success() { + let body = match bytes { + Ok(bytes) => String::from_utf8(bytes.to_vec()).ok(), + Err(error @ TransportError::ResponseTooLarge { .. }) => return Err(error), + // Keep bounded diagnostic-body failures from hiding HTTP auth/retry status. + Err(_) if response_body_limit_bytes.is_some() => None, + Err(error) => return Err(error), + }; + return Err(TransportError::Http { + status, + url: Some(url), + headers: Some(headers), + body, + }); + } + Ok(Response { + status, + headers, + body: bytes?, + }) + } + + async fn stream(&self, req: Request) -> Result { + self.trace_request(&req); + + let url = req.url.clone(); + let response_body_limit_bytes = req.response_body_limit_bytes; + let builder = self.build(req)?; + let resp = builder.send().await.map_err(Self::map_error)?; + let status = resp.status(); + let headers = resp.headers().clone(); + if !status.is_success() { + let body = match response_body_limit_bytes { + Some(max_bytes) => match bounded_response_bytes(resp, max_bytes).await { + Ok(bytes) => { + // Reuse the unbounded path's charset, BOM and replacement decoding, + // but only after the body has passed the byte limit. + let mut buffered = http::Response::new(bytes); + *buffered.headers_mut() = headers.clone(); + reqwest::Response::from(buffered).text().await.ok() + } + Err(error @ TransportError::ResponseTooLarge { .. }) => return Err(error), + // A failed diagnostic body must not hide HTTP auth or retry semantics. + Err(_) => None, + }, + None => resp.text().await.ok(), + }; + return Err(TransportError::Http { + status, + url: Some(url), + headers: Some(headers), + body, + }); + } + let bytes = match response_body_limit_bytes { + Some(max_bytes) => bounded_response_stream(resp, max_bytes)?, + None => Box::pin( + resp.bytes_stream() + .map(|result| result.map_err(Self::map_error)), + ), + }; + Ok(StreamResponse { + status, + headers, + bytes, + }) + } +} + +fn bounded_response_stream( + response: reqwest::Response, + max_bytes: usize, +) -> Result { + if response + .content_length() + .is_some_and(|length| length > max_bytes as u64) + { + return Err(TransportError::ResponseTooLarge { max_bytes }); + } + + let stream = futures::stream::try_unfold( + (response, max_bytes), + move |(mut response, remaining)| async move { + match response + .chunk() + .await + .map_err(|error| ReqwestTransport::map_error(error.without_url()))? + { + Some(chunk) if chunk.len() <= remaining => { + let remaining = remaining - chunk.len(); + Ok(Some((chunk, (response, remaining)))) + } + Some(_) => Err(TransportError::ResponseTooLarge { max_bytes }), + None => Ok(None), + } + }, + ); + Ok(Box::pin(stream)) +} + +async fn bounded_response_bytes( + response: reqwest::Response, + max_bytes: usize, +) -> Result { + let mut stream = bounded_response_stream(response, max_bytes)?; + let mut body = BytesMut::new(); + while let Some(chunk) = stream.next().await { + body.extend_from_slice(&chunk?); + } + Ok(body.freeze()) +} + +#[cfg(test)] +#[path = "transport_tests.rs"] +mod tests; + +#[cfg(test)] +#[path = "transport_limit_tests.rs"] +mod limit_tests; diff --git a/codex-rs/http-client/src/transport_limit_tests.rs b/codex-rs/http-client/src/transport_limit_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..1c977854a1af6483328251f7805ee69f646e0ccd --- /dev/null +++ b/codex-rs/http-client/src/transport_limit_tests.rs @@ -0,0 +1,452 @@ +//! HTTP fixtures exercise opt-in body limits for both buffered and streaming callers. + +use super::*; +use pretty_assertions::assert_eq; +use std::io::Read; +use std::io::Write; +use std::net::TcpListener; +use std::net::TcpStream; +use std::time::Duration; +use std::time::Instant; + +#[derive(Clone, Copy, Debug)] +enum Delivery { + Buffered, + Streamed, +} + +#[tokio::test] +async fn accepts_complete_bodies_at_the_limit() { + for delivery in [Delivery::Buffered, Delivery::Streamed] { + for (framing, wire_body, max_bytes, expected_body) in [ + ("Content-Length: 4\r\n", "test", 4, "test"), + ("", "test", 4, "test"), + ( + "Transfer-Encoding: chunked\r\n", + "2\r\nte\r\n2\r\nst\r\n0\r\n\r\n", + 4, + "test", + ), + ("Content-Length: 0\r\n", "", 0, ""), + ] { + let response = receive_response( + format!("HTTP/1.1 200 OK\r\n{framing}Connection: close\r\n\r\n{wire_body}"), + &transport(), + Some(max_bytes), + delivery, + ) + .await + .expect("a complete body within the limit should succeed"); + let mut expected_headers = HeaderMap::new(); + expected_headers.insert("connection", "close".parse().unwrap()); + for line in framing.lines() { + let (name, value) = line.split_once(':').unwrap(); + expected_headers.insert( + http::HeaderName::from_bytes(name.as_bytes()).unwrap(), + value.trim().parse().unwrap(), + ); + } + assert_eq!( + (response.status, response.headers, response.body), + ( + StatusCode::OK, + expected_headers, + Bytes::copy_from_slice(expected_body.as_bytes()), + ), + "{delivery:?} with framing {framing:?}", + ); + } + } +} + +#[tokio::test] +async fn rejects_oversized_success_and_error_bodies_without_sensitive_diagnostics() { + for delivery in [Delivery::Buffered, Delivery::Streamed] { + for status in ["200 OK", "500 Internal Server Error"] { + for (framing, wire_body) in [ + ("Content-Length: 11\r\n", "body-secret"), + ("", "body-secret"), + ( + "Transfer-Encoding: chunked\r\n", + "4\r\nbody\r\n7\r\n-secret\r\n0\r\n\r\n", + ), + ( + "Transfer-Encoding: chunked\r\nContent-Length: 1\r\n", + "4\r\nbody\r\n7\r\n-secret\r\n0\r\n\r\n", + ), + ] { + let error = receive_response( + format!( + "HTTP/1.1 {status}\r\n{framing}X-Secret: header-secret\r\nConnection: close\r\n\r\n{wire_body}" + ), + &transport(), + Some(4), + delivery, + ) + .await + .expect_err("the observed body must not exceed the limit"); + assert_eq!( + (error.to_string(), format!("{error:?}")), + ( + "response body exceeds the 4 byte limit".to_string(), + "ResponseTooLarge { max_bytes: 4 }".to_string(), + ), + "{delivery:?}, {status}, framing {framing:?}", + ); + } + } + } +} + +#[tokio::test] +async fn rejects_oversized_declared_lengths_before_reading_the_body() { + for delivery in [Delivery::Buffered, Delivery::Streamed] { + for status in [ + StatusCode::OK, + StatusCode::UNAUTHORIZED, + StatusCode::TOO_MANY_REQUESTS, + ] { + let error = receive_response( + format!( + "HTTP/1.1 {status}\r\nContent-Length: 1000000\r\nConnection: close\r\n\r\n" + ), + &transport(), + Some(4), + delivery, + ) + .await + .expect_err("an oversized declaration should fail before the incomplete body"); + assert!(matches!( + error, + TransportError::ResponseTooLarge { max_bytes: 4 } + )); + } + } +} + +#[tokio::test] +async fn streaming_limit_counts_across_chunks_and_terminates_on_error() { + let (continue_tx, continue_rx) = std::sync::mpsc::channel(); + let (url, server) = start_server(move |connection| { + connection + .write_all(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\npart\r\n") + .expect("write the first chunk"); + continue_rx + .recv_timeout(Duration::from_secs(5)) + .expect("client should consume the first chunk"); + connection + .write_all(b"4\r\nnext\r\n0\r\n\r\n") + .expect("write the second chunk"); + }); + let mut request = Request::new(Method::GET, url); + request.response_body_limit_bytes = Some(6); + let mut response = transport() + .stream(request) + .await + .expect("stream should open before the body limit is reached"); + assert_eq!( + response.bytes.next().await.unwrap().unwrap(), + Bytes::from_static(b"part"), + ); + continue_tx.send(()).expect("release the second chunk"); + let error = response + .bytes + .next() + .await + .expect("second chunk should produce the limit error") + .expect_err("two individually small chunks should exceed the total limit"); + assert!(matches!( + error, + TransportError::ResponseTooLarge { max_bytes: 6 } + )); + assert!(response.bytes.next().await.is_none()); + server.join().expect("HTTP fixture should finish"); +} + +#[tokio::test] +async fn preserves_http_failures_within_the_limit() { + for delivery in [Delivery::Buffered, Delivery::Streamed] { + let error = receive_response( + "HTTP/1.1 403 Forbidden\r\nContent-Length: 6\r\nConnection: close\r\n\r\ndenied" + .to_string(), + &transport(), + Some(6), + delivery, + ) + .await + .expect_err("a small error body should retain HTTP error semantics"); + let TransportError::Http { + status, + headers, + body, + .. + } = error + else { + panic!("expected HTTP error, got {error:?}"); + }; + assert_eq!( + (status, headers, body), + ( + StatusCode::FORBIDDEN, + Some(HeaderMap::from_iter([ + (http::header::CONTENT_LENGTH, "6".parse().unwrap()), + (http::header::CONNECTION, "close".parse().unwrap()), + ])), + Some("denied".to_string()), + ), + ); + } +} + +#[tokio::test] +async fn bounded_errors_preserve_status_when_the_body_is_interrupted() { + for delivery in [Delivery::Buffered, Delivery::Streamed] { + for status in [StatusCode::UNAUTHORIZED, StatusCode::TOO_MANY_REQUESTS] { + for max_bytes in [None, Some(5), Some(4), Some(3)] { + // The first chunk is complete, but the peer closes before the next chunk. + let error = receive_response( + format!( + "HTTP/1.1 {status}\r\nTransfer-Encoding: chunked\r\nRetry-After: 7\r\nWWW-Authenticate: Bearer\r\nConnection: close\r\n\r\n4\r\npart\r\n" + ), + &transport(), + max_bytes, + delivery, + ) + .await + .expect_err("an interrupted error body must still fail"); + if matches!(delivery, Delivery::Buffered) && max_bytes.is_none() { + assert!(matches!(error, TransportError::Network(_))); + continue; + } + if max_bytes == Some(3) { + assert!(matches!( + error, + TransportError::ResponseTooLarge { max_bytes: 3 } + )); + continue; + } + let TransportError::Http { + status: actual_status, + headers, + body, + .. + } = error + else { + panic!("expected preserved HTTP error, got {error:?}"); + }; + assert_eq!( + (actual_status, headers, body), + ( + status, + Some(HeaderMap::from_iter([ + (http::header::TRANSFER_ENCODING, "chunked".parse().unwrap()), + (http::header::RETRY_AFTER, "7".parse().unwrap()), + (http::header::WWW_AUTHENTICATE, "Bearer".parse().unwrap()), + (http::header::CONNECTION, "close".parse().unwrap()), + ])), + None, + ), + "{delivery:?}, max_bytes={max_bytes:?}", + ); + } + } + } +} + +#[tokio::test] +async fn bounded_successes_report_interrupted_bodies_as_network_errors() { + for delivery in [Delivery::Buffered, Delivery::Streamed] { + let error = receive_response( + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n4\r\npart\r\n" + .to_string(), + &transport(), + Some(4), + delivery, + ) + .await + .expect_err("a successful status cannot hide an incomplete response body"); + assert!(matches!(error, TransportError::Network(_))); + } +} + +#[tokio::test] +async fn bounded_streamed_errors_preserve_text_decoding() { + for (content_type, bytes, expected) in [ + ( + "text/plain; charset=windows-1252", + [0x63, 0x61, 0x66, 0xe9].as_slice(), + "café", + ), + ( + "text/plain; charset=utf-8", + b"error: \xff".as_slice(), + "error: �", + ), + ( + "text/plain; charset=unknown", + b"error: \xff".as_slice(), + "error: �", + ), + ("text/plain", b"\xef\xbb\xbfdenied".as_slice(), "denied"), + ] { + for bounded in [false, true] { + let mut wire = format!( + "HTTP/1.1 403 Forbidden\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + bytes.len() + ) + .into_bytes(); + wire.extend_from_slice(bytes); + let error = receive_response( + wire, + &transport(), + bounded.then_some(bytes.len()), + Delivery::Streamed, + ) + .await + .expect_err("retain the HTTP error and decoded diagnostics"); + let TransportError::Http { + status, + headers, + body, + .. + } = error + else { + panic!("expected HTTP error, got {error:?}"); + }; + assert_eq!( + (status, headers, body), + ( + StatusCode::FORBIDDEN, + Some(HeaderMap::from_iter([ + (http::header::CONTENT_TYPE, content_type.parse().unwrap()), + ( + http::header::CONTENT_LENGTH, + bytes.len().to_string().parse().unwrap() + ), + (http::header::CONNECTION, "close".parse().unwrap()), + ])), + Some(expected.to_string()), + ), + "{content_type}, bounded={bounded}", + ); + } + } +} + +#[tokio::test] +async fn shared_transport_keeps_each_requests_response_limit_independent() { + let transport = transport(); + for delivery in [Delivery::Buffered, Delivery::Streamed] { + for max_bytes in [Some(4), None, Some(11), Some(0), None] { + let response = receive_response( + "HTTP/1.1 200 OK\r\nContent-Length: 11\r\nConnection: close\r\n\r\nbody-secret" + .to_string(), + &transport, + max_bytes, + delivery, + ) + .await; + match max_bytes { + Some(max_bytes) if max_bytes < 11 => assert!(matches!( + response, + Err(TransportError::ResponseTooLarge { max_bytes: actual }) if actual == max_bytes + )), + _ => assert_eq!( + response.expect("only this request's limit applies").body, + Bytes::from_static(b"body-secret"), + ), + } + } + } +} + +fn transport() -> ReqwestTransport { + ReqwestTransport::new( + reqwest::Client::builder() + .no_proxy() + .build() + .expect("test client should build"), + ) +} + +async fn receive_response( + wire_response: impl Into>, + transport: &ReqwestTransport, + response_body_limit_bytes: Option, + delivery: Delivery, +) -> Result { + let wire_response = wire_response.into(); + let (url, server) = start_server(move |connection| { + connection + .write_all(&wire_response) + .expect("write HTTP response"); + }); + let mut request = Request::new(Method::GET, url); + request.headers.insert( + http::header::AUTHORIZATION, + "Bearer auth-secret".parse().unwrap(), + ); + request.timeout = Some(Duration::from_secs(5)); + request.response_body_limit_bytes = response_body_limit_bytes; + let result = match delivery { + Delivery::Buffered => transport.execute(request).await, + Delivery::Streamed => { + async { + let mut response = transport.stream(request).await?; + let mut body = BytesMut::new(); + while let Some(chunk) = response.bytes.next().await { + body.extend_from_slice(&chunk?); + } + Ok(Response { + status: response.status, + headers: response.headers, + body: body.freeze(), + }) + } + .await + } + }; + server.join().expect("HTTP fixture should finish"); + result +} + +fn start_server( + respond: impl FnOnce(&mut TcpStream) + Send + 'static, +) -> (String, std::thread::JoinHandle<()>) { + let listener = TcpListener::bind(("127.0.0.1", 0)).expect("HTTP fixture should bind"); + let address = listener.local_addr().expect("HTTP fixture address"); + listener + .set_nonblocking(true) + .expect("nonblocking listener"); + let server = std::thread::spawn(move || { + let deadline = Instant::now() + Duration::from_secs(5); + let (mut connection, _) = loop { + match listener.accept() { + Ok(connection) => break connection, + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { + assert!(Instant::now() < deadline, "HTTP fixture accept timed out"); + std::thread::sleep(Duration::from_millis(1)); + } + Err(error) => panic!("HTTP fixture should accept: {error}"), + } + }; + connection + .set_nonblocking(false) + .expect("HTTP fixture connection should use blocking reads"); + connection + .set_read_timeout(Some(Duration::from_secs(5))) + .expect("HTTP fixture read timeout"); + let mut request = Vec::new(); + let mut buffer = [0; 1024]; + while !request.windows(4).any(|window| window == b"\r\n\r\n") { + let count = connection.read(&mut buffer).expect("read HTTP request"); + assert_ne!(count, 0, "HTTP request should contain complete headers"); + request.extend_from_slice(&buffer[..count]); + } + respond(&mut connection); + }); + ( + format!("http://{address}/models?credential=url-secret"), + server, + ) +} diff --git a/codex-rs/http-client/src/transport_tests.rs b/codex-rs/http-client/src/transport_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..a8ab11aba5cfe6830daaf1dd00a7db698fe73f4e --- /dev/null +++ b/codex-rs/http-client/src/transport_tests.rs @@ -0,0 +1,114 @@ +use super::*; +use serde_json::json; +use std::io::Write; +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; +use tracing_subscriber::Layer; +use tracing_subscriber::layer::SubscriberExt; + +#[tokio::test] +async fn enabled_request_logging_emits_transport_url_and_body() { + let logs = capture_transport_logs(HttpClient::new(test_reqwest_client())).await; + + assert!(logs.contains("log capture sentinel")); + assert!(logs.contains("url-secret")); + assert!(logs.contains("body-secret")); +} + +#[tokio::test] +async fn disabled_request_logging_suppresses_transport_url_and_body() { + let logs = capture_transport_logs(HttpClient::new_without_request_logging( + test_reqwest_client(), + )) + .await; + + assert!(logs.contains("log capture sentinel")); + assert!(!logs.contains("url-secret")); + assert!(!logs.contains("body-secret")); +} + +#[tokio::test] +async fn connection_failures_are_classified_without_exposing_request_urls() { + let unavailable_server = + std::net::TcpListener::bind(("127.0.0.1", 0)).expect("server port should bind"); + let server_addr = unavailable_server + .local_addr() + .expect("server listener should have an address"); + drop(unavailable_server); + let transport = ReqwestTransport::from_http_client(HttpClient::new(test_reqwest_client())); + let request = Request::new( + Method::POST, + format!("http://{server_addr}/responses?token=url-secret"), + ); + + let error = match transport.stream(request).await { + Err(TransportError::Connection(error)) => error, + Err(error) => panic!("expected a connection failure, got {error}"), + Ok(_) => panic!("an unavailable server should not return a response"), + }; + assert!(!error.to_string().contains("url-secret")); +} + +fn test_reqwest_client() -> reqwest::Client { + reqwest::Client::builder() + .no_proxy() + .build() + .expect("HTTP client should build") +} + +async fn capture_transport_logs(client: HttpClient) -> String { + let unavailable_server = + std::net::TcpListener::bind(("127.0.0.1", 0)).expect("server port should bind"); + let server_addr = unavailable_server + .local_addr() + .expect("server listener should have an address"); + drop(unavailable_server); + let transport = ReqwestTransport::from_http_client(client); + let log_buffer = Arc::new(Mutex::new(Vec::new())); + let writer_buffer = Arc::clone(&log_buffer); + let subscriber = tracing_subscriber::registry().with( + tracing_subscriber::fmt::layer() + .with_ansi(false) + .with_writer(move || TestLogWriter(Arc::clone(&writer_buffer))) + .with_filter( + tracing_subscriber::filter::Targets::new() + .with_target("codex_http_client::transport", tracing::Level::TRACE), + ), + ); + let _guard = tracing::subscriber::set_default(subscriber); + tracing::trace!(target: "codex_http_client::transport", "log capture sentinel"); + let mut request = Request::new( + Method::POST, + format!("http://{server_addr}/request?token=url-secret"), + ) + .with_json(&json!({"token": "body-secret"})); + request.timeout = Some(Duration::from_secs(1)); + + let _ = transport.execute(request).await; + + String::from_utf8( + log_buffer + .lock() + .expect("log buffer should not be poisoned") + .clone(), + ) + .expect("captured logs should be UTF-8") +} + +#[derive(Clone)] +struct TestLogWriter(Arc>>); + +impl Write for TestLogWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.0 + .lock() + .map_err(|_| std::io::Error::other("log buffer should not be poisoned"))? + .extend_from_slice(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} diff --git a/codex-rs/http-client/src/windows_tls.rs b/codex-rs/http-client/src/windows_tls.rs new file mode 100644 index 0000000000000000000000000000000000000000..5b42b37fc31a4b5fb604d866e93200b19d1dd5b3 --- /dev/null +++ b/codex-rs/http-client/src/windows_tls.rs @@ -0,0 +1,22 @@ +//! Windows certificate validation for callers without a custom CA bundle. +//! +//! Unlike a snapshot of installed roots, Windows can retrieve missing trusted roots on demand. +//! Callers must retain the custom-CA configuration path when a bundle is configured. + +use std::sync::Arc; + +use codex_utils_rustls_provider::ensure_rustls_crypto_provider; +use rustls::ClientConfig; +use rustls_platform_verifier::ConfigVerifierExt; + +/// Builds a TLS configuration that validates server certificates using Windows trust policy. +/// +/// This does not load Codex custom CA settings; callers must check those before selecting it. +pub fn build_windows_platform_tls_config() -> Result, rustls::Error> { + ensure_rustls_crypto_provider(); + ClientConfig::with_platform_verifier().map(Arc::new) +} + +#[cfg(test)] +#[path = "windows_tls_tests.rs"] +mod tests; diff --git a/codex-rs/http-client/src/windows_tls_tests.rs b/codex-rs/http-client/src/windows_tls_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..6ddf3b50f754a3c2ffa61a906600472daf2cc1f9 --- /dev/null +++ b/codex-rs/http-client/src/windows_tls_tests.rs @@ -0,0 +1,58 @@ +//! Certificate rejection coverage for Windows platform TLS. + +use super::build_windows_platform_tls_config; +use std::sync::Arc; + +#[test] +fn platform_tls_rejects_a_self_signed_server() { + let client_config = + build_windows_platform_tls_config().expect("configure platform verification"); + let certified_key = rcgen::generate_simple_self_signed(vec!["localhost".to_string()]) + .expect("generate untrusted certificate"); + let server_config = rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert( + vec![certified_key.cert.der().clone()], + certified_key.signing_key.into(), + ) + .expect("configure test server"); + let mut client = rustls::ClientConnection::new( + client_config, + "localhost".try_into().expect("valid server name"), + ) + .expect("create client"); + let mut server = rustls::ServerConnection::new(Arc::new(server_config)).expect("create server"); + + // Drive the real handshake without opening ports or changing system trust. + for _ in 0..10 { + let mut outbound = Vec::new(); + client + .write_tls(&mut outbound) + .expect("write client records"); + server + .read_tls(&mut outbound.as_slice()) + .expect("read client records"); + server + .process_new_packets() + .expect("process client records"); + outbound.clear(); + server + .write_tls(&mut outbound) + .expect("write server records"); + client + .read_tls(&mut outbound.as_slice()) + .expect("read server records"); + if let Err(error) = client.process_new_packets() { + assert!( + matches!(error, rustls::Error::InvalidCertificate(_)), + "{error:?}" + ); + return; + } + assert!( + client.is_handshaking(), + "untrusted server must not complete TLS" + ); + } + panic!("expected certificate rejection during handshake"); +} diff --git a/codex-rs/http-client/tests/ca_env.rs b/codex-rs/http-client/tests/ca_env.rs new file mode 100644 index 0000000000000000000000000000000000000000..f35225fee778e80858daca7e5c9791ee0ec3d3ed --- /dev/null +++ b/codex-rs/http-client/tests/ca_env.rs @@ -0,0 +1,540 @@ +#![allow(clippy::expect_used)] +//! Subprocess coverage for custom CA behavior that must build a real reqwest client. +//! +//! These tests intentionally run through `custom_ca_probe` and +//! `build_reqwest_client_for_subprocess_tests` instead of calling the helper in-process. The +//! detailed explanation of what "hermetic" means here lives in `codex_http_client::custom_ca`; these +//! tests add the process-level half of that contract by scrubbing inherited CA environment +//! variables before each subprocess launch. Most assertions here cover CA file selection, PEM +//! parsing, and user-facing errors. The HTTPS probes go further and perform real POSTs against +//! locally generated certificates, including through a TLS-intercepting CONNECT proxy. + +use codex_utils_cargo_bin::cargo_bin; +use rcgen::BasicConstraints; +use rcgen::CertificateParams; +use rcgen::CertifiedIssuer; +use rcgen::DistinguishedName; +use rcgen::DnType; +use rcgen::ExtendedKeyUsagePurpose; +use rcgen::IsCa; +use rcgen::KeyPair; +use rcgen::KeyUsagePurpose; +use rcgen::PKCS_ECDSA_P256_SHA256; +use rustls_pki_types::CertificateDer; +use rustls_pki_types::PrivateKeyDer; +use std::fs; +use std::io; +use std::io::Read; +use std::io::Write; +use std::net::TcpListener; +use std::net::TcpStream; +use std::path::Path; +use std::path::PathBuf; +use std::process::Command; +use std::sync::Arc; +use std::sync::mpsc; +use std::thread; +use std::time::Duration; +use std::time::Instant; +use tempfile::TempDir; + +const CODEX_CA_CERT_ENV: &str = "CODEX_CA_CERTIFICATE"; +const PROBE_PROXY_ENV: &str = "CODEX_CUSTOM_CA_PROBE_PROXY"; +const PROBE_TLS13_ENV: &str = "CODEX_CUSTOM_CA_PROBE_TLS13"; +const PROBE_URL_ENV: &str = "CODEX_CUSTOM_CA_PROBE_URL"; +const SSL_CERT_FILE_ENV: &str = "SSL_CERT_FILE"; +const PROXY_ENV_VARS: &[&str] = &[ + "HTTP_PROXY", + "http_proxy", + "HTTPS_PROXY", + "https_proxy", + "ALL_PROXY", + "all_proxy", + "NO_PROXY", + "no_proxy", +]; + +const TEST_CERT_1: &str = include_str!("fixtures/test-ca.pem"); +const TEST_CERT_2: &str = include_str!("fixtures/test-intermediate.pem"); +const TRUSTED_TEST_CERT: &str = include_str!("fixtures/test-ca-trusted.pem"); + +struct Tls13Material { + ca_cert_pem: String, + server_cert: CertificateDer<'static>, + server_key: PrivateKeyDer<'static>, +} + +struct Tls13TestServer { + ca_cert_pem: String, + request_rx: mpsc::Receiver>, + url: String, +} + +struct PlainHttpOrigin { + request_rx: mpsc::Receiver>, + url: String, +} + +struct TlsInterceptingProxy { + ca_cert_pem: String, + request_rx: mpsc::Receiver>, + url: String, +} + +fn write_cert_file(temp_dir: &TempDir, name: &str, contents: &str) -> PathBuf { + let path = temp_dir.path().join(name); + fs::write(&path, contents).expect("certificate fixture should be writable"); + path +} + +fn probe_command() -> Command { + let mut cmd = Command::new( + cargo_bin("custom_ca_probe").expect("custom_ca_probe binary should be available"), + ); + // `Command` inherits the parent environment by default, so scrub CA-related variables first or + // these tests can accidentally pass/fail based on the developer shell or CI runner. + cmd.env_remove(CODEX_CA_CERT_ENV); + cmd.env_remove(PROBE_PROXY_ENV); + cmd.env_remove(PROBE_TLS13_ENV); + cmd.env_remove(PROBE_URL_ENV); + cmd.env_remove(SSL_CERT_FILE_ENV); + for env_var in PROXY_ENV_VARS { + cmd.env_remove(env_var); + } + cmd +} + +fn run_probe(envs: &[(&str, &Path)]) -> std::process::Output { + let mut cmd = probe_command(); + for (key, value) in envs { + cmd.env(key, value); + } + cmd.output().expect("custom_ca_probe should run") +} + +fn run_probe_posting_to_tls13_server(envs: &[(&str, &Path)], url: &str) -> std::process::Output { + let mut cmd = probe_command(); + for (key, value) in envs { + cmd.env(key, value); + } + cmd.env(PROBE_TLS13_ENV, "1"); + cmd.env(PROBE_URL_ENV, url); + cmd.output().expect("custom_ca_probe should run") +} + +fn run_probe_posting_through_tls_intercepting_proxy( + envs: &[(&str, &Path)], + url: &str, + proxy_url: &str, +) -> std::process::Output { + let mut cmd = probe_command(); + for (key, value) in envs { + cmd.env(key, value); + } + cmd.env(PROBE_PROXY_ENV, proxy_url); + cmd.env(PROBE_TLS13_ENV, "1"); + cmd.env(PROBE_URL_ENV, url); + cmd.output().expect("custom_ca_probe should run") +} + +fn spawn_tls13_test_server() -> Tls13TestServer { + codex_utils_rustls_provider::ensure_rustls_crypto_provider(); + let material = generate_tls13_material(); + let listener = TcpListener::bind(("127.0.0.1", 0)).expect("TLS test server should bind"); + listener + .set_nonblocking(true) + .expect("TLS test server should become nonblocking"); + let port = listener + .local_addr() + .expect("TLS test server should have a local address") + .port(); + let config = Arc::new( + rustls::ServerConfig::builder_with_protocol_versions(&[&rustls::version::TLS13]) + .with_no_client_auth() + .with_single_cert(vec![material.server_cert], material.server_key) + .expect("TLS 1.3 server config should be valid"), + ); + let (request_tx, request_rx) = mpsc::channel(); + + thread::spawn(move || { + let result = accept_tls13_request(listener, config); + let _ = request_tx.send(result.map_err(|error| error.to_string())); + }); + + Tls13TestServer { + ca_cert_pem: material.ca_cert_pem, + request_rx, + url: format!("https://127.0.0.1:{port}/oauth/token"), + } +} + +fn spawn_plain_http_origin() -> PlainHttpOrigin { + let listener = TcpListener::bind(("127.0.0.1", 0)).expect("plain HTTP origin should bind"); + listener + .set_nonblocking(true) + .expect("plain HTTP origin should become nonblocking"); + let port = listener + .local_addr() + .expect("plain HTTP origin should have a local address") + .port(); + let (request_tx, request_rx) = mpsc::channel(); + + thread::spawn(move || { + let result = accept_plain_http_origin_request(listener); + let _ = request_tx.send(result.map_err(|error| error.to_string())); + }); + + PlainHttpOrigin { + request_rx, + url: format!("https://127.0.0.1:{port}/oauth/token"), + } +} + +fn spawn_tls_intercepting_proxy() -> TlsInterceptingProxy { + codex_utils_rustls_provider::ensure_rustls_crypto_provider(); + let material = generate_tls13_material(); + let listener = TcpListener::bind(("127.0.0.1", 0)).expect("TLS intercepting proxy should bind"); + listener + .set_nonblocking(true) + .expect("TLS intercepting proxy should become nonblocking"); + let port = listener + .local_addr() + .expect("TLS intercepting proxy should have a local address") + .port(); + let config = Arc::new( + rustls::ServerConfig::builder_with_protocol_versions(&[&rustls::version::TLS13]) + .with_no_client_auth() + .with_single_cert(vec![material.server_cert], material.server_key) + .expect("TLS intercepting proxy config should be valid"), + ); + let (request_tx, request_rx) = mpsc::channel(); + + thread::spawn(move || { + let result = accept_tls_intercepting_proxy_request(listener, config); + let _ = request_tx.send(result.map_err(|error| error.to_string())); + }); + + TlsInterceptingProxy { + ca_cert_pem: material.ca_cert_pem, + request_rx, + url: format!("http://127.0.0.1:{port}"), + } +} + +fn generate_tls13_material() -> Tls13Material { + let mut ca_params = CertificateParams::default(); + ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + ca_params.key_usages = vec![KeyUsagePurpose::KeyCertSign, KeyUsagePurpose::CrlSign]; + let mut ca_distinguished_name = DistinguishedName::new(); + ca_distinguished_name.push(DnType::CommonName, "codex test CA"); + ca_params.distinguished_name = ca_distinguished_name; + let ca_key_pair = + KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256).expect("test CA key pair should generate"); + let ca = CertifiedIssuer::self_signed(ca_params, ca_key_pair) + .expect("test CA certificate should generate"); + + let mut server_params = + CertificateParams::new(vec!["localhost".to_string(), "127.0.0.1".to_string()]) + .expect("test server certificate params should be valid"); + server_params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ServerAuth]; + server_params.key_usages = vec![ + KeyUsagePurpose::DigitalSignature, + KeyUsagePurpose::KeyEncipherment, + ]; + let server_key_pair = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256) + .expect("test server key pair should generate"); + let server_cert = server_params + .signed_by(&server_key_pair, &ca) + .expect("test server certificate should generate"); + + Tls13Material { + ca_cert_pem: ca.pem(), + server_cert: server_cert.der().clone(), + server_key: PrivateKeyDer::from(server_key_pair), + } +} + +fn accept_plain_http_origin_request(listener: TcpListener) -> io::Result { + let mut stream = accept_with_timeout(listener, Duration::from_secs(5))?; + stream.set_nonblocking(false)?; + stream.set_read_timeout(Some(Duration::from_secs(5)))?; + stream.set_write_timeout(Some(Duration::from_secs(5)))?; + + let request = read_http_message(&mut stream)?; + stream.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok")?; + stream.flush()?; + Ok(request) +} + +fn accept_tls13_request( + listener: TcpListener, + config: Arc, +) -> io::Result { + let stream = accept_with_timeout(listener, Duration::from_secs(5))?; + stream.set_nonblocking(false)?; + stream.set_read_timeout(Some(Duration::from_secs(5)))?; + stream.set_write_timeout(Some(Duration::from_secs(5)))?; + + let connection = rustls::ServerConnection::new(config).map_err(io::Error::other)?; + let mut tls = rustls::StreamOwned::new(connection, stream); + let request = read_http_message(&mut tls)?; + tls.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok")?; + tls.flush()?; + Ok(request) +} + +fn accept_tls_intercepting_proxy_request( + listener: TcpListener, + config: Arc, +) -> io::Result { + let mut stream = accept_with_timeout(listener, Duration::from_secs(5))?; + stream.set_nonblocking(false)?; + stream.set_read_timeout(Some(Duration::from_secs(5)))?; + stream.set_write_timeout(Some(Duration::from_secs(5)))?; + + let connect_request = read_http_message(&mut stream)?; + let origin_authority = connect_authority_from_request(&connect_request)?; + stream.write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")?; + stream.flush()?; + + let connection = rustls::ServerConnection::new(config).map_err(io::Error::other)?; + let mut tls = rustls::StreamOwned::new(connection, stream); + let request = read_http_message(&mut tls)?; + + let mut origin = TcpStream::connect(origin_authority.as_str())?; + origin.set_read_timeout(Some(Duration::from_secs(5)))?; + origin.set_write_timeout(Some(Duration::from_secs(5)))?; + origin.write_all(request.as_bytes())?; + origin.flush()?; + let response = read_http_message(&mut origin)?; + + tls.write_all(response.as_bytes())?; + tls.flush()?; + Ok(request) +} + +fn connect_authority_from_request(request: &str) -> io::Result { + let request_line = request + .lines() + .next() + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "empty CONNECT request"))?; + let mut parts = request_line.split_whitespace(); + match (parts.next(), parts.next(), parts.next()) { + (Some("CONNECT"), Some(authority), Some(_version)) => Ok(authority.to_string()), + _ => Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("invalid CONNECT request line: {request_line}"), + )), + } +} + +fn accept_with_timeout(listener: TcpListener, timeout: Duration) -> io::Result { + let deadline = Instant::now() + timeout; + loop { + match listener.accept() { + Ok((stream, _)) => return Ok(stream), + Err(error) if error.kind() == io::ErrorKind::WouldBlock => { + if Instant::now() >= deadline { + return Err(io::Error::new( + io::ErrorKind::TimedOut, + "timed out waiting for TLS test client", + )); + } + thread::sleep(Duration::from_millis(10)); + } + Err(error) => return Err(error), + } + } +} + +fn read_http_message(stream: &mut impl Read) -> io::Result { + let mut buffer = Vec::new(); + let mut chunk = [0; 1024]; + loop { + let bytes_read = stream.read(&mut chunk)?; + if bytes_read == 0 { + break; + } + buffer.extend_from_slice(&chunk[..bytes_read]); + if let Some(header_end) = buffer.windows(4).position(|window| window == b"\r\n\r\n") { + let body_start = header_end + 4; + let headers = String::from_utf8_lossy(&buffer[..body_start]); + let content_length = headers + .lines() + .filter_map(|line| line.split_once(':')) + .find_map(|(name, value)| { + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + if buffer.len() >= body_start + content_length { + break; + } + } + } + Ok(String::from_utf8_lossy(&buffer).into_owned()) +} + +fn assert_token_exchange_request(request: &str) { + assert!( + request.starts_with("POST /oauth/token HTTP/1.1"), + "unexpected request:\n{request}" + ); + assert!( + request.contains("grant_type=authorization_code&code=test"), + "unexpected request body:\n{request}" + ); +} + +#[test] +fn uses_codex_ca_cert_env() { + let temp_dir = TempDir::new().expect("tempdir"); + let cert_path = write_cert_file(&temp_dir, "ca.pem", TEST_CERT_1); + + let output = run_probe(&[(CODEX_CA_CERT_ENV, cert_path.as_path())]); + + assert!(output.status.success()); +} + +#[test] +fn falls_back_to_ssl_cert_file() { + let temp_dir = TempDir::new().expect("tempdir"); + let cert_path = write_cert_file(&temp_dir, "ssl.pem", TEST_CERT_1); + + let output = run_probe(&[(SSL_CERT_FILE_ENV, cert_path.as_path())]); + + assert!(output.status.success()); +} + +#[test] +fn prefers_codex_ca_cert_over_ssl_cert_file() { + let temp_dir = TempDir::new().expect("tempdir"); + let cert_path = write_cert_file(&temp_dir, "ca.pem", TEST_CERT_1); + let bad_path = write_cert_file(&temp_dir, "bad.pem", ""); + + let output = run_probe(&[ + (CODEX_CA_CERT_ENV, cert_path.as_path()), + (SSL_CERT_FILE_ENV, bad_path.as_path()), + ]); + + assert!(output.status.success()); +} + +#[test] +fn handles_multi_certificate_bundle() { + let temp_dir = TempDir::new().expect("tempdir"); + let bundle = format!("{TEST_CERT_1}\n{TEST_CERT_2}"); + let cert_path = write_cert_file(&temp_dir, "bundle.pem", &bundle); + + let output = run_probe(&[(CODEX_CA_CERT_ENV, cert_path.as_path())]); + + assert!(output.status.success()); +} + +#[test] +fn posts_to_tls13_server_using_custom_ca_bundle() { + let temp_dir = TempDir::new().expect("tempdir"); + let server = spawn_tls13_test_server(); + let cert_path = write_cert_file(&temp_dir, "tls-ca.pem", &server.ca_cert_pem); + + let output = + run_probe_posting_to_tls13_server(&[(CODEX_CA_CERT_ENV, cert_path.as_path())], &server.url); + let server_result = server.request_rx.recv_timeout(Duration::from_secs(5)); + + assert!( + output.status.success(), + "custom_ca_probe failed\nstdout:\n{}\nstderr:\n{}\nserver:\n{server_result:?}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + let request = server_result + .expect("TLS test server should report a request") + .expect("TLS test server should accept the probe request"); + assert_token_exchange_request(&request); +} + +#[test] +fn posts_to_token_origin_through_tls_intercepting_proxy_with_custom_ca_bundle() { + let temp_dir = TempDir::new().expect("tempdir"); + let origin = spawn_plain_http_origin(); + let proxy = spawn_tls_intercepting_proxy(); + let cert_path = write_cert_file(&temp_dir, "proxy-ca.pem", &proxy.ca_cert_pem); + + let output = run_probe_posting_through_tls_intercepting_proxy( + &[(CODEX_CA_CERT_ENV, cert_path.as_path())], + &origin.url, + &proxy.url, + ); + let proxy_result = proxy.request_rx.recv_timeout(Duration::from_secs(5)); + let origin_result = origin.request_rx.recv_timeout(Duration::from_secs(5)); + + assert!( + output.status.success(), + "custom_ca_probe failed\nstdout:\n{}\nstderr:\n{}\nproxy:\n{proxy_result:?}\norigin:\n{origin_result:?}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + let proxy_request = proxy_result + .expect("TLS intercepting proxy should report a request") + .expect("TLS intercepting proxy should accept the probe request"); + let origin_request = origin_result + .expect("plain HTTP origin should report a request") + .expect("plain HTTP origin should accept the forwarded request"); + assert_token_exchange_request(&proxy_request); + assert_token_exchange_request(&origin_request); +} + +#[test] +fn rejects_empty_pem_file_with_hint() { + let temp_dir = TempDir::new().expect("tempdir"); + let cert_path = write_cert_file(&temp_dir, "empty.pem", ""); + + let output = run_probe(&[(CODEX_CA_CERT_ENV, cert_path.as_path())]); + + assert!(!output.status.success()); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!(stderr.contains("no certificates found in PEM file")); + assert!(stderr.contains("CODEX_CA_CERTIFICATE")); + assert!(stderr.contains("SSL_CERT_FILE")); +} + +#[test] +fn rejects_malformed_pem_with_hint() { + let temp_dir = TempDir::new().expect("tempdir"); + let cert_path = write_cert_file( + &temp_dir, + "malformed.pem", + "-----BEGIN CERTIFICATE-----\nMIIBroken", + ); + + let output = run_probe(&[(CODEX_CA_CERT_ENV, cert_path.as_path())]); + + assert!(!output.status.success()); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!(stderr.contains("failed to parse PEM file")); + assert!(stderr.contains("CODEX_CA_CERTIFICATE")); + assert!(stderr.contains("SSL_CERT_FILE")); +} + +#[test] +fn accepts_openssl_trusted_certificate() { + let temp_dir = TempDir::new().expect("tempdir"); + let cert_path = write_cert_file(&temp_dir, "trusted.pem", TRUSTED_TEST_CERT); + + let output = run_probe(&[(CODEX_CA_CERT_ENV, cert_path.as_path())]); + + assert!(output.status.success()); +} + +#[test] +fn accepts_bundle_with_crl() { + let temp_dir = TempDir::new().expect("tempdir"); + let crl = "-----BEGIN X509 CRL-----\nMIIC\n-----END X509 CRL-----"; + let bundle = format!("{TEST_CERT_1}\n{crl}"); + let cert_path = write_cert_file(&temp_dir, "bundle_crl.pem", &bundle); + + let output = run_probe(&[(CODEX_CA_CERT_ENV, cert_path.as_path())]); + + assert!(output.status.success()); +} diff --git a/codex-rs/http-client/tests/fixtures/test-ca-trusted.pem b/codex-rs/http-client/tests/fixtures/test-ca-trusted.pem new file mode 100644 index 0000000000000000000000000000000000000000..0b394ce84fea4d169172ba669f322d521f8a9773 --- /dev/null +++ b/codex-rs/http-client/tests/fixtures/test-ca-trusted.pem @@ -0,0 +1,25 @@ +# Test-only OpenSSL trusted-certificate fixture generated from test-ca.pem with +# `openssl x509 -addtrust serverAuth -trustout`. +# The extra trailing bytes model the OpenSSL X509_AUX data that follows the +# certificate DER in real TRUSTED CERTIFICATE bundles. +# This fixture exists to validate the X509_AUX trimming path against a real +# OpenSSL-generated artifact, not just label normalization. +-----BEGIN TRUSTED CERTIFICATE----- +MIIDBTCCAe2gAwIBAgIUZYhGvBUG7SucNzYh9VIeZ7b9zHowDQYJKoZIhvcNAQEL +BQAwEjEQMA4GA1UEAwwHdGVzdC1jYTAeFw0yNTEyMTEyMzEyNTFaFw0zNTEyMDky +MzEyNTFaMBIxEDAOBgNVBAMMB3Rlc3QtY2EwggEiMA0GCSqGSIb3DQEBAQUAA4IB +DwAwggEKAoIBAQC+NJRZAdn15FFBN8eR1HTAe+LMVpO19kKtiCsQjyqHONfhfHcF +7zQfwmH6MqeNpC/5k5m8V1uSIhyHBskQm83Jv8/vHlffNxE/hl0Na/Yd1bc+2kxH +twIAsF32GKnSKnFva/iGczV81+/ETgG6RXfTfy/Xs6fXL8On8SRRmTcMw0bEfwko +ziid87VOHg2JfdRKN5QpS9lvQ8q4q2M3jMftolpUTpwlR0u8j9OXnZfn+ja33X0l +kjkoCbXE2fVbAzO/jhUHQX1H5RbTGGUnrrCWAj84Rq/E80KK1nrRF91K+vgZmilM +gOZosLMMI1PeqTakwg1yIRngpTyk0eJP+haxAgMBAAGjUzBRMB0GA1UdDgQWBBT6 +sqvfjMIl0DFZkeu8LU577YqMVDAfBgNVHSMEGDAWgBT6sqvfjMIl0DFZkeu8LU57 +7YqMVDAPBgNVHRMBAf8EBTADAQH/MA0GCSqGSIb3DQEBCwUAA4IBAQBQ1sYs2RvB +TZ+xSBglLwH/S7zXVJIDwQ23Rlj11dgnVvcilSJCX24Rr+pfIVLpYNDdZzc/DIJd +S1dt2JuLnvXnle29rU7cxuzYUkUkRtaeY2Sj210vsE3lqUFyIy8XCc/lteb+FiJ7 +zo/gPk7P+y4ihK9Mm6SBqkDVEYSFSn9bgoemK+0e93jGe2182PyuTwfTmZgENSBO +2f9dSuay4C7e5UO8bhVccQJg6f4d70zUNG0oPHrnVxJLjwCd++jx25Gh4U7+ek13 +CW57pxJrpPMDWb2YK64rT2juHMKF73YuplW92SInd+QLpI2ekTLc+bRw8JvqzXg+ +SprtRUBjlWzjMAwwCgYIKwYBBQUHAwE= +-----END TRUSTED CERTIFICATE----- diff --git a/codex-rs/http-client/tests/fixtures/test-ca.pem b/codex-rs/http-client/tests/fixtures/test-ca.pem new file mode 100644 index 0000000000000000000000000000000000000000..7c9a9883c81bcd16fc27fa9df9aad3a65d8396fa --- /dev/null +++ b/codex-rs/http-client/tests/fixtures/test-ca.pem @@ -0,0 +1,21 @@ +# Test-only self-signed CA fixture used for single-certificate loading. +# These tests only verify PEM parsing and root-certificate registration, not a TLS handshake. +-----BEGIN CERTIFICATE----- +MIIDBTCCAe2gAwIBAgIUZYhGvBUG7SucNzYh9VIeZ7b9zHowDQYJKoZIhvcNAQEL +BQAwEjEQMA4GA1UEAwwHdGVzdC1jYTAeFw0yNTEyMTEyMzEyNTFaFw0zNTEyMDky +MzEyNTFaMBIxEDAOBgNVBAMMB3Rlc3QtY2EwggEiMA0GCSqGSIb3DQEBAQUAA4IB +DwAwggEKAoIBAQC+NJRZAdn15FFBN8eR1HTAe+LMVpO19kKtiCsQjyqHONfhfHcF +7zQfwmH6MqeNpC/5k5m8V1uSIhyHBskQm83Jv8/vHlffNxE/hl0Na/Yd1bc+2kxH +twIAsF32GKnSKnFva/iGczV81+/ETgG6RXfTfy/Xs6fXL8On8SRRmTcMw0bEfwko +ziid87VOHg2JfdRKN5QpS9lvQ8q4q2M3jMftolpUTpwlR0u8j9OXnZfn+ja33X0l +kjkoCbXE2fVbAzO/jhUHQX1H5RbTGGUnrrCWAj84Rq/E80KK1nrRF91K+vgZmilM +gOZosLMMI1PeqTakwg1yIRngpTyk0eJP+haxAgMBAAGjUzBRMB0GA1UdDgQWBBT6 +sqvfjMIl0DFZkeu8LU577YqMVDAfBgNVHSMEGDAWgBT6sqvfjMIl0DFZkeu8LU57 +7YqMVDAPBgNVHRMBAf8EBTADAQH/MA0GCSqGSIb3DQEBCwUAA4IBAQBQ1sYs2RvB +TZ+xSBglLwH/S7zXVJIDwQ23Rlj11dgnVvcilSJCX24Rr+pfIVLpYNDdZzc/DIJd +S1dt2JuLnvXnle29rU7cxuzYUkUkRtaeY2Sj210vsE3lqUFyIy8XCc/lteb+FiJ7 +zo/gPk7P+y4ihK9Mm6SBqkDVEYSFSn9bgoemK+0e93jGe2182PyuTwfTmZgENSBO +2f9dSuay4C7e5UO8bhVccQJg6f4d70zUNG0oPHrnVxJLjwCd++jx25Gh4U7+ek13 +CW57pxJrpPMDWb2YK64rT2juHMKF73YuplW92SInd+QLpI2ekTLc+bRw8JvqzXg+ +SprtRUBjlWzj +-----END CERTIFICATE----- diff --git a/codex-rs/http-client/tests/fixtures/test-intermediate.pem b/codex-rs/http-client/tests/fixtures/test-intermediate.pem new file mode 100644 index 0000000000000000000000000000000000000000..f29e69d63fd584a4dd4df8ad9100d79ae8108a07 --- /dev/null +++ b/codex-rs/http-client/tests/fixtures/test-intermediate.pem @@ -0,0 +1,21 @@ +# Second valid test-only certificate used for multi-certificate bundle coverage. +# It is intentionally distinct from test-ca.pem; chain validation is not part of these tests. +-----BEGIN CERTIFICATE----- +MIIDGTCCAgGgAwIBAgIUWxlcvHzwITWAHWHbKMFUTgeDmjwwDQYJKoZIhvcNAQEL +BQAwHDEaMBgGA1UEAwwRdGVzdC1pbnRlcm1lZGlhdGUwHhcNMjUxMTE5MTU1MDIz +WhcNMjYxMTE5MTU1MDIzWjAcMRowGAYDVQQDDBF0ZXN0LWludGVybWVkaWF0ZTCC +ASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBANq7xbeYpC2GaXANqD1nLk0t +j9j2sOk6e7DqTapxnIUijS7z4DF0Vo1xHM07wK1m+wsB/t9CubNYRvtn6hrIzx7K +jjlmvxo4/YluwO1EDMQWZAXkaY2O28ESKVx7QLfBPYAc4bf/5B4Nmt6KX5sQyyyH +2qTfzVBUCAl3sI+Ydd3mx7NOye1yNNkCNqyK3Hj45F1JuH8NZxcb4OlKssZhMlD+ +EQx4G46AzKE9Ho8AqlQvg/tiWrMHRluw7zolMJ/AXzedAXedNIrX4fCOmZwcTkA1 +a8eLPP8oM9VFrr67a7on6p4zPqugUEQ4fawp7A5KqSjUAVCt1FXmn2V8N8V6W/sC +AwEAAaNTMFEwHQYDVR0OBBYEFBEwRwW0gm3IjhLw1U3eOAvR0r6SMB8GA1UdIwQY +MBaAFBEwRwW0gm3IjhLw1U3eOAvR0r6SMA8GA1UdEwEB/wQFMAMBAf8wDQYJKoZI +hvcNAQELBQADggEBAB2fjAlpevK42Odv8XUEgV6VWlEP9HAmkRvugW9hjhzx1Iz9 +Vh/l9VcxL7PcqdpyGH+BIRvQIMokcYF5TXzf/KV1T2y56U8AWaSd2/xSjYNWwkgE +TLE5V+H/YDKzvTe58UrOaxa5N3URscQL9f+ZKworODmfMlkJ1mlREK130ZMlBexB +p9w5wo1M1fjx76Rqzq9MkpwBSbIO2zx/8+qy4BAH23MPGW+9OOnnq2DiIX3qUu1v +hnjYOxYpCB28MZEJmqsjFJQQ9RF+Te4U2/oknVcf8lZIMJ2ZBOwt2zg8RqCtM52/ +IbATwYj77wg3CFLFKcDYs3tdUqpiniabKcf6zAs= +-----END CERTIFICATE----- diff --git a/codex-rs/mermaid/examples/render.rs b/codex-rs/mermaid/examples/render.rs new file mode 100644 index 0000000000000000000000000000000000000000..2a19e74f0b6d1f5f7c8c057c922bed862d10dfbc --- /dev/null +++ b/codex-rs/mermaid/examples/render.rs @@ -0,0 +1,17 @@ +//! Render a diagram from standard input for manual prototype evaluation. + +use std::io::Read; + +fn main() -> Result<(), Box> { + let mut source = String::new(); + std::io::stdin() + .take(16 * 1024 + 1) + .read_to_string(&mut source)?; + let max_width = std::env::args() + .nth(1) + .map(|value| value.parse()) + .transpose()? + .unwrap_or(120); + println!("{}", codex_mermaid::render(&source, max_width)?); + Ok(()) +} diff --git a/codex-rs/mermaid/src/draw.rs b/codex-rs/mermaid/src/draw.rs new file mode 100644 index 0000000000000000000000000000000000000000..0068b047a7776efc2397533cddd1dbfe6dba78b7 --- /dev/null +++ b/codex-rs/mermaid/src/draw.rs @@ -0,0 +1,253 @@ +//! Orthogonal routes with unique lanes and endpoint positions, in all four directions. +//! +//! Routes never share segments. Crossings are explicitly marked; labels occupy reserved gutters. + +use super::Direction; +use super::Graph; +use super::RenderError; +use super::Role; +use super::Span; +use super::output::Cell; +use super::output::finish; +use unicode_width::UnicodeWidthChar; +use unicode_width::UnicodeWidthStr; + +pub(super) fn render(graph: &Graph, max_width: usize) -> Result>, RenderError> { + let horizontal = matches!(graph.direction, Direction::Right | Direction::Left); + let labels = graph + .nodes + .iter() + .map(|node| { + let mut lines = vec![if node.decision { + format!("◇ {}", node.label) + } else { + node.label.clone() + }]; + if !node.members.is_empty() { + lines.push("─".repeat(node.label.width())); + lines.extend(node.members.iter().cloned()); + } + lines + }) + .collect::>(); + let label_width = graph + .edges + .iter() + .flat_map(|edge| [&edge.label, &edge.target_label]) + .map(|label| label.width()) + .max() + .unwrap_or(0); + let box_cross = if horizontal { + labels.iter().map(Vec::len).max().unwrap_or(0) + 2 + } else { + labels + .iter() + .flatten() + .map(|label| label.width()) + .max() + .unwrap_or(0) + + 4 + }; + let mut counts = vec![0; graph.nodes.len()]; + let mut ports = Vec::new(); + for edge in &graph.edges { + let source = counts[edge.from]; + counts[edge.from] += 1; + let target = counts[edge.to]; + counts[edge.to] += 1; + ports.push((source, target)); + } + let stride = if horizontal { label_width + 2 } else { 1 }; + let sizes = labels + .iter() + .enumerate() + .map(|(i, lines)| { + if horizontal { + (lines.iter().map(|line| line.width()).max().unwrap_or(0) + 4) + .max(counts[i] * stride + 2) + } else { + lines.len() + counts[i] + 2 + } + }) + .collect::>(); + let mut starts = vec![0; graph.nodes.len()]; + let mut along = 0; + for i in 0..graph.nodes.len() { + let i = if matches!(graph.direction, Direction::Up | Direction::Left) { + graph.nodes.len() - 1 - i + } else { + i + }; + starts[i] = along; + along += sizes[i] + 2; + } + let first_lane = box_cross + if horizontal { 4 } else { label_width + 5 }; + let across = if graph.edges.is_empty() { + box_cross + } else { + first_lane + graph.edges.len() * 2 - 1 + }; + let (width, height) = if horizontal { + (along - 2, across) + } else { + (across, along - 2) + }; + if width > max_width { + return Err(RenderError::TooWide); + } + if width * height > super::MAX_CELLS { + return Err(RenderError::Limit); + } + let mut canvas = Canvas { + cells: vec![vec![Cell::edge(' '); width]; height], + horizontal, + }; + for (i, lines) in labels.iter().enumerate() { + let start = starts[i]; + let end = start + sizes[i] - 1; + canvas.set(/*across*/ 0, start, Cell::node('┌')); + canvas.set(box_cross - 1, start, Cell::node('┐')); + canvas.set(/*across*/ 0, end, Cell::node('└')); + canvas.set(box_cross - 1, end, Cell::node('┘')); + for x in 1..box_cross - 1 { + canvas.set(x, start, Cell::node('─')); + canvas.set(x, end, Cell::node('─')); + } + for y in start + 1..end { + canvas.set(/*across*/ 0, y, Cell::node('│')); + canvas.set(box_cross - 1, y, Cell::node('│')); + } + for (j, line) in lines.iter().enumerate() { + let (x, y) = if horizontal { + (start + 2, j + 1) + } else { + (2, start + j + 1) + }; + put_text(&mut canvas.cells[y], x, line)?; + } + } + let endpoints = graph + .edges + .iter() + .enumerate() + .map(|(i, edge)| { + let offset = |node: usize, port: usize| { + starts[node] + + if horizontal { + 1 + port * stride + } else { + labels[node].len() + 1 + port + } + }; + (offset(edge.from, ports[i].0), offset(edge.to, ports[i].1)) + }) + .collect::>(); + // Paint lanes first so every crossing is independent of iteration order. + for (i, edge) in graph.edges.iter().enumerate() { + let (source, target) = endpoints[i]; + for y in source.min(target) + 1..source.max(target) { + canvas.set( + first_lane + 2 * i, + y, + Cell::edge(if edge.dashed { '┆' } else { '│' }), + ); + } + } + for (i, edge) in graph.edges.iter().enumerate() { + let (source, target) = endpoints[i]; + let lane = first_lane + 2 * i; + for y in [source, target] { + for x in box_cross..lane { + let ch = if matches!(canvas.get(x, y), '│' | '┆') { + '╪' + } else if edge.dashed { + '┄' + } else { + '─' + }; + canvas.set(x, y, Cell::edge(ch)); + } + } + canvas.set(box_cross - 1, source, Cell::edge('├')); + canvas.set(box_cross - 1, target, Cell::edge('├')); + canvas.set(box_cross, source, Cell::edge(edge.source_tip)); + canvas.set(box_cross, target, Cell::edge(edge.target_tip)); + canvas.set( + lane, + source, + Cell::edge(if source < target { '┐' } else { '┘' }), + ); + canvas.set( + lane, + target, + Cell::edge(if source < target { '┘' } else { '┐' }), + ); + for (port, label) in [(source, &edge.label), (target, &edge.target_label)] { + let (x, y) = if horizontal { + (port + 1, box_cross + 1) + } else { + (box_cross + 2, port) + }; + put_text(&mut canvas.cells[y], x, label)?; + } + } + Ok(finish(canvas.cells)) +} + +struct Canvas { + cells: Vec>, + horizontal: bool, +} + +impl Canvas { + fn set(&mut self, across: usize, along: usize, mut cell: Cell) { + if self.horizontal { + cell.symbol = transpose(cell.symbol); + self.cells[across][along] = cell; + } else { + self.cells[along][across] = cell; + } + } + + fn get(&self, across: usize, along: usize) -> char { + if self.horizontal { + transpose(self.cells[across][along].symbol) + } else { + self.cells[along][across].symbol + } + } +} + +fn transpose(ch: char) -> char { + match ch { + '─' => '│', + '│' => '─', + '┄' => '┆', + '┆' => '┄', + '┐' => '└', + '└' => '┐', + '├' => '┬', + '┬' => '├', + '◄' => '▲', + '▲' => '◄', + '◁' => '△', + '△' => '◁', + other => other, + } +} + +pub(super) fn put_text(row: &mut [Cell], mut column: usize, text: &str) -> Result<(), RenderError> { + for ch in text.chars() { + let width = UnicodeWidthChar::width(ch).ok_or(RenderError::Unsupported)?; + row[column] = Cell { + symbol: ch, + role: Role::Text, + }; + row[column + 1..column + width].fill(Cell { + symbol: '\0', + role: Role::Text, + }); + column += width; + } + Ok(()) +} diff --git a/codex-rs/mermaid/src/families_tests.rs b/codex-rs/mermaid/src/families_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..176c031e3a80bfe6ee1f2ddc1edd6a12905f83b9 --- /dev/null +++ b/codex-rs/mermaid/src/families_tests.rs @@ -0,0 +1,359 @@ +//! End-to-end grammar, topology, and rendered output checks for each supported family. + +use super::RenderError; +use super::render; +use pretty_assertions::assert_eq; +use unicode_width::UnicodeWidthStr; + +const SEQUENCE: &str = "sequenceDiagram + actor U as Buyer + participant A as API + participant S as 库存 + participant P as Payments + U->>A: Place order + A->>S: Reserve items + S-->>A: Reservation + opt Items reserved + loop Up to 3 attempts + A->>P: Charge card + P->>P: Check fraud + P-->>A: Payment status + alt Approved + Note over A,P: Payment recorded + A-->>U: Order confirmed + else Declined + A->>S: Release items + A-->>U: Payment failed + end + end + end"; + +const STATE: &str = "stateDiagram-v2 + state \"Payment pending\" as Charging + [*] --> Draft + Draft --> Validating: submit + Validating --> Charging: valid + Validating --> Rejected: invalid + Charging --> Packing: paid + Charging --> Rejected: declined + Packing --> Shipped: dispatch + Shipped --> Delivered: received + Delivered --> [*] + Rejected --> Draft: revise + Charging: Retry up to 3 times"; + +const CLASS: &str = "classDiagram + class Order { + +String id + +Status status + +submit() + +cancel() + } + class LineItem { + +int quantity + +Decimal price + +subtotal() + } + class Payment { + +Decimal amount + +authorize() + } + class CardPayment { + +String lastFour + +authorize() + } + Order \"1\" *-- \"1..*\" LineItem : contains + Order \"1\" --> \"1\" Payment : pays with + Payment <|-- CardPayment + CardPayment ..> Order : updates"; + +const ER: &str = "erDiagram + CUSTOMER ||--o{ ORDER : places + ORDER ||--|{ LINE_ITEM : contains + PRODUCT ||..o{ LINE_ITEM : appears_in + CUSTOMER { + int id PK + string email UK + string name + } + ORDER { + int id PK + int customer_id FK + string status + } + LINE_ITEM { + int order_id PK, FK \"order key\" + int product_id PK, FK + int quantity + } + PRODUCT { + int id PK + string name + decimal price + }"; + +#[test] +fn complex_families() { + for (name, source) in [ + ("sequence", SEQUENCE), + ("state", STATE), + ("class", CLASS), + ("er", ER), + ] { + let output = render(source, /*max_width*/ 180).unwrap(); + let width = output.lines().map(UnicodeWidthStr::width).max().unwrap(); + assert_eq!(render(source, width), Ok(output.clone())); + assert_eq!(render(source, width - 1), Err(RenderError::TooWide)); + insta::assert_snapshot!(name, output); + } +} + +#[test] +fn unicode_labels_and_later_declarations() { + for direction in ["TD", "BT", "LR", "RL"] { + let source = format!( + "%% heading\ngraph {direction}; A -->|准备| B; A[请求]; B{{Réponse?}}; B -->|retry| A; B --> C[Ship 🚀]" + ); + let output = render(&source, /*max_width*/ 160).unwrap(); + for label in ["请求", "Réponse?", "Ship 🚀", "准备", "retry"] { + assert!(output.contains(label), "{direction}: {label}"); + } + let width = output.lines().map(UnicodeWidthStr::width).max().unwrap(); + assert_eq!(render(&source, width), Ok(output.clone())); + assert_eq!(render(&source, width - 1), Err(RenderError::TooWide)); + if direction == "LR" { + insta::assert_snapshot!("LR", output); + } + } +} + +#[test] +fn class_relationship_endpoints() { + for (operator, source_tip, target_tip, dashed) in [ + ("<|--", '◁', '─', false), + ("*--", '◆', '─', false), + ("o--", '◇', '─', false), + ("-->", '─', '◄', false), + ("--", '─', '─', false), + ("..>", '─', '◄', true), + ("..|>", '─', '◁', true), + ("..", '─', '─', true), + ("<--", '◄', '─', false), + ("--*", '─', '◆', false), + ("--o", '─', '◇', false), + ("--|>", '─', '◁', false), + ] { + let output = render( + &format!("classDiagram; A \"one\" {operator} \"many\" B : uses"), + /*max_width*/ 100, + ) + .unwrap(); + let ports = output + .lines() + .filter(|line| line.contains('├')) + .collect::>(); + assert!( + ports[0].contains(&format!("├{source_tip}")), + "{operator}: {output}" + ); + assert!( + ports[1].contains(&format!("├{target_tip}")), + "{operator}: {output}" + ); + assert!(ports[0].contains("(one) uses")); + assert!(ports[1].contains("(many)")); + assert_eq!(output.contains('┆'), dashed); + } +} + +#[test] +fn er_cardinalities() { + for (left, source_card) in [ + ("||", "1"), + ("|o", "0..1"), + ("}|", "1..many"), + ("}o", "0..many"), + ] { + for (right, target_card) in [ + ("||", "1"), + ("o|", "0..1"), + ("|{", "1..many"), + ("o{", "0..many"), + ] { + let output = render( + &format!("erDiagram; A {left}--{right} B : owns"), + /*max_width*/ 100, + ) + .unwrap(); + let ports = output + .lines() + .filter(|line| line.contains('├')) + .collect::>(); + assert!(ports[0].contains(&format!("({source_card}) owns"))); + assert!(ports[1].contains(&format!("({target_card})"))); + } + } +} + +#[test] +fn rejects_incomplete_and_unsupported_families() { + for source in [ + "sequenceDiagram; A->>B: hello; nonsense", + "sequenceDiagram; A->>B: hello; activate B", + "sequenceDiagram; participant A as x; participant A as y", + "sequenceDiagram; alt ready; A->>B: hi", + "sequenceDiagram; A->>B: hi; end", + "sequenceDiagram; loop retry; else no; end", + "sequenceDiagram; alt one; else two; else three; end", + "sequenceDiagram; A->>B:
", + "sequenceDiagram; A->>B: \u{1b}[31m", + "sequenceDiagram; participant A as e\u{301}", + "stateDiagram-v2; state Processing {; A-->B; }", + "stateDiagram-v2; A --> B: ok; garbage text", + "stateDiagram-v2; state \"First\" as A; state \"Second\" as A", + "stateDiagram-v2; accDescr: Order lifecycle; [*] --> Ready", + "stateDiagram; ACCDESCR: Order lifecycle; [*] --> Ready", + "stateDiagram-v2; A:::highlight; A --> B", + "stateDiagram-v2; A --> B:::highlight", + "classDiagram; accTitle: Account model; class Account", + "classDiagram; class A {; +foo()", + "classDiagram; class A; A --? B", + "classDiagram; A --> B; click A", + "classDiagram; class A {; +; }", + "erDiagram; A ||--o{ B", + "erDiagram; A ||--?? B : owns", + "erDiagram; A {; int x KEY; }", + "erDiagram; A {; int x PK \"open comment; }", + "erDiagram; A {; int x; }; A ||--|| B : uses; trailing junk", + "erDiagram; accDescr {; A model; }; CUSTOMER", + "erDiagram; ACCDESCR {; A model; }; CUSTOMER", + "erDiagram; A ||--|| B:::highlight : owns", + "%%{init: {}}%%\nclassDiagram; class A", + ] { + assert_eq!( + render(source, /*max_width*/ 200), + Err(RenderError::Unsupported), + "{source}" + ); + } +} + +#[test] +fn metadata_names_and_style_text_remain_valid_in_content() { + for source in [ + "classDiagram; class accTitle {; +id; }; class accDescr {; +id; }", + "classDiagram; class A; A: ::: literal; A --> B: ::: literal", + "erDiagram; accTitle {; int id; }", + "erDiagram; A ||--|| B: ::: literal", + "stateDiagram-v2; A: text ::: literal; A --> B: ::: literal", + ] { + assert!(render(source, /*max_width*/ 200).is_ok(), "{source}"); + } +} + +#[test] +fn family_limits() { + for source in [ + format!( + "sequenceDiagram; {}", + (0..9) + .map(|i| format!("participant P{i};")) + .collect::() + ), + format!("sequenceDiagram; {}", "A->>B: msg;".repeat(65)), + format!( + "sequenceDiagram; {} A->>B: msg; {}", + "loop retry;".repeat(5), + "end;".repeat(5) + ), + format!("classDiagram; class A {{; {} }}", "+field;".repeat(17)), + format!("erDiagram; A {{; {} }}", "int id;".repeat(17)), + format!("stateDiagram-v2; {}", "A --> B;".repeat(25)), + ] { + assert_eq!( + render(&source, /*max_width*/ usize::MAX), + Err(RenderError::Limit), + "{source}" + ); + } +} + +#[test] +fn truncated_sources_and_terminal_widths() { + // Exercise incomplete streamed source at every UTF-8 boundary, including inside control blocks. + for source in [SEQUENCE, STATE, CLASS, ER, "sequenceDiagram; A->>B: 请求"] { + for end in source.char_indices().map(|(index, _)| index) { + if let Ok(output) = render(&source[..end], /*max_width*/ 180) { + assert!(output.lines().all(|line| line.width() <= 180)); + assert!(!output.chars().any(|ch| ch.is_control() && ch != '\n')); + } + } + } + for source in [ + "sequenceDiagram; A->>B: hi", + "sequenceDiagram; A->>A: self", + "sequenceDiagram; Note over A: memo", + ] { + let output = render(source, /*max_width*/ 100).unwrap(); + let width = output.lines().map(UnicodeWidthStr::width).max().unwrap(); + assert_eq!(render(source, width), Ok(output)); + assert_eq!(render(source, width - 1), Err(RenderError::TooWide)); + } +} + +#[test] +fn sequence_arrows_preserve_sender_recipient_and_style() { + for (operator, forward, reverse, dashed) in [ + ("->>", '▶', '◀', false), + ("-->>", '▶', '◀', true), + ("->", '─', '─', false), + ("-->", '┄', '┄', true), + ("-x", 'x', 'x', false), + ("--x", 'x', 'x', true), + ] { + for (from, to, tip) in [ + ("A", "B", forward), + ("B", "A", reverse), + ("B", "B", reverse), + ] { + let output = render( + &format!( + "sequenceDiagram; participant A; participant B; {from}{operator}{to}: msg" + ), + /*max_width*/ 100, + ) + .unwrap(); + let rows = output + .lines() + .map(|line| line.chars().collect::>()) + .collect::>(); + let a = rows[1].iter().position(|ch| *ch == 'A').unwrap(); + let b = rows[1].iter().position(|ch| *ch == 'B').unwrap(); + let recipient = if to == "A" { a } else { b }; + let row = if from == to { 5 } else { 4 }; + assert_eq!(rows[row][recipient], tip, "{from}{operator}{to}"); + assert_eq!(rows[row].contains(&'┄'), dashed); + } + } +} + +#[test] +fn canvas_limit_applies_even_with_unlimited_caller_width() { + let graph = format!( + "graph LR; {}", + format!("A-->|{}|B;", "x".repeat(40)).repeat(24) + ); + let sequence = format!( + "sequenceDiagram; {} {}", + (0..8) + .map(|i| format!("participant P{i};")) + .collect::(), + format!("P0->>P7: {};", "x".repeat(40)).repeat(64) + ); + for source in [graph, sequence] { + assert_eq!( + render(&source, /*max_width*/ usize::MAX), + Err(RenderError::Limit) + ); + } +} diff --git a/codex-rs/mermaid/src/lib.rs b/codex-rs/mermaid/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..a6475a3a35b91015fbad37d08d0b01db0576f06e --- /dev/null +++ b/codex-rs/mermaid/src/lib.rs @@ -0,0 +1,176 @@ +//! Bounded, terminal-native Mermaid diagram prototype. +//! +//! Supports flowcharts, sequences, states, classes, and entity relationships. +//! Graph edges own separate lanes and endpoint positions. Crossings never join routes. Unsupported +//! syntax and outputs exceeding the caller's width return errors, leaving source fallback to +//! the caller. This crate performs no I/O and does not depend on a Mermaid implementation. + +mod draw; +mod output; +mod parse; +mod relations; +mod sequence; +mod state; + +pub use output::Role; +pub use output::Span; +use std::fmt; + +const MAX_SOURCE: usize = 16 * 1024; +const MAX_NODES: usize = 16; +const MAX_EDGES: usize = 24; +const MAX_LABEL: usize = 40; +const MAX_CELLS: usize = 64 * 1024; + +/// A diagram cannot be faithfully represented by this bounded prototype. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RenderError { + /// Syntax or text is outside the explicitly supported subset. + Unsupported, + /// A source, diagram, label, or canvas limit was exceeded. + Limit, + /// The complete diagram exceeds the supplied display width. + TooWide, +} + +impl fmt::Display for RenderError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + Self::Unsupported => "unsupported Mermaid syntax or label", + Self::Limit => "diagram exceeds prototype limits", + Self::TooWide => "diagram exceeds available display width", + }) + } +} + +impl std::error::Error for RenderError {} + +/// Render a bounded subset of Mermaid as plain Unicode text. +/// +/// Supports flowcharts, sequence diagrams, flat state diagrams, class diagrams, and ER diagrams. +/// See the crate README for each grammar and its limits. Unknown syntax, unsafe terminal text, +/// and diagrams exceeding `max_width` return errors; no partial result is returned. +pub fn render(source: &str, max_width: usize) -> Result { + Ok(render_spans(source, max_width)? + .into_iter() + .map(|line| line.into_iter().map(|span| span.text).collect::()) + .collect::>() + .join("\n")) +} + +/// Render the same bounded diagram as lines of semantic spans for caller-provided styling. +pub fn render_spans(source: &str, max_width: usize) -> Result>, RenderError> { + if source.len() > MAX_SOURCE { + return Err(RenderError::Limit); + } + let statements = source + .lines() + .filter(|line| !line.trim_start().starts_with("%%")) + .flat_map(|line| line.split(';')) + .map(str::trim) + .filter(|line| !line.is_empty()) + .collect::>(); + if source + .lines() + .any(|line| line.trim_start().starts_with("%%{")) + { + return Err(RenderError::Unsupported); + } + let (header, body) = statements.split_first().ok_or(RenderError::Unsupported)?; + match *header { + "sequenceDiagram" => sequence::render(body, max_width), + "stateDiagram-v2" | "stateDiagram" | "classDiagram" | "erDiagram" => { + draw::render(&relations::parse(header, body)?, max_width) + } + _ => draw::render(&parse::parse(header, body)?, max_width), + } +} + +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +enum Direction { + #[default] + Down, + Up, + Right, + Left, +} + +impl Direction { + fn parse(text: &str) -> Result { + match text { + "TD" | "TB" => Ok(Self::Down), + "BT" => Ok(Self::Up), + "LR" => Ok(Self::Right), + "RL" => Ok(Self::Left), + _ => Err(RenderError::Unsupported), + } + } +} + +#[derive(Debug, PartialEq, Eq)] +struct Node { + id: String, + label: String, + decision: bool, + declared: bool, + members: Vec, +} + +#[derive(Debug, PartialEq, Eq)] +struct Edge { + from: usize, + to: usize, + label: String, + target_label: String, + source_tip: char, + target_tip: char, + dashed: bool, +} + +impl Edge { + fn directed(from: usize, to: usize, label: String) -> Self { + Self { + from, + to, + label, + target_label: String::new(), + source_tip: '─', + target_tip: '◄', + dashed: false, + } + } +} + +#[derive(Debug, Default, PartialEq, Eq)] +struct Graph { + direction: Direction, + nodes: Vec, + edges: Vec, +} + +impl Graph { + fn node(&mut self, id: &str) -> Result { + if let Some(index) = self.nodes.iter().position(|node| node.id == id) { + return Ok(index); + } + if self.nodes.len() == MAX_NODES { + return Err(RenderError::Limit); + } + self.nodes.push(Node { + id: id.to_owned(), + label: id.to_owned(), + decision: false, + declared: false, + members: Vec::new(), + }); + Ok(self.nodes.len() - 1) + } +} + +#[cfg(test)] +#[path = "tests.rs"] +mod tests; + +#[cfg(test)] +#[path = "families_tests.rs"] +mod families_tests; diff --git a/codex-rs/mermaid/src/output.rs b/codex-rs/mermaid/src/output.rs new file mode 100644 index 0000000000000000000000000000000000000000..8e78af515e22c92226b1337691343171a106f0f3 --- /dev/null +++ b/codex-rs/mermaid/src/output.rs @@ -0,0 +1,69 @@ +//! Semantic spans for theme-independent diagram output; drawing cells never contain ANSI escapes. + +/// The diagram element a caller can style with its own theme. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Role { + Node, + Edge, + Text, +} + +/// Adjacent characters with one semantic role, independent of terminal styling. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Span { + pub text: String, + pub role: Role, +} + +#[derive(Clone, Copy)] +pub(super) struct Cell { + pub symbol: char, + pub role: Role, +} + +impl Cell { + pub(super) fn edge(symbol: char) -> Self { + Self { + symbol, + role: Role::Edge, + } + } + + pub(super) fn node(symbol: char) -> Self { + Self { + symbol, + role: Role::Node, + } + } +} + +pub(super) fn finish(rows: Vec>) -> Vec> { + rows.into_iter() + .map(|row| { + let mut spans: Vec = Vec::new(); + for Cell { symbol, role } in row { + if symbol == '\0' { + continue; + } + if let Some(last) = spans.last_mut() + && last.role == role + { + last.text.push(symbol); + } else { + spans.push(Span { + text: symbol.to_string(), + role, + }); + } + } + while let Some(last) = spans.last_mut() { + last.text.truncate(last.text.trim_end().len()); + if !last.text.is_empty() { + break; + } + spans.pop(); + } + spans + }) + .collect() +} diff --git a/codex-rs/mermaid/src/parse.rs b/codex-rs/mermaid/src/parse.rs new file mode 100644 index 0000000000000000000000000000000000000000..157e69fbbf062a3b4aaeedc2f4b38f24e47ee3cc --- /dev/null +++ b/codex-rs/mermaid/src/parse.rs @@ -0,0 +1,143 @@ +//! Strict parser for a small flowchart grammar; every non-comment byte must be consumed. + +use super::Direction; +use super::Edge; +use super::Graph; +use super::MAX_EDGES; +use super::MAX_LABEL; +use super::RenderError; +use unicode_width::UnicodeWidthChar; +use unicode_width::UnicodeWidthStr; + +pub(super) fn parse(header: &str, body: &[&str]) -> Result { + let tokens = header.split_whitespace().collect::>(); + let ["flowchart" | "graph", direction] = tokens.as_slice() else { + return Err(RenderError::Unsupported); + }; + let mut graph = Graph { + direction: Direction::parse(direction)?, + ..Graph::default() + }; + for statement in body { + let mut rest = *statement; + let mut from = node(&mut rest, &mut graph)?; + while !rest.trim_start().is_empty() { + rest = rest + .trim_start() + .strip_prefix("-->") + .ok_or(RenderError::Unsupported)?; + rest = rest.trim_start(); + let label = if let Some(after) = rest.strip_prefix('|') { + let (label, remaining) = after.split_once('|').ok_or(RenderError::Unsupported)?; + check_label(label)?; + rest = remaining; + label.to_owned() + } else { + String::new() + }; + let to = node(&mut rest, &mut graph)?; + if graph.edges.len() == MAX_EDGES { + return Err(RenderError::Limit); + } + graph.edges.push(Edge::directed(from, to, label)); + from = to; + } + } + if graph.nodes.is_empty() { + return Err(RenderError::Unsupported); + } + Ok(graph) +} + +fn node(rest: &mut &str, graph: &mut Graph) -> Result { + let id = identifier(rest)?; + // Reserved constructs must not be interpreted as ordinary node declarations. + if matches!( + id, + "end" | "subgraph" | "direction" | "style" | "class" | "classDef" | "linkStyle" | "click" + ) { + return Err(RenderError::Unsupported); + } + let declaration = match rest.chars().next() { + Some(open @ ('[' | '{')) => { + let close = if open == '[' { ']' } else { '}' }; + let (label, remaining) = rest[1..] + .split_once(close) + .ok_or(RenderError::Unsupported)?; + check_label(label)?; + *rest = remaining; + Some((label, open == '{')) + } + _ => None, + }; + let index = graph.node(id)?; + if let Some((label, decision)) = declaration { + let node = &mut graph.nodes[index]; + if node.declared && (node.label != label || node.decision != decision) { + return Err(RenderError::Unsupported); + } + node.label = label.to_owned(); + node.decision = decision; + node.declared = true; + } + Ok(index) +} + +pub(super) fn check_label(label: &str) -> Result<(), RenderError> { + if label.trim().is_empty() + || label.chars().any(|ch| { + ch.is_control() + || matches!( + ch, + '[' | ']' + | '{' + | '}' + | '|' + | '<' + | '>' + | '&' + | '"' + | '\\' + | '┌' + | '┐' + | '└' + | '┘' + | '├' + | '┤' + | '╪' + | '◄' + ) + || UnicodeWidthChar::width(ch).is_none_or(|width| width == 0) + }) + { + return Err(RenderError::Unsupported); + } + // Labels are drawn one Unicode scalar at a time. Reject ligatures whose string width differs + // from those scalar widths rather than misaligning borders or underallocating the canvas. + if label + .chars() + .filter_map(UnicodeWidthChar::width) + .sum::() + != label.width() + { + return Err(RenderError::Unsupported); + } + if UnicodeWidthStr::width(label) > MAX_LABEL { + return Err(RenderError::Limit); + } + Ok(()) +} + +pub(super) fn identifier<'a>(rest: &mut &'a str) -> Result<&'a str, RenderError> { + *rest = rest.trim_start(); + let len = rest + .bytes() + .take_while(|b| b.is_ascii_alphanumeric() || *b == b'_') + .count(); + let id = &rest[..len]; + if id.is_empty() || !id.as_bytes()[0].is_ascii_alphabetic() || id.len() > MAX_LABEL { + return Err(RenderError::Unsupported); + } + *rest = &rest[len..]; + Ok(id) +} diff --git a/codex-rs/mermaid/src/relations.rs b/codex-rs/mermaid/src/relations.rs new file mode 100644 index 0000000000000000000000000000000000000000..c13f79e151a599b9914ada2227414af3158c006b --- /dev/null +++ b/codex-rs/mermaid/src/relations.rs @@ -0,0 +1,225 @@ +//! Class and ER grammars, preserving members, endpoint cardinalities, and relationship kinds. + +use super::Direction; +use super::Edge; +use super::Graph; +use super::MAX_EDGES; +use super::RenderError; +use super::parse::check_label; +use super::parse::identifier; + +pub(super) fn parse(header: &str, body: &[&str]) -> Result { + if matches!(header, "stateDiagram" | "stateDiagram-v2") { + return super::state::parse(body); + } + let er = header == "erDiagram"; + let mut graph = Graph::default(); + let mut block = None; + let mut direction = false; + for &line in body { + if let Some(index) = block { + if line == "}" { + block = None; + } else { + let member = if er { + attribute(line)? + } else { + check_label(line)?; + line.to_owned() + }; + let node: &mut super::Node = &mut graph.nodes[index]; + if node.members.len() == 16 { + return Err(RenderError::Limit); + } + node.members.push(member); + } + continue; + } + if let Some(value) = line.strip_prefix("direction ") { + if direction { + return Err(RenderError::Unsupported); + } + graph.direction = Direction::parse(value.trim())?; + direction = true; + continue; + } + let declaration = !er && line.starts_with("class "); + let mut rest = line.strip_prefix("class ").filter(|_| !er).unwrap_or(line); + let id = identifier(&mut rest)?; + rest = rest.trim_start(); + let title = id == "accTitle" || er && id.eq_ignore_ascii_case("accTitle"); + let description = id == "accDescr" || er && id.eq_ignore_ascii_case("accDescr"); + if !declaration + && (title && rest.starts_with(':') || description && rest.starts_with([':', '{'])) + { + return Err(RenderError::Unsupported); + } + let from = graph.node(id)?; + if rest == "{" && (er || declaration) { + block = Some(from); + continue; + } + if declaration || (er && rest.is_empty()) { + if !rest.is_empty() { + return Err(RenderError::Unsupported); + } + continue; + } + if !er + && let Some(member) = rest + .strip_prefix(':') + .filter(|text| !text.starts_with("::")) + { + let member = member.trim(); + check_label(member)?; + if graph.nodes[from].members.len() == 16 { + return Err(RenderError::Limit); + } + graph.nodes[from].members.push(member.to_owned()); + continue; + } + let source_card = if er { + String::new() + } else { + cardinality(&mut rest)? + }; + rest = rest.trim_start(); + let split = rest + .find("--") + .into_iter() + .chain(rest.find("..")) + .min() + .ok_or(RenderError::Unsupported)?; + if split > 2 { + return Err(RenderError::Unsupported); + } + let left = &rest[..split]; + let dashed = &rest[split..split + 2] == ".."; + rest = &rest[split + 2..]; + let (source_tip, target_tip, source_card, target_card) = if er { + let source_card = match left { + "||" => "1", + "|o" => "0..1", + "}|" => "1..many", + "}o" => "0..many", + _ => return Err(RenderError::Unsupported), + } + .to_owned(); + let right = rest.get(..2).ok_or(RenderError::Unsupported)?; + let target_card = match right { + "||" => "1", + "o|" => "0..1", + "|{" => "1..many", + "o{" => "0..many", + _ => return Err(RenderError::Unsupported), + } + .to_owned(); + rest = &rest[2..]; + ('─', '─', source_card, target_card) + } else { + let source_tip = match left { + "" => '─', + "<" => '◄', + "<|" => '◁', + "*" => '◆', + "o" => '◇', + _ => return Err(RenderError::Unsupported), + }; + let mut target_tip = '─'; + for (token, tip) in [("|>", '◁'), (">", '◄'), ("*", '◆'), ("o", '◇')] { + if let Some(after) = rest.strip_prefix(token) { + target_tip = tip; + rest = after; + break; + } + } + let target_card = cardinality(&mut rest)?; + (source_tip, target_tip, source_card, target_card) + }; + let to = graph.node(identifier(&mut rest)?)?; + rest = rest.trim(); + let label = if let Some(label) = rest + .strip_prefix(':') + .filter(|text| !text.starts_with("::")) + { + let label = label.trim(); + check_label(label)?; + label.to_owned() + } else if !rest.is_empty() || er { + return Err(RenderError::Unsupported); + } else { + String::new() + }; + if graph.edges.len() == MAX_EDGES { + return Err(RenderError::Limit); + } + let label = match (source_card.is_empty(), label.is_empty()) { + (true, _) => label, + (false, true) => format!("({source_card})"), + (false, false) => format!("({source_card}) {label}"), + }; + graph.edges.push(Edge { + from, + to, + label, + source_tip, + target_tip, + dashed, + target_label: if target_card.is_empty() { + target_card + } else { + format!("({target_card})") + }, + }); + } + if block.is_some() || graph.nodes.is_empty() { + return Err(RenderError::Unsupported); + } + Ok(graph) +} + +fn cardinality(rest: &mut &str) -> Result { + *rest = rest.trim_start(); + if let Some(after) = rest.strip_prefix('"') { + let (value, after) = after.split_once('"').ok_or(RenderError::Unsupported)?; + check_label(value)?; + *rest = after; + Ok(value.to_owned()) + } else { + Ok(String::new()) + } +} + +fn attribute(line: &str) -> Result { + let mut rest = line; + let data_type = identifier(&mut rest)?; + if !rest.starts_with(char::is_whitespace) { + return Err(RenderError::Unsupported); + } + let name = identifier(&mut rest)?; + let (keys, comment) = if let Some((keys, comment)) = rest.trim().split_once('"') { + let comment = comment.strip_suffix('"').ok_or(RenderError::Unsupported)?; + check_label(comment)?; + (keys.trim(), Some(comment)) + } else { + (rest.trim(), None) + }; + if !keys.is_empty() + && !keys + .split(',') + .all(|key| matches!(key.trim(), "PK" | "FK" | "UK")) + { + return Err(RenderError::Unsupported); + } + let mut result = format!("{data_type} {name}"); + if !keys.is_empty() { + result.push(' '); + result.push_str(keys); + } + if let Some(comment) = comment { + result.push_str(" — "); + result.push_str(comment); + } + check_label(&result)?; + Ok(result) +} diff --git a/codex-rs/mermaid/src/sequence.rs b/codex-rs/mermaid/src/sequence.rs new file mode 100644 index 0000000000000000000000000000000000000000..f075117e7a01a3c499545400a4513016907b23e4 --- /dev/null +++ b/codex-rs/mermaid/src/sequence.rs @@ -0,0 +1,306 @@ +//! Bounded sequence timelines with ordered messages and explicitly nested control fragments. + +use super::RenderError; +use super::Span; +use super::draw::put_text; +use super::output::Cell; +use super::output::finish; +use super::parse::check_label; +use super::parse::identifier; +use unicode_width::UnicodeWidthStr; + +#[derive(Debug)] +enum Event { + Message { + from: usize, + to: usize, + text: String, + dashed: bool, + arrow: char, + }, + Open(String), + Branch(String), + Close, + Note { + left: usize, + right: usize, + text: String, + }, +} + +pub(super) fn render(body: &[&str], max_width: usize) -> Result>, RenderError> { + let mut people: Vec<(String, String, bool)> = Vec::new(); + let mut events = Vec::new(); + let mut blocks = Vec::new(); + for &line in body { + let mut rest = line; + if let Some(after) = line + .strip_prefix("participant ") + .or_else(|| line.strip_prefix("actor ")) + { + rest = after; + let id = identifier(&mut rest)?; + let label = if let Some(label) = rest.trim_start().strip_prefix("as ") { + check_label(label)?; + label + } else if rest.trim().is_empty() { + id + } else { + return Err(RenderError::Unsupported); + }; + let index = participant(&mut people, id)?; + if people[index].2 { + return Err(RenderError::Unsupported); + } + people[index].1 = if line.starts_with("actor ") { + format!("{label} (actor)") + } else { + label.to_owned() + }; + people[index].2 = true; + continue; + } + let event = if line == "end" { + blocks.pop().ok_or(RenderError::Unsupported)?; + Event::Close + } else if let Some(text) = line.strip_prefix("else ") { + let Some(("alt", branched)) = blocks.last_mut() else { + return Err(RenderError::Unsupported); + }; + if *branched { + return Err(RenderError::Unsupported); + } + *branched = true; + check_label(text)?; + Event::Branch(format!("else {text}")) + } else if let Some((kind @ ("loop" | "alt" | "opt" | "critical" | "break"), text)) = + line.split_once(' ') + { + check_label(text)?; + if blocks.len() == 4 { + return Err(RenderError::Limit); + } + blocks.push((kind, false)); + Event::Open(format!("{kind} {text}")) + } else if let Some(after) = line + .strip_prefix("Note over ") + .or_else(|| line.strip_prefix("note over ")) + { + rest = after; + let left = participant(&mut people, identifier(&mut rest)?)?; + let right = if let Some(after) = rest.trim_start().strip_prefix(',') { + rest = after; + participant(&mut people, identifier(&mut rest)?)? + } else { + left + }; + let text = rest + .trim_start() + .strip_prefix(':') + .ok_or(RenderError::Unsupported)? + .trim(); + check_label(text)?; + Event::Note { + left: left.min(right), + right: left.max(right), + text: text.to_owned(), + } + } else { + let from = participant(&mut people, identifier(&mut rest)?)?; + rest = rest.trim_start(); + let mut operator = None; + for (token, dashed, arrow) in [ + ("-->>", true, '▶'), + ("->>", false, '▶'), + ("-->", true, '┄'), + ("->", false, '─'), + ("--x", true, 'x'), + ("-x", false, 'x'), + ] { + if let Some(after) = rest.strip_prefix(token) { + rest = after; + operator = Some((dashed, arrow)); + break; + } + } + let (dashed, arrow) = operator.ok_or(RenderError::Unsupported)?; + let to = participant(&mut people, identifier(&mut rest)?)?; + let text = rest + .trim_start() + .strip_prefix(':') + .ok_or(RenderError::Unsupported)? + .trim(); + check_label(text)?; + Event::Message { + from, + to, + text: text.to_owned(), + dashed, + arrow, + } + }; + if events.len() == 64 { + return Err(RenderError::Limit); + } + events.push(event); + } + if !blocks.is_empty() || people.is_empty() { + return Err(RenderError::Unsupported); + } + let name_width = people + .iter() + .map(|(_, label, _)| label.width()) + .max() + .unwrap_or(0); + let text_width = events + .iter() + .map(|event| match event { + Event::Message { text, .. } | Event::Open(text) | Event::Branch(text) => text.width(), + Event::Note { text, .. } => text.width() + 6, + Event::Close => 0, + }) + .max() + .unwrap_or(0); + let box_width = name_width + 4; + let stride = (box_width + 2).max(text_width + 4); + // Four frame levels consume eight columns on each side, separate from participant lifelines. + let centers = (0..people.len()) + .map(|i| 10 + box_width / 2 + stride * i) + .collect::>(); + let last = centers[people.len() - 1]; + let width = (last + box_width - box_width / 2 + 10).max(last + text_width + 14); + if width * (3 + 4 * events.len()) > super::MAX_CELLS { + return Err(RenderError::Limit); + } + let mut rows = vec![vec![Cell::edge(' '); width]; 3]; + for (i, (_, name, _)) in people.iter().enumerate() { + let left = centers[i] - box_width / 2; + let right = left + box_width - 1; + rows[0][left] = Cell::node('┌'); + rows[0][right] = Cell::node('┐'); + rows[2][left] = Cell::node('└'); + rows[2][right] = Cell::node('┘'); + rows[0][left + 1..right].fill(Cell::node('─')); + rows[2][left + 1..right].fill(Cell::node('─')); + rows[1][left] = Cell::node('│'); + rows[1][right] = Cell::node('│'); + put_text(&mut rows[1], left + 2, name)?; + } + let mut depth = 0; + for event in events { + let count = if matches!(event, Event::Message { from, to, .. } if from == to) { + 4 + } else { + 3 + }; + let top = rows.len(); + rows.extend((0..count).map(|_| { + let mut row = vec![Cell::edge(' '); width]; + for x in ¢ers { + row[*x] = Cell::edge('│'); + } + for level in 0..depth { + row[level * 2] = Cell::node('│'); + row[width - 1 - level * 2] = Cell::node('│'); + } + row + })); + match event { + Event::Message { + from, + to, + text, + dashed, + arrow, + } => { + let (a, b) = (centers[from], centers[to]); + let arrow = if a >= b && arrow == '▶' { + '◀' + } else { + arrow + }; + let stroke = if dashed { '┄' } else { '─' }; + put_text(&mut rows[top], a.min(b) + 2, &text)?; + if a == b { + rows[top + 1][a..a + 4].fill(Cell::edge(stroke)); + rows[top + 1][a] = Cell::edge('├'); + rows[top + 1][a + 4] = Cell::edge('┐'); + rows[top + 2][a..a + 4].fill(Cell::edge(stroke)); + rows[top + 2][a + 4] = Cell::edge('┘'); + rows[top + 2][a] = Cell::edge(arrow); + } else { + rows[top + 1][a.min(b)..=a.max(b)].fill(Cell::edge(stroke)); + for x in ¢ers { + if *x > a.min(b) && *x < a.max(b) { + rows[top + 1][*x] = Cell::edge('┼'); + } + } + rows[top + 1][a] = Cell::edge(if a < b { '├' } else { '┤' }); + rows[top + 1][b] = Cell::edge(arrow); + } + } + Event::Note { left, right, text } => { + let label = format!("Note: {text}"); + let start = centers[left]; + let end = centers[right].max(start + label.width() + 3); + rows[top][start..=end].fill(Cell::node('─')); + rows[top + 2][start..=end].fill(Cell::node('─')); + rows[top][start] = Cell::node('┌'); + rows[top][end] = Cell::node('┐'); + rows[top + 2][start] = Cell::node('└'); + rows[top + 2][end] = Cell::node('┘'); + rows[top + 1][start..=end].fill(Cell::node(' ')); + rows[top + 1][start] = Cell::node('│'); + rows[top + 1][end] = Cell::node('│'); + put_text(&mut rows[top + 1], start + 2, &label)?; + } + Event::Open(label) | Event::Branch(label) => { + let opening = !label.starts_with("else "); + if opening { + depth += 1; + } + let left = (depth - 1) * 2; + let right = width - 1 - left; + rows[top][left..=right].fill(Cell::node('─')); + rows[top][left] = Cell::node(if opening { '┌' } else { '├' }); + rows[top][right] = Cell::node(if opening { '┐' } else { '┤' }); + put_text(&mut rows[top], left + 2, &label)?; + for row in &mut rows[top + 1..top + 3] { + row[left] = Cell::node('│'); + row[right] = Cell::node('│'); + } + } + Event::Close => { + depth -= 1; + let left = depth * 2; + let right = width - 1 - left; + rows[top][left..=right].fill(Cell::node('─')); + rows[top][left] = Cell::node('└'); + rows[top][right] = Cell::node('┘'); + for row in &mut rows[top + 1..top + 3] { + row[left] = Cell::node(' '); + row[right] = Cell::node(' '); + } + } + } + } + if rows.iter().any(|row| { + row.iter() + .rposition(|cell| cell.symbol != ' ') + .is_some_and(|x| x >= max_width) + }) { + return Err(RenderError::TooWide); + } + Ok(finish(rows)) +} + +fn participant(people: &mut Vec<(String, String, bool)>, id: &str) -> Result { + if let Some(index) = people.iter().position(|(name, _, _)| name == id) { + return Ok(index); + } + if people.len() == 8 { + return Err(RenderError::Limit); + } + people.push((id.to_owned(), id.to_owned(), false)); + Ok(people.len() - 1) +} diff --git a/codex-rs/mermaid/src/snapshots/codex_mermaid__families_tests__LR.snap b/codex-rs/mermaid/src/snapshots/codex_mermaid__families_tests__LR.snap new file mode 100644 index 0000000000000000000000000000000000000000..a06e3e246afa27f9f4f60aa538b0da40fe382494 --- /dev/null +++ b/codex-rs/mermaid/src/snapshots/codex_mermaid__families_tests__LR.snap @@ -0,0 +1,17 @@ +--- +source: mermaid/src/families_tests.rs +assertion_line: 125 +expression: output +--- +┌──────────────┐ ┌─────────────────────┐ ┌─────────┐ +│ 请求 │ │ ◇ Réponse? │ │ Ship 🚀 │ +└┬──────┬──────┘ └┬──────┬──────┬──────┘ └┬────────┘ + │ ▲ ▲ │ │ ▲ + │准备 │ │ │retry │ │ + │ │ │ │ │ │ + │ │ │ │ │ │ + └──────╪──────────┘ │ │ │ + │ │ │ │ + └─────────────────┘ │ │ + │ │ + └──────────┘ diff --git a/codex-rs/mermaid/src/snapshots/codex_mermaid__families_tests__class.snap b/codex-rs/mermaid/src/snapshots/codex_mermaid__families_tests__class.snap new file mode 100644 index 0000000000000000000000000000000000000000..054ec8f9468daa38ff5b594d1d2b1d6fb6331908 --- /dev/null +++ b/codex-rs/mermaid/src/snapshots/codex_mermaid__families_tests__class.snap @@ -0,0 +1,46 @@ +--- +source: mermaid/src/families_tests.rs +assertion_line: 100 +expression: output +--- +┌──────────────────┐ +│ Order │ +│ ───── │ +│ +String id │ +│ +Status status │ +│ +submit() │ +│ +cancel() │ +│ ├◆─(1) contains────┐ +│ ├──(1) pays with───╪─┐ +│ ├◄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄╪┄╪┄┄┄┐ +└──────────────────┘ │ │ ┆ + │ │ ┆ + │ │ ┆ +┌──────────────────┐ │ │ ┆ +│ LineItem │ │ │ ┆ +│ ──────── │ │ │ ┆ +│ +int quantity │ │ │ ┆ +│ +Decimal price │ │ │ ┆ +│ +subtotal() │ │ │ ┆ +│ ├──(1..*)──────────┘ │ ┆ +└──────────────────┘ │ ┆ + │ ┆ + │ ┆ +┌──────────────────┐ │ ┆ +│ Payment │ │ ┆ +│ ─────── │ │ ┆ +│ +Decimal amount │ │ ┆ +│ +authorize() │ │ ┆ +│ ├◄─(1)───────────────┘ ┆ +│ ├◁─────────────────────┐ ┆ +└──────────────────┘ │ ┆ + │ ┆ + │ ┆ +┌──────────────────┐ │ ┆ +│ CardPayment │ │ ┆ +│ ─────────── │ │ ┆ +│ +String lastFour │ │ ┆ +│ +authorize() │ │ ┆ +│ ├──────────────────────┘ ┆ +│ ├─┄updates┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┘ +└──────────────────┘ diff --git a/codex-rs/mermaid/src/snapshots/codex_mermaid__families_tests__er.snap b/codex-rs/mermaid/src/snapshots/codex_mermaid__families_tests__er.snap new file mode 100644 index 0000000000000000000000000000000000000000..ca7699d1f0f3ab90bcb04034fd623466643ce9e5 --- /dev/null +++ b/codex-rs/mermaid/src/snapshots/codex_mermaid__families_tests__er.snap @@ -0,0 +1,45 @@ +--- +source: mermaid/src/families_tests.rs +assertion_line: 107 +expression: output +--- +┌─────────────────────────────────┐ +│ CUSTOMER │ +│ ──────── │ +│ int id PK │ +│ string email UK │ +│ string name │ +│ ├──(1) places───────┐ +└─────────────────────────────────┘ │ + │ + │ +┌─────────────────────────────────┐ │ +│ ORDER │ │ +│ ───── │ │ +│ int id PK │ │ +│ int customer_id FK │ │ +│ string status │ │ +│ ├──(0..many)────────┘ +│ ├──(1) contains───────┐ +└─────────────────────────────────┘ │ + │ + │ +┌─────────────────────────────────┐ │ +│ LINE_ITEM │ │ +│ ───────── │ │ +│ int order_id PK, FK — order key │ │ +│ int product_id PK, FK │ │ +│ int quantity │ │ +│ ├──(1..many)──────────┘ +│ ├─┄(0..many)┄┄┄┄┄┄┄┄┄┄┄┄┐ +└─────────────────────────────────┘ ┆ + ┆ + ┆ +┌─────────────────────────────────┐ ┆ +│ PRODUCT │ ┆ +│ ─────── │ ┆ +│ int id PK │ ┆ +│ string name │ ┆ +│ decimal price │ ┆ +│ ├─┄(1) appears_in┄┄┄┄┄┄┄┘ +└─────────────────────────────────┘ diff --git a/codex-rs/mermaid/src/snapshots/codex_mermaid__families_tests__sequence.snap b/codex-rs/mermaid/src/snapshots/codex_mermaid__families_tests__sequence.snap new file mode 100644 index 0000000000000000000000000000000000000000..9646e74e1312c741b2936181540a0d4343757aae --- /dev/null +++ b/codex-rs/mermaid/src/snapshots/codex_mermaid__families_tests__sequence.snap @@ -0,0 +1,60 @@ +--- +source: mermaid/src/families_tests.rs +assertion_line: 107 +expression: output +--- + ┌───────────────┐ ┌───────────────┐ ┌───────────────┐ ┌───────────────┐ + │ Buyer (actor) │ │ API │ │ 库存 │ │ Payments │ + └───────────────┘ └───────────────┘ └───────────────┘ └───────────────┘ + │ Place order │ │ │ + ├─────────────────────────▶ │ │ + │ │ │ │ + │ │ Reserve items │ │ + │ ├─────────────────────────▶ │ + │ │ │ │ + │ │ Reservation │ │ + │ ◀┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┤ │ + │ │ │ │ +┌─opt Items reserved───────────────────────────────────────────────────────────────────────────────────────────────────────────────┐ +│ │ │ │ │ │ +│ │ │ │ │ │ +│ ┌─loop Up to 3 attempts────────────────────────────────────────────────────────────────────────────────────────────────────────┐ │ +│ │ │ │ │ │ │ │ +│ │ │ │ │ │ │ │ +│ │ │ │ Charge card │ │ │ │ +│ │ │ ├─────────────────────────┼─────────────────────────▶ │ │ +│ │ │ │ │ │ │ │ +│ │ │ │ │ │ Check fraud │ │ +│ │ │ │ │ ├───┐ │ │ +│ │ │ │ │ ◀───┘ │ │ +│ │ │ │ │ │ │ │ +│ │ │ │ Payment status │ │ │ │ +│ │ │ ◀┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┼┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┤ │ │ +│ │ │ │ │ │ │ │ +│ │ ┌─alt Approved─────────────────────────────────────────────────────────────────────────────────────────────────────────────┐ │ │ +│ │ │ │ │ │ │ │ │ │ +│ │ │ │ │ │ │ │ │ │ +│ │ │ │ ┌───────────────────────────────────────────────────┐ │ │ │ +│ │ │ │ │ Note: Payment recorded │ │ │ │ +│ │ │ │ └───────────────────────────────────────────────────┘ │ │ │ +│ │ │ │ Order confirmed │ │ │ │ │ │ +│ │ │ ◀┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┤ │ │ │ │ │ +│ │ │ │ │ │ │ │ │ │ +│ │ ├─else Declined────────────────────────────────────────────────────────────────────────────────────────────────────────────┤ │ │ +│ │ │ │ │ │ │ │ │ │ +│ │ │ │ │ │ │ │ │ │ +│ │ │ │ │ Release items │ │ │ │ │ +│ │ │ │ ├─────────────────────────▶ │ │ │ │ +│ │ │ │ │ │ │ │ │ │ +│ │ │ │ Payment failed │ │ │ │ │ │ +│ │ │ ◀┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┤ │ │ │ │ │ +│ │ │ │ │ │ │ │ │ │ +│ │ └──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘ │ │ +│ │ │ │ │ │ │ │ +│ │ │ │ │ │ │ │ +│ └──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘ │ +│ │ │ │ │ │ +│ │ │ │ │ │ +└──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘ + │ │ │ │ + │ │ │ │ diff --git a/codex-rs/mermaid/src/snapshots/codex_mermaid__families_tests__state.snap b/codex-rs/mermaid/src/snapshots/codex_mermaid__families_tests__state.snap new file mode 100644 index 0000000000000000000000000000000000000000..23ccd97f5d6e83945f375f3ca12aad5463a829ce --- /dev/null +++ b/codex-rs/mermaid/src/snapshots/codex_mermaid__families_tests__state.snap @@ -0,0 +1,70 @@ +--- +source: mermaid/src/families_tests.rs +assertion_line: 100 +expression: output +--- +┌─────────────────────┐ +│ Payment pending │ +│ ─────────────── │ +│ Retry up to 3 times │ +│ ├◄────────────────┐ +│ ├──paid───────────╪───┐ +│ ├──declined───────╪───╪─┐ +└─────────────────────┘ │ │ │ + │ │ │ + │ │ │ +┌─────────────────────┐ │ │ │ +│ ● initial │ │ │ │ +│ ├─────────────┐ │ │ │ +└─────────────────────┘ │ │ │ │ + │ │ │ │ + │ │ │ │ +┌─────────────────────┐ │ │ │ │ +│ Draft │ │ │ │ │ +│ ├◄────────────┘ │ │ │ +│ ├──submit───────┐ │ │ │ +│ ├◄──────────────╪─╪───╪─╪───────┐ +└─────────────────────┘ │ │ │ │ │ + │ │ │ │ │ + │ │ │ │ │ +┌─────────────────────┐ │ │ │ │ │ +│ Validating │ │ │ │ │ │ +│ ├◄──────────────┘ │ │ │ │ +│ ├──valid──────────┘ │ │ │ +│ ├──invalid──────────┐ │ │ │ +└─────────────────────┘ │ │ │ │ + │ │ │ │ + │ │ │ │ +┌─────────────────────┐ │ │ │ │ +│ Rejected │ │ │ │ │ +│ ├◄──────────────────┘ │ │ │ +│ ├◄────────────────────╪─┘ │ +│ ├──revise─────────────╪─────────┘ +└─────────────────────┘ │ + │ + │ +┌─────────────────────┐ │ +│ Packing │ │ +│ ├◄────────────────────┘ +│ ├──dispatch───────────────┐ +└─────────────────────┘ │ + │ + │ +┌─────────────────────┐ │ +│ Shipped │ │ +│ ├◄────────────────────────┘ +│ ├──received─────────────────┐ +└─────────────────────┘ │ + │ + │ +┌─────────────────────┐ │ +│ Delivered │ │ +│ ├◄──────────────────────────┘ +│ ├─────────────────────────────┐ +└─────────────────────┘ │ + │ + │ +┌─────────────────────┐ │ +│ ◎ final │ │ +│ ├◄────────────────────────────┘ +└─────────────────────┘ diff --git a/codex-rs/mermaid/src/snapshots/codex_mermaid__tests__branches_merges_and_retry_loop.snap b/codex-rs/mermaid/src/snapshots/codex_mermaid__tests__branches_merges_and_retry_loop.snap new file mode 100644 index 0000000000000000000000000000000000000000..70ceaeb6a06eec87b5e5d2becfddc0d9d6a46d58 --- /dev/null +++ b/codex-rs/mermaid/src/snapshots/codex_mermaid__tests__branches_merges_and_retry_loop.snap @@ -0,0 +1,61 @@ +--- +source: mermaid/src/tests.rs +assertion_line: 10 +expression: "render(source, 100).unwrap()" +--- +┌───────────────┐ +│ Checkout │ +│ ├────────┐ +└───────────────┘ │ + │ + │ +┌───────────────┐ │ +│ ◇ In stock? │ │ +│ ├◄───────┘ +│ ├──yes─────┐ +│ ├──no──────╪─┐ +└───────────────┘ │ │ + │ │ + │ │ +┌───────────────┐ │ │ +│ Reserve │ │ │ +│ ├◄─────────┘ │ +│ ├────────────╪─┐ +└───────────────┘ │ │ + │ │ + │ │ +┌───────────────┐ │ │ +│ Waitlist │ │ │ +│ ├◄───────────┘ │ +│ ├──────────────╪───────┐ +└───────────────┘ │ │ + │ │ + │ │ +┌───────────────┐ │ │ +│ ◇ Paid? │ │ │ +│ ├◄─────────────┘ │ +│ ├──yes───────────┐ │ +│ ├──no────────────╪─┐ │ +│ ├◄───────────────╪─╪─┐ │ +└───────────────┘ │ │ │ │ + │ │ │ │ + │ │ │ │ +┌───────────────┐ │ │ │ │ +│ Ship │ │ │ │ │ +│ ├◄───────────────┘ │ │ │ +│ ├──────────────────╪─╪─╪─┐ +└───────────────┘ │ │ │ │ + │ │ │ │ + │ │ │ │ +┌───────────────┐ │ │ │ │ +│ Retry payment │ │ │ │ │ +│ ├◄─────────────────┘ │ │ │ +│ ├────────────────────┘ │ │ +└───────────────┘ │ │ + │ │ + │ │ +┌───────────────┐ │ │ +│ Notify buyer │ │ │ +│ ├◄─────────────────────┘ │ +│ ├◄───────────────────────┘ +└───────────────┘ diff --git a/codex-rs/mermaid/src/state.rs b/codex-rs/mermaid/src/state.rs new file mode 100644 index 0000000000000000000000000000000000000000..077fd987bee02915c6c81376e95987012153c5aa --- /dev/null +++ b/codex-rs/mermaid/src/state.rs @@ -0,0 +1,109 @@ +//! Flat state machines, with distinct initial/final pseudostates and labeled transitions. + +use super::Direction; +use super::Edge; +use super::Graph; +use super::MAX_EDGES; +use super::RenderError; +use super::parse::check_label; +use super::parse::identifier; + +pub(super) fn parse(body: &[&str]) -> Result { + let mut graph = Graph::default(); + let mut direction = false; + for &line in body { + if matches!(line, "state" | "direction" | "note" | "end" | "hide") { + return Err(RenderError::Unsupported); + } + if let Some(value) = line.strip_prefix("direction ") { + if direction { + return Err(RenderError::Unsupported); + } + graph.direction = Direction::parse(value.trim())?; + direction = true; + continue; + } + let mut rest = line; + if let Some(after) = rest.strip_prefix("state \"") { + let (label, after) = after.split_once('"').ok_or(RenderError::Unsupported)?; + check_label(label)?; + rest = after + .trim_start() + .strip_prefix("as ") + .ok_or(RenderError::Unsupported)?; + let index = graph.node(identifier(&mut rest)?)?; + if !rest.trim().is_empty() || graph.nodes[index].declared { + return Err(RenderError::Unsupported); + } + graph.nodes[index].label = label.to_owned(); + graph.nodes[index].declared = true; + continue; + } + let from = if let Some(after) = rest.strip_prefix("[*]") { + rest = after; + let index = graph.node("[initial]")?; + graph.nodes[index].label = "● initial".to_owned(); + index + } else { + let id = identifier(&mut rest)?; + if (id.eq_ignore_ascii_case("accTitle") || id.eq_ignore_ascii_case("accDescr")) + && rest.trim_start().starts_with(':') + { + return Err(RenderError::Unsupported); + } + graph.node(id)? + }; + rest = rest.trim(); + if rest.is_empty() && graph.nodes[from].id != "[initial]" { + continue; + } + if let Some(description) = rest + .strip_prefix(':') + .filter(|text| !text.starts_with("::")) + { + if graph.nodes[from].id == "[initial]" { + return Err(RenderError::Unsupported); + } + let description = description.trim(); + check_label(description)?; + if graph.nodes[from].members.len() == 16 { + return Err(RenderError::Limit); + } + graph.nodes[from].members.push(description.to_owned()); + continue; + } + rest = rest + .strip_prefix("-->") + .ok_or(RenderError::Unsupported)? + .trim_start(); + let to = if let Some(after) = rest.strip_prefix("[*]") { + rest = after; + let index = graph.node("[final]")?; + graph.nodes[index].label = "◎ final".to_owned(); + index + } else { + graph.node(identifier(&mut rest)?)? + }; + let label = if let Some(label) = rest + .trim() + .strip_prefix(':') + .filter(|text| !text.starts_with("::")) + { + let label = label.trim(); + check_label(label)?; + label.to_owned() + } else if rest.trim().is_empty() { + String::new() + } else { + return Err(RenderError::Unsupported); + }; + if graph.edges.len() == MAX_EDGES { + return Err(RenderError::Limit); + } + graph.edges.push(Edge::directed(from, to, label)); + } + if graph.nodes.is_empty() { + return Err(RenderError::Unsupported); + } + Ok(graph) +} diff --git a/codex-rs/mermaid/src/tests.rs b/codex-rs/mermaid/src/tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..7daf9935ca8b515cd6a67e720a1975104ee9c3e7 --- /dev/null +++ b/codex-rs/mermaid/src/tests.rs @@ -0,0 +1,193 @@ +use super::RenderError; +use super::render; +use insta::assert_snapshot; +use pretty_assertions::assert_eq; +use unicode_width::UnicodeWidthStr; + +#[test] +fn branches_merges_and_retry_loop() { + let source = "flowchart TD\nA[Checkout] --> B{In stock?}\nB -->|yes| C[Reserve]\nB -->|no| D[Waitlist]\nC --> E{Paid?}\nE -->|yes| F[Ship]\nE -->|no| G[Retry payment]\nG --> E\nD --> H[Notify buyer]\nF --> H"; + assert_snapshot!(render(source, /*max_width*/ 100).unwrap()); +} + +#[test] +fn rejects_partial_or_unsupported_input() { + for source in [ + "flowchart TD; subgraph X; A; end", + "flowchart TD; A --> B; garbage syntax", + "flowchart TD; A -.-> B", + "flowchart TD; A & B --> C", + "flowchart TD; A[one]; A[two]", + "flowchart TD; A[HTML]", + "flowchart TD; A[]", + "flowchart TD; A[\u{1b}]", + "flowchart TD; A[e\u{301}]", + "flowchart TD; A[👍🏽]", + "flowchart TD; A[👩\u{200d}💻]", + "flowchart TD; A[✈\u{fe0f}]", + "flowchart TD; A[zero\u{200b}width]", + "flowchart TD; A[left\u{202e}right]", + "flowchart TD; A[\u{2066}isolated\u{2069}]", + "flowchart TD; A[unclosed", + "flowchart TD; click A", + "flowchart TD; A((circle))", + "flowchart TD; A -->|unclosed B", + "flowchart TD; A[\"quoted\"]", + "flowchart TD; A[foo;bar]", + "flowchart TD; A[لا]", + "flowchart TD; A -->|yes┐| B", + ] { + assert_eq!( + render(source, /*max_width*/ 100), + Err(RenderError::Unsupported), + "{source:?}" + ); + } +} + +#[test] +fn source_graph_and_width_limits() { + for source in [ + " ".repeat(16 * 1024 + 1), + format!("graph TD; A[{}]", "x".repeat(41)), + format!( + "graph TD; {}", + (0..17).map(|n| format!("N{n};")).collect::() + ), + format!("graph TD; {}", "A-->B;".repeat(25)), + ] { + assert_eq!(render(&source, /*max_width*/ 200), Err(RenderError::Limit)); + } + let output = render("graph TD; A --> B", /*max_width*/ 100).unwrap(); + let width = output.lines().map(UnicodeWidthStr::width).max().unwrap(); + assert_eq!(render("graph TD; A --> B", width), Ok(output)); + assert_eq!( + render("graph TD; A --> B", width - 1), + Err(RenderError::TooWide) + ); + assert_eq!( + render("graph TD; A", /*max_width*/ 0), + Err(RenderError::TooWide) + ); +} + +#[test] +fn reconstruct_every_edge_from_rendered_paths() { + // All 512 directed graphs on three nodes, including self-loops, cycles, fan-in and fan-out. + // Reconstruct connections from the emitted glyphs without consulting the renderer's layout. + for (mask, direction) in + (0u16..512).flat_map(|mask| ["TD", "BT", "LR", "RL"].map(|direction| (mask, direction))) + { + let mut source = format!("graph {direction}; A; B; C;"); + let mut expected = Vec::new(); + for from in 0..3 { + for to in 0..3 { + if mask & (1 << (from * 3 + to)) != 0 { + let a = char::from(b'A' + from); + let b = char::from(b'A' + to); + source.push_str(&format!("{a}-->{b};")); + expected.push((a, b)); + } + } + } + let output = render(&source, /*max_width*/ 100).unwrap(); + let mut rows = output + .lines() + .map(|line| line.chars().collect::>()) + .collect::>(); + if matches!(direction, "LR" | "RL") { + let width = rows.iter().map(Vec::len).max().unwrap(); + rows = (0..width) + .map(|x| { + rows.iter() + .map(|row| match row.get(x).copied().unwrap_or(' ') { + '─' => '│', + '│' => '─', + '┐' => '└', + '└' => '┐', + '┬' => '├', + '▲' => '◄', + other => other, + }) + .collect() + }) + .collect(); + } + let mut order = Vec::new(); + let mut owners = vec![' '; rows.len()]; + for (top, row) in rows.iter().enumerate() { + if row.first() != Some(&'┌') { + continue; + } + let bottom = (top + 1..rows.len()) + .find(|y| rows[*y].first() == Some(&'└')) + .unwrap(); + let owner = rows[top..=bottom] + .iter() + .flatten() + .find(|ch| matches!(ch, 'A' | 'B' | 'C')) + .unwrap(); + owners[top..=bottom].fill(*owner); + order.push(*owner); + } + assert_eq!( + order, + if matches!(direction, "BT" | "RL") { + vec!['C', 'B', 'A'] + } else { + vec!['A', 'B', 'C'] + } + ); + let mut actual = Vec::new(); + for (y, row) in rows.iter().enumerate() { + if let Some(port) = row.windows(2).position(|pair| pair == ['├', '─']) { + let lane = (port + 1..row.len()) + .find(|x| matches!(row[*x], '┐' | '┘')) + .unwrap(); + let mut target = y; + loop { + target = if row[lane] == '┐' { + target + 1 + } else { + target - 1 + }; + let ch = rows[target][lane]; + if matches!(ch, '┘' | '┐') { + break; + } + assert!(matches!(ch, '│' | '╪'), "broken vertical path: {source}"); + } + assert_eq!(&rows[target][port..port + 2], &['├', '◄']); + assert!( + rows[target][port + 2..lane] + .iter() + .all(|ch| matches!(ch, '─' | '╪')) + ); + actual.push((owners[y], owners[target])); + } + } + actual.sort_unstable(); + assert_eq!(actual, expected, "{source}"); + } +} + +#[test] +fn semantic_spans_distinguish_labels_from_matching_endpoint_glyphs() { + use super::Role; + + let lines = super::render_spans("sequenceDiagram; A-xB: x", /*max_width*/ 100).unwrap(); + let roles = lines + .iter() + .flatten() + .flat_map(|span| span.text.chars().filter(|ch| *ch == 'x').map(|_| span.role)) + .collect::>(); + assert_eq!(roles, vec![Role::Text, Role::Edge]); + assert_eq!( + lines[0] + .iter() + .find(|span| span.text.contains('┌')) + .unwrap() + .role, + Role::Node + ); +} diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..9c27b15cc4800fbc8138423806e5e794122a211e --- /dev/null +++ b/codex-rs/model-provider-info/src/lib.rs @@ -0,0 +1,710 @@ +//! Registry of model providers supported by Codex. +//! +//! Providers can be defined in two places: +//! 1. Built-in defaults compiled into the binary so Codex works out-of-the-box. +//! 2. User-defined entries inside `~/.codex/config.toml` under the `model_providers` +//! key. These override or extend the defaults at runtime. + +use codex_api::Provider as ApiProvider; +use codex_api::RetryConfig as ApiRetryConfig; +use codex_protocol::auth::AuthMode; +use codex_protocol::config_types::ModelProviderAuthInfo; +use codex_protocol::error::CodexErr; +use codex_protocol::error::EnvVarError; +use codex_protocol::error::Result as CodexResult; +use codex_utils_redacted_string::RedactedString; +use http::HeaderMap; +use http::header::HeaderName; +use http::header::HeaderValue; +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; +use std::collections::HashMap; +use std::fmt; +use std::num::NonZeroU64; +use std::path::Component; +use std::path::Path; +use std::time::Duration; + +const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000; +const DEFAULT_STREAM_MAX_RETRIES: u64 = 5; +const DEFAULT_REQUEST_MAX_RETRIES: u64 = 4; +const DEFAULT_AWS_CREDENTIAL_EXPORT_TIMEOUT_MS: u64 = 30_000; +const DEFAULT_AWS_AUTH_REFRESH_TIMEOUT_MS: u64 = 300_000; +pub const DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS: u64 = 15_000; +/// Hard cap for user-configured `stream_max_retries`. +const MAX_STREAM_MAX_RETRIES: u64 = 100; +/// Hard cap for user-configured `request_max_retries`. +const MAX_REQUEST_MAX_RETRIES: u64 = 100; + +const OPENAI_PROVIDER_NAME: &str = "OpenAI"; +const OPENAI_ACTOR_AUTHORIZATION_HEADER: &str = "x-openai-actor-authorization"; +pub const OPENAI_PROVIDER_ID: &str = "openai"; +pub const CHATGPT_CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api/codex"; +const AMAZON_BEDROCK_PROVIDER_NAME: &str = "Amazon Bedrock"; +pub const AMAZON_BEDROCK_PROVIDER_ID: &str = "amazon-bedrock"; +const AMAZON_BEDROCK_RUNTIME_PROVIDER_NAME: &str = "Amazon Bedrock Runtime"; +pub const AMAZON_BEDROCK_RUNTIME_PROVIDER_ID: &str = "amazon-bedrock-runtime"; +pub const AMAZON_BEDROCK_GPT_5_5_MODEL_ID: &str = "openai.gpt-5.5"; +pub const AMAZON_BEDROCK_GPT_5_4_MODEL_ID: &str = "openai.gpt-5.4"; +pub const AMAZON_BEDROCK_GPT_5_6_SOL_MODEL_ID: &str = "openai.gpt-5.6-sol"; +pub const AMAZON_BEDROCK_GPT_6_ASTRA_MODEL_ID: &str = "openai.gpt-6-astra"; +pub const AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID: &str = "openai.gpt-5.6-terra"; +pub const AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID: &str = "openai.gpt-5.6-luna"; +pub const AMAZON_BEDROCK_RUNTIME_GLOBAL_GPT_5_6_TERRA_MODEL_ID: &str = + "global.openai.gpt-5.6-terra"; +pub const AMAZON_BEDROCK_RUNTIME_GLOBAL_GPT_5_6_LUNA_MODEL_ID: &str = "global.openai.gpt-5.6-luna"; +pub const AMAZON_BEDROCK_DEFAULT_BASE_URL: &str = + "https://bedrock-mantle.us-east-1.api.aws/openai/v1"; +const AMAZON_BEDROCK_MANTLE_CLIENT_AGENT_HEADER: &str = "x-amzn-mantle-client-agent"; +const AMAZON_BEDROCK_MANTLE_CLIENT_AGENT_VALUE: &str = "codex"; +const CHAT_WIRE_API_REMOVED_ERROR: &str = "`wire_api = \"chat\"` is no longer supported.\nHow to fix: set `wire_api = \"responses\"` in your provider config.\nMore info: https://github.com/openai/codex/discussions/7782"; +pub const LEGACY_OLLAMA_CHAT_PROVIDER_ID: &str = "ollama-chat"; +pub const OLLAMA_CHAT_PROVIDER_REMOVED_ERROR: &str = "`ollama-chat` is no longer supported.\nHow to fix: replace `ollama-chat` with `ollama` in `model_provider`, `oss_provider`, or `--local-provider`.\nMore info: https://github.com/openai/codex/discussions/7782"; + +/// Wire protocol that the provider speaks. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, JsonSchema)] +#[serde(rename_all = "lowercase")] +pub enum WireApi { + /// The Responses API exposed by OpenAI at `/v1/responses`. + #[default] + Responses, +} + +impl fmt::Display for WireApi { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let value = match self { + Self::Responses => "responses", + }; + f.write_str(value) + } +} + +impl<'de> Deserialize<'de> for WireApi { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let value = String::deserialize(deserializer)?; + match value.as_str() { + "responses" => Ok(Self::Responses), + "chat" => Err(serde::de::Error::custom(CHAT_WIRE_API_REMOVED_ERROR)), + _ => Err(serde::de::Error::unknown_variant(&value, &["responses"])), + } + } +} + +/// Serializable representation of a provider definition. +#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, JsonSchema)] +#[schemars(deny_unknown_fields)] +pub struct ModelProviderInfo { + /// Friendly display name. + #[serde(default)] + pub name: String, + /// Base URL for the provider's OpenAI-compatible API. + pub base_url: Option, + /// Environment variable that stores the user's API key for this provider. + pub env_key: Option, + + /// Optional instructions to help the user get a valid value for the + /// variable and set it. + pub env_key_instructions: Option, + /// Value to use with `Authorization: Bearer ` header. Use of this + /// config is discouraged in favor of `env_key` for security reasons, but + /// this may be necessary when using this programmatically. + pub experimental_bearer_token: Option, + /// Command-backed bearer-token configuration for this provider. + pub auth: Option, + /// AWS SigV4 auth configuration for this provider. + pub aws: Option, + /// Which wire protocol this provider expects. + #[serde(default)] + pub wire_api: WireApi, + /// Optional query parameters to append to the base URL. + pub query_params: Option>, + /// Additional HTTP headers to include in requests to this provider where + /// the (key, value) pairs are the header name and value. + pub http_headers: Option>, + /// Optional HTTP headers to include in requests to this provider where the + /// (key, value) pairs are the header name and _environment variable_ whose + /// value should be used. If the environment variable is not set, or the + /// value is empty, the header will not be included in the request. + pub env_http_headers: Option>, + /// Maximum number of times to retry a failed HTTP request to this provider. + pub request_max_retries: Option, + /// Number of times to retry reconnecting a dropped streaming response before failing. + pub stream_max_retries: Option, + /// Idle timeout (in milliseconds) to wait for activity on a streaming response before treating + /// the connection as lost. + pub stream_idle_timeout_ms: Option, + /// Maximum time (in milliseconds) to wait for a websocket connection attempt before treating + /// it as failed. + pub websocket_connect_timeout_ms: Option, + /// Does this provider require an OpenAI API Key or ChatGPT login token? If true, + /// user is presented with login screen on first run, and login preference and token/key + /// are stored in auth.json. If false (which is the default), login screen is skipped, + /// and API key (if needed) comes from the "env_key" environment variable. + #[serde(default)] + pub requires_openai_auth: bool, + /// Whether this provider supports the Responses API WebSocket transport. + #[serde(default)] + pub supports_websockets: bool, + /// Whether this provider supports the standalone web-search endpoint. + #[serde(default)] + pub supports_standalone_web_search: bool, +} + +/// AWS SigV4 auth configuration for a model provider. +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema)] +#[schemars(deny_unknown_fields)] +pub struct ModelProviderAwsAuthInfo { + /// AWS profile name to use. When unset, the AWS SDK default chain decides. + pub profile: Option, + /// AWS region to use for provider-specific endpoints. + pub region: Option, + /// Optional command whose exported credentials replace the AWS SDK credential chain. + pub credential_export: Option, + /// Optional command used to reauthenticate after a refreshable AWS auth failure. + pub auth_refresh: Option, +} + +/// Command used to export AWS signing credentials for a model provider. +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema)] +#[schemars(deny_unknown_fields)] +pub struct AwsCredentialExportConfig { + /// Executable to invoke directly, without a shell. + pub command: String, + /// Arguments passed to the credential export command. + #[serde(default)] + pub args: Vec, + /// Maximum time to wait for the credential export command to complete. + #[serde(default = "default_aws_credential_export_timeout_ms")] + pub timeout_ms: NonZeroU64, +} + +impl AwsCredentialExportConfig { + pub fn timeout(&self) -> Duration { + Duration::from_millis(self.timeout_ms.get()) + } +} + +fn default_aws_credential_export_timeout_ms() -> NonZeroU64 { + match NonZeroU64::new(DEFAULT_AWS_CREDENTIAL_EXPORT_TIMEOUT_MS) { + Some(timeout_ms) => timeout_ms, + None => panic!("AWS credential export timeout must be non-zero"), + } +} + +/// Command used to refresh AWS credentials for a model provider. +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema)] +#[schemars(deny_unknown_fields)] +pub struct AwsAuthRefreshConfig { + /// Executable to invoke directly, without a shell. + pub command: String, + /// Arguments passed to the refresh command. + #[serde(default)] + pub args: Vec, + /// Maximum time to wait for the refresh command to complete. + #[serde(default = "default_aws_auth_refresh_timeout_ms")] + pub timeout_ms: NonZeroU64, +} + +impl AwsAuthRefreshConfig { + pub fn timeout(&self) -> Duration { + Duration::from_millis(self.timeout_ms.get()) + } +} + +fn default_aws_auth_refresh_timeout_ms() -> NonZeroU64 { + match NonZeroU64::new(DEFAULT_AWS_AUTH_REFRESH_TIMEOUT_MS) { + Some(timeout_ms) => timeout_ms, + None => panic!("AWS auth refresh timeout must be non-zero"), + } +} + +impl ModelProviderInfo { + /// Checks that a configured Bedrock entry only customizes supported fields. + /// Call this on the override before merging it with the built-in provider. + pub fn validate_bedrock_override(&self) -> Result<(), String> { + let unsupported_fields = Self { + base_url: None, + auth: None, + aws: None, + http_headers: None, + ..self.clone() + }; + if unsupported_fields != Self::default() { + return Err("only supports changing \ +`base_url`, `auth`, `http_headers`, `aws.profile`, `aws.region`, `aws.credential_export`, \ +and `aws.auth_refresh`; \ +other non-default provider fields are not supported" + .to_string()); + } + Ok(()) + } + + pub fn validate(&self) -> std::result::Result<(), String> { + if let Some(aws) = self.aws.as_ref() { + if self.supports_websockets { + // TODO(celia-oai): Support AWS SigV4 signing for WebSocket + // upgrade requests before allowing AWS-authenticated providers + // to enable Responses-over-WebSocket. + return Err("provider aws cannot be combined with supports_websockets".to_string()); + } + + let mut conflicts = Vec::new(); + if self.env_key.is_some() { + conflicts.push("env_key"); + } + if self.experimental_bearer_token.is_some() { + conflicts.push("experimental_bearer_token"); + } + if self.auth.is_some() { + conflicts.push("auth"); + } + if self.requires_openai_auth { + conflicts.push("requires_openai_auth"); + } + + if !conflicts.is_empty() { + return Err(format!( + "provider aws cannot be combined with {}", + conflicts.join(", ") + )); + } + + if let Some(credential_export) = aws.credential_export.as_ref() { + if aws.profile.is_some() { + return Err( + "provider aws.credential_export cannot be combined with aws.profile" + .to_string(), + ); + } + if credential_export.command.trim().is_empty() { + return Err( + "provider aws.credential_export.command must not be empty".to_string() + ); + } + let command = Path::new(&credential_export.command); + let mut components = command.components(); + let is_bare_command = matches!( + (components.next(), components.next()), + (Some(Component::Normal(name)), None) if name == command.as_os_str() + ); + if !command.is_absolute() && !is_bare_command { + return Err( + "provider aws.credential_export.command must be an absolute path or a bare executable name" + .to_string(), + ); + } + } + + if let Some(auth_refresh) = aws.auth_refresh.as_ref() { + if auth_refresh.command.trim().is_empty() { + return Err("provider aws.auth_refresh.command must not be empty".to_string()); + } + if auth_refresh.command != "aws" { + return Err("provider aws.auth_refresh.command must be `aws`".to_string()); + } + } + } + + let Some(auth) = self.auth.as_ref() else { + return Ok(()); + }; + + if auth.command.trim().is_empty() { + return Err("provider auth.command must not be empty".to_string()); + } + + let mut conflicts = Vec::new(); + if self.env_key.is_some() { + conflicts.push("env_key"); + } + if self.experimental_bearer_token.is_some() { + conflicts.push("experimental_bearer_token"); + } + if self.requires_openai_auth { + conflicts.push("requires_openai_auth"); + } + + if conflicts.is_empty() { + Ok(()) + } else { + Err(format!( + "provider auth cannot be combined with {}", + conflicts.join(", ") + )) + } + } + + fn build_header_map(&self) -> CodexResult { + let capacity = self.http_headers.as_ref().map_or(0, HashMap::len) + + self.env_http_headers.as_ref().map_or(0, HashMap::len); + let mut headers = HeaderMap::with_capacity(capacity); + if let Some(extra) = &self.http_headers { + for (k, v) in extra { + if let (Ok(name), Ok(value)) = + (HeaderName::try_from(k), HeaderValue::try_from(v.as_str())) + { + headers.insert(name, value); + } + } + } + + if let Some(env_headers) = &self.env_http_headers { + for (header, env_var) in env_headers { + if let Ok(val) = std::env::var(env_var) + && !val.trim().is_empty() + && let (Ok(name), Ok(value)) = + (HeaderName::try_from(header), HeaderValue::try_from(val)) + { + headers.insert(name, value); + } + } + } + + Ok(headers) + } + + pub fn to_api_provider(&self, auth_mode: Option) -> CodexResult { + let default_base_url = if matches!( + auth_mode, + Some( + AuthMode::Chatgpt + | AuthMode::ChatgptAuthTokens + | AuthMode::Headers + | AuthMode::AgentIdentity + | AuthMode::PersonalAccessToken + ) + ) { + CHATGPT_CODEX_BASE_URL + } else { + "https://api.openai.com/v1" + }; + let base_url = self + .base_url + .clone() + .unwrap_or_else(|| default_base_url.to_string()); + + let headers = self.build_header_map()?; + let retry = ApiRetryConfig { + max_attempts: self.request_max_retries(), + base_delay: Duration::from_millis(200), + retry_429: false, + retry_5xx: true, + retry_transport: true, + }; + + Ok(ApiProvider { + name: self.name.clone(), + base_url, + query_params: self.query_params.clone().map(|params| { + params + .into_iter() + .map(|(name, value)| (name, value.into_inner())) + .collect() + }), + headers, + retry, + stream_idle_timeout: self.stream_idle_timeout(), + }) + } + + /// If `env_key` is Some, returns the API key for this provider if present + /// (and non-empty) in the environment. If `env_key` is required but + /// cannot be found, returns an error. + pub fn api_key(&self) -> CodexResult> { + match &self.env_key { + Some(env_key) => { + let api_key = std::env::var(env_key) + .ok() + .filter(|v| !v.trim().is_empty()) + .ok_or_else(|| { + CodexErr::EnvVar(EnvVarError { + var: env_key.clone(), + instructions: self.env_key_instructions.clone(), + }) + })?; + Ok(Some(api_key)) + } + None => Ok(None), + } + } + + /// Effective maximum number of request retries for this provider. + pub fn request_max_retries(&self) -> u64 { + self.request_max_retries + .unwrap_or(DEFAULT_REQUEST_MAX_RETRIES) + .min(MAX_REQUEST_MAX_RETRIES) + } + + /// Effective maximum number of stream reconnection attempts for this provider. + pub fn stream_max_retries(&self) -> u64 { + self.stream_max_retries + .unwrap_or(DEFAULT_STREAM_MAX_RETRIES) + .min(MAX_STREAM_MAX_RETRIES) + } + + /// Effective idle timeout for streaming responses. + pub fn stream_idle_timeout(&self) -> Duration { + self.stream_idle_timeout_ms + .map(Duration::from_millis) + .unwrap_or(Duration::from_millis(DEFAULT_STREAM_IDLE_TIMEOUT_MS)) + } + + /// Effective timeout for websocket connect attempts. + pub fn websocket_connect_timeout(&self) -> Duration { + self.websocket_connect_timeout_ms + .map(Duration::from_millis) + .unwrap_or(Duration::from_millis(DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS)) + } + + pub fn create_openai_provider(base_url: Option) -> ModelProviderInfo { + ModelProviderInfo { + name: OPENAI_PROVIDER_NAME.into(), + base_url, + env_key: None, + env_key_instructions: None, + experimental_bearer_token: None, + auth: None, + aws: None, + wire_api: WireApi::Responses, + query_params: None, + http_headers: Some( + [("version".to_string(), env!("CARGO_PKG_VERSION").into())] + .into_iter() + .collect(), + ), + env_http_headers: Some( + [ + ( + "OpenAI-Organization".to_string(), + "OPENAI_ORGANIZATION".to_string(), + ), + ("OpenAI-Project".to_string(), "OPENAI_PROJECT".to_string()), + ] + .into_iter() + .collect(), + ), + // Use global defaults for retry/timeout unless overridden in config.toml. + request_max_retries: None, + stream_max_retries: None, + stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, + requires_openai_auth: true, + supports_websockets: true, + supports_standalone_web_search: true, + } + } + + pub fn create_amazon_bedrock_provider( + aws: Option, + ) -> ModelProviderInfo { + ModelProviderInfo { + name: AMAZON_BEDROCK_PROVIDER_NAME.into(), + // The runtime provider derives the regional Mantle endpoint when + // this is unset. A configured value is therefore unambiguously an + // endpoint override. + base_url: None, + env_key: None, + env_key_instructions: None, + experimental_bearer_token: None, + auth: None, + aws: Some(aws.unwrap_or(ModelProviderAwsAuthInfo { + profile: None, + region: None, + credential_export: None, + auth_refresh: None, + })), + wire_api: WireApi::Responses, + query_params: None, + http_headers: Some(HashMap::from([( + AMAZON_BEDROCK_MANTLE_CLIENT_AGENT_HEADER.to_string(), + AMAZON_BEDROCK_MANTLE_CLIENT_AGENT_VALUE.into(), + )])), + env_http_headers: None, + request_max_retries: None, + stream_max_retries: None, + stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, + requires_openai_auth: false, + supports_websockets: false, + supports_standalone_web_search: false, + } + } + + pub fn create_amazon_bedrock_runtime_provider( + aws: Option, + ) -> ModelProviderInfo { + let mut provider = Self::create_amazon_bedrock_provider(aws); + provider.name = AMAZON_BEDROCK_RUNTIME_PROVIDER_NAME.into(); + provider.http_headers = None; + provider + } + + pub fn is_openai(&self) -> bool { + self.name == OPENAI_PROVIDER_NAME + } + + pub fn supports_codex_backend_routes(&self) -> bool { + self.is_openai() + && self.base_url.as_deref().is_none_or(|base_url| { + base_url + .trim_end_matches('/') + .ends_with("/backend-api/codex") + }) + } + + pub fn uses_openai_actor_authorization(&self) -> bool { + !self.requires_openai_auth + && self.http_headers.as_ref().is_some_and(|headers| { + headers.iter().any(|(name, value)| { + name.eq_ignore_ascii_case(OPENAI_ACTOR_AUTHORIZATION_HEADER) + && !value.trim().is_empty() + }) + }) + } + + pub fn is_amazon_bedrock(&self) -> bool { + self.name == AMAZON_BEDROCK_PROVIDER_NAME + || self.name == AMAZON_BEDROCK_RUNTIME_PROVIDER_NAME + } + + pub fn is_amazon_bedrock_runtime(&self) -> bool { + self.name == AMAZON_BEDROCK_RUNTIME_PROVIDER_NAME + } + + pub fn has_command_auth(&self) -> bool { + self.auth.is_some() + } +} + +pub const DEFAULT_LMSTUDIO_PORT: u16 = 1234; +pub const DEFAULT_OLLAMA_PORT: u16 = 11434; + +pub const LMSTUDIO_OSS_PROVIDER_ID: &str = "lmstudio"; +pub const OLLAMA_OSS_PROVIDER_ID: &str = "ollama"; + +/// Built-in default provider list. +pub fn built_in_model_providers( + openai_base_url: Option, +) -> HashMap { + use ModelProviderInfo as P; + let openai_provider = P::create_openai_provider(openai_base_url); + let amazon_bedrock_provider = P::create_amazon_bedrock_provider(/*aws*/ None); + let amazon_bedrock_runtime_provider = + P::create_amazon_bedrock_runtime_provider(/*aws*/ None); + + // We do not want to be in the business of adjucating which third-party + // providers are bundled with Codex CLI, so we only include the OpenAI and + // open source ("oss") providers by default. Users are encouraged to add to + // `model_providers` in config.toml to add their own providers. + [ + (OPENAI_PROVIDER_ID, openai_provider), + (AMAZON_BEDROCK_PROVIDER_ID, amazon_bedrock_provider), + ( + AMAZON_BEDROCK_RUNTIME_PROVIDER_ID, + amazon_bedrock_runtime_provider, + ), + ( + OLLAMA_OSS_PROVIDER_ID, + create_oss_provider(DEFAULT_OLLAMA_PORT, WireApi::Responses), + ), + ( + LMSTUDIO_OSS_PROVIDER_ID, + create_oss_provider(DEFAULT_LMSTUDIO_PORT, WireApi::Responses), + ), + ] + .into_iter() + .map(|(k, v)| (k.to_string(), v)) + .collect() +} + +/// Merge configured providers into the built-in provider catalog. +/// +/// Configured providers extend the built-in set. Built-in providers are not +/// generally overridable, but built-in Amazon Bedrock providers allow the user +/// to customize their endpoint, authentication, headers, and AWS settings. +pub fn merge_configured_model_providers( + mut model_providers: HashMap, + configured_model_providers: HashMap, +) -> Result, String> { + for (key, mut provider) in configured_model_providers { + if matches!( + key.as_str(), + AMAZON_BEDROCK_PROVIDER_ID | AMAZON_BEDROCK_RUNTIME_PROVIDER_ID + ) { + provider + .validate_bedrock_override() + .map_err(|message| format!("model_providers.{key} {message}"))?; + let base_url_override = provider.base_url.take(); + let auth_override = provider.auth.take(); + let aws_override = provider.aws.take(); + let http_headers_override = provider.http_headers.take(); + if let Some(built_in_provider) = model_providers.get_mut(&key) { + built_in_provider.base_url = base_url_override; + built_in_provider.auth = auth_override; + if let Some(aws_override) = aws_override { + built_in_provider.aws = Some(aws_override); + } + if let Some(http_headers_override) = http_headers_override { + built_in_provider + .http_headers + .get_or_insert_default() + .extend(http_headers_override); + } + } + } else { + model_providers.entry(key).or_insert(provider); + } + } + + Ok(model_providers) +} + +pub fn create_oss_provider(default_provider_port: u16, wire_api: WireApi) -> ModelProviderInfo { + // These CODEX_OSS_ environment variables are experimental: we may + // switch to reading values from config.toml instead. + let default_codex_oss_base_url = format!( + "http://localhost:{codex_oss_port}/v1", + codex_oss_port = std::env::var("CODEX_OSS_PORT") + .ok() + .filter(|value| !value.trim().is_empty()) + .and_then(|value| value.parse::().ok()) + .unwrap_or(default_provider_port) + ); + + let codex_oss_base_url = std::env::var("CODEX_OSS_BASE_URL") + .ok() + .filter(|v| !v.trim().is_empty()) + .unwrap_or(default_codex_oss_base_url); + create_oss_provider_with_base_url(&codex_oss_base_url, wire_api) +} + +pub fn create_oss_provider_with_base_url(base_url: &str, wire_api: WireApi) -> ModelProviderInfo { + ModelProviderInfo { + name: "gpt-oss".into(), + base_url: Some(base_url.into()), + env_key: None, + env_key_instructions: None, + experimental_bearer_token: None, + auth: None, + aws: None, + wire_api, + query_params: None, + http_headers: None, + env_http_headers: None, + request_max_retries: None, + stream_max_retries: None, + stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, + requires_openai_auth: false, + supports_websockets: false, + supports_standalone_web_search: false, + } +} + +#[cfg(test)] +#[path = "model_provider_info_tests.rs"] +mod tests; diff --git a/codex-rs/model-provider-info/src/model_provider_info_tests.rs b/codex-rs/model-provider-info/src/model_provider_info_tests.rs new file mode 100644 index 0000000000000000000000000000000000000000..8e13869cb2fc22c11be79c4e9a8adc7f794d921f --- /dev/null +++ b/codex-rs/model-provider-info/src/model_provider_info_tests.rs @@ -0,0 +1,764 @@ +use super::*; +use codex_utils_absolute_path::AbsolutePathBuf; +use codex_utils_absolute_path::AbsolutePathBufGuard; +use pretty_assertions::assert_eq; +use std::num::NonZeroU64; +use tempfile::tempdir; + +#[test] +fn test_deserialize_ollama_model_provider_toml() { + let azure_provider_toml = r#" +name = "Ollama" +base_url = "http://localhost:11434/v1" + "#; + let expected_provider = ModelProviderInfo { + name: "Ollama".into(), + base_url: Some("http://localhost:11434/v1".into()), + env_key: None, + env_key_instructions: None, + experimental_bearer_token: None, + auth: None, + aws: None, + wire_api: WireApi::Responses, + query_params: None, + http_headers: None, + env_http_headers: None, + request_max_retries: None, + stream_max_retries: None, + stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, + requires_openai_auth: false, + supports_websockets: false, + supports_standalone_web_search: false, + }; + + let provider: ModelProviderInfo = toml::from_str(azure_provider_toml).unwrap(); + assert_eq!(expected_provider, provider); +} + +#[test] +fn test_deserialize_azure_model_provider_toml() { + let azure_provider_toml = r#" +name = "Azure" +base_url = "https://xxxxx.openai.azure.com/openai" +env_key = "AZURE_OPENAI_API_KEY" +query_params = { api-version = "2025-04-01-preview" } + "#; + let expected_provider = ModelProviderInfo { + name: "Azure".into(), + base_url: Some("https://xxxxx.openai.azure.com/openai".into()), + env_key: Some("AZURE_OPENAI_API_KEY".into()), + env_key_instructions: None, + experimental_bearer_token: None, + auth: None, + aws: None, + wire_api: WireApi::Responses, + query_params: Some(maplit::hashmap! { + "api-version".to_string() => "2025-04-01-preview".into(), + }), + http_headers: None, + env_http_headers: None, + request_max_retries: None, + stream_max_retries: None, + stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, + requires_openai_auth: false, + supports_websockets: false, + supports_standalone_web_search: false, + }; + + let provider: ModelProviderInfo = toml::from_str(azure_provider_toml).unwrap(); + assert_eq!(expected_provider, provider); +} + +#[test] +fn test_deserialize_example_model_provider_toml() { + let azure_provider_toml = r#" +name = "Example" +base_url = "https://example.com" +env_key = "API_KEY" +http_headers = { "X-Example-Header" = "example-value" } +env_http_headers = { "X-Example-Env-Header" = "EXAMPLE_ENV_VAR" } +supports_standalone_web_search = true + "#; + let expected_provider = ModelProviderInfo { + name: "Example".into(), + base_url: Some("https://example.com".into()), + env_key: Some("API_KEY".into()), + env_key_instructions: None, + experimental_bearer_token: None, + auth: None, + aws: None, + wire_api: WireApi::Responses, + query_params: None, + http_headers: Some(maplit::hashmap! { + "X-Example-Header".to_string() => "example-value".into(), + }), + env_http_headers: Some(maplit::hashmap! { + "X-Example-Env-Header".to_string() => "EXAMPLE_ENV_VAR".to_string(), + }), + request_max_retries: None, + stream_max_retries: None, + stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, + requires_openai_auth: false, + supports_websockets: false, + supports_standalone_web_search: true, + }; + + let provider: ModelProviderInfo = toml::from_str(azure_provider_toml).unwrap(); + assert_eq!(expected_provider, provider); +} + +#[test] +fn test_deserialize_chat_wire_api_shows_helpful_error() { + let provider_toml = r#" +name = "OpenAI using Chat Completions" +base_url = "https://api.openai.com/v1" +env_key = "OPENAI_API_KEY" +wire_api = "chat" + "#; + + let err = toml::from_str::(provider_toml).unwrap_err(); + assert!(err.to_string().contains(CHAT_WIRE_API_REMOVED_ERROR)); +} + +#[test] +fn test_deserialize_websocket_connect_timeout() { + let provider_toml = r#" +name = "OpenAI" +base_url = "https://api.openai.com/v1" +websocket_connect_timeout_ms = 15000 +supports_websockets = true + "#; + + let provider: ModelProviderInfo = toml::from_str(provider_toml).unwrap(); + assert_eq!(provider.websocket_connect_timeout_ms, Some(15_000)); +} + +#[test] +fn test_personal_access_token_uses_chatgpt_codex_base_url() { + let api_provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None) + .to_api_provider(Some(AuthMode::PersonalAccessToken)) + .expect("OpenAI provider should build API provider"); + + assert_eq!(api_provider.base_url, CHATGPT_CODEX_BASE_URL); +} + +#[test] +fn test_header_auth_uses_chatgpt_codex_base_url() { + let api_provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None) + .to_api_provider(Some(AuthMode::Headers)) + .expect("OpenAI provider should build API provider"); + + assert_eq!(api_provider.base_url, CHATGPT_CODEX_BASE_URL); +} + +#[test] +fn codex_backend_routes_require_codex_base_url() { + for (base_url, expected) in [ + (None, true), + (Some(CHATGPT_CODEX_BASE_URL), true), + (Some("https://chatgpt-staging.com/backend-api/codex/"), true), + (Some("https://proxy.example.com/v1"), false), + ] { + let provider = ModelProviderInfo::create_openai_provider(base_url.map(str::to_owned)); + assert_eq!(provider.supports_codex_backend_routes(), expected); + } +} + +#[test] +fn test_uses_openai_actor_authorization() { + let mut provider = ModelProviderInfo { + http_headers: Some(maplit::hashmap! { + "X-OpenAI-Actor-Authorization".to_string() => "actor-token".into(), + }), + ..ModelProviderInfo::default() + }; + assert!(provider.uses_openai_actor_authorization()); + + provider.http_headers = None; + assert!(!provider.uses_openai_actor_authorization()); + + provider.http_headers = Some(maplit::hashmap! { + OPENAI_ACTOR_AUTHORIZATION_HEADER.to_string() => " ".into(), + }); + assert!(!provider.uses_openai_actor_authorization()); + + provider.http_headers = Some(maplit::hashmap! { + OPENAI_ACTOR_AUTHORIZATION_HEADER.to_string() => "actor-token".into(), + }); + provider.requires_openai_auth = true; + assert!(!provider.uses_openai_actor_authorization()); +} + +#[test] +fn test_deserialize_provider_auth_config_defaults() { + let base_dir = tempdir().unwrap(); + let provider_toml = r#" +name = "Corp" + +[auth] +command = "./scripts/print-token" +args = ["--format=text"] + "#; + + let provider: ModelProviderInfo = { + let _guard = AbsolutePathBufGuard::new(base_dir.path()); + toml::from_str(provider_toml).unwrap() + }; + + assert_eq!( + provider.auth, + Some(ModelProviderAuthInfo { + command: "./scripts/print-token".to_string(), + args: vec!["--format=text".into()], + timeout_ms: NonZeroU64::new(5_000).unwrap(), + refresh_interval_ms: 300_000, + cwd: AbsolutePathBuf::resolve_path_against_base(".", base_dir.path()), + }) + ); +} + +#[test] +fn test_deserialize_provider_aws_config() { + let provider_toml = r#" +name = "Amazon Bedrock" +base_url = "https://bedrock.example.com/v1" + +[aws] +region = "us-west-2" + +[aws.credential_export] +command = "aws-vault" +args = ["--profile", "codex-bedrock"] + +[aws.auth_refresh] +command = "aws" +args = ["login", "--profile", "codex-bedrock"] + "#; + + let provider: ModelProviderInfo = toml::from_str(provider_toml).unwrap(); + + assert_eq!( + provider.aws, + Some(ModelProviderAwsAuthInfo { + profile: None, + region: Some("us-west-2".to_string()), + credential_export: Some(AwsCredentialExportConfig { + command: "aws-vault".to_string(), + args: vec!["--profile".into(), "codex-bedrock".into()], + timeout_ms: NonZeroU64::new(30_000).expect("timeout should be non-zero"), + }), + auth_refresh: Some(AwsAuthRefreshConfig { + command: "aws".to_string(), + args: vec!["login".into(), "--profile".into(), "codex-bedrock".into()], + timeout_ms: NonZeroU64::new(300_000).expect("timeout should be non-zero"), + }), + }) + ); +} + +#[test] +fn test_create_amazon_bedrock_provider() { + assert_eq!( + ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), + ModelProviderInfo { + name: "Amazon Bedrock".to_string(), + base_url: None, + env_key: None, + env_key_instructions: None, + experimental_bearer_token: None, + auth: None, + aws: Some(ModelProviderAwsAuthInfo { + profile: None, + region: None, + credential_export: None, + auth_refresh: None, + }), + wire_api: WireApi::Responses, + query_params: None, + http_headers: Some(maplit::hashmap! { + AMAZON_BEDROCK_MANTLE_CLIENT_AGENT_HEADER.to_string() => + AMAZON_BEDROCK_MANTLE_CLIENT_AGENT_VALUE.into(), + }), + env_http_headers: None, + request_max_retries: None, + stream_max_retries: None, + stream_idle_timeout_ms: None, + websocket_connect_timeout_ms: None, + requires_openai_auth: false, + supports_websockets: false, + supports_standalone_web_search: false, + } + ); +} + +#[test] +fn test_create_amazon_bedrock_runtime_provider() { + let mut expected = ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None); + expected.name = "Amazon Bedrock Runtime".to_string(); + expected.http_headers = None; + + assert_eq!( + ModelProviderInfo::create_amazon_bedrock_runtime_provider(/*aws*/ None), + expected + ); +} + +#[test] +fn test_create_amazon_bedrock_runtime_provider_with_aws_configuration() { + let provider = + ModelProviderInfo::create_amazon_bedrock_runtime_provider(Some(ModelProviderAwsAuthInfo { + profile: Some("runtime-profile".to_string()), + region: Some("us-west-2".to_string()), + credential_export: None, + auth_refresh: None, + })); + + assert_eq!( + ( + provider.name.as_str(), + provider.aws, + provider.http_headers, + provider.supports_standalone_web_search, + ), + ( + "Amazon Bedrock Runtime", + Some(ModelProviderAwsAuthInfo { + profile: Some("runtime-profile".to_string()), + region: Some("us-west-2".to_string()), + credential_export: None, + auth_refresh: None, + }), + None, + false, + ) + ); +} + +fn provider_auth_for_test() -> ModelProviderAuthInfo { + ModelProviderAuthInfo { + command: "token-fetcher".to_string(), + args: vec!["fetch".into()], + timeout_ms: NonZeroU64::new(5_000).expect("timeout should be non-zero"), + refresh_interval_ms: 300_000, + cwd: std::env::current_dir() + .expect("current directory should be available") + .try_into() + .expect("current directory should be absolute"), + } +} + +#[test] +fn test_amazon_bedrock_provider_adds_mantle_client_agent_header() { + let api_provider = ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None) + .to_api_provider(/*auth_mode*/ None) + .expect("Amazon Bedrock provider should build API provider"); + + assert_eq!( + api_provider + .headers + .get(AMAZON_BEDROCK_MANTLE_CLIENT_AGENT_HEADER) + .and_then(|value| value.to_str().ok()), + Some(AMAZON_BEDROCK_MANTLE_CLIENT_AGENT_VALUE) + ); +} + +#[test] +fn test_built_in_model_providers_include_amazon_bedrock_endpoints() { + let providers = built_in_model_providers(/*openai_base_url*/ None); + + assert_eq!( + [ + AMAZON_BEDROCK_PROVIDER_ID, + AMAZON_BEDROCK_RUNTIME_PROVIDER_ID + ] + .into_iter() + .map(|provider_id| { + providers + .get(provider_id) + .map(ModelProviderInfo::is_amazon_bedrock) + }) + .collect::>(), + vec![Some(true), Some(true)] + ); +} + +#[test] +fn test_built_in_model_providers_include_amazon_bedrock_runtime() { + let providers = built_in_model_providers(/*openai_base_url*/ None); + let runtime = providers + .get(AMAZON_BEDROCK_RUNTIME_PROVIDER_ID) + .expect("Amazon Bedrock Runtime provider should be built in"); + + assert!(runtime.is_amazon_bedrock()); + assert!(runtime.is_amazon_bedrock_runtime()); + assert!( + !providers + .get(AMAZON_BEDROCK_PROVIDER_ID) + .expect("Amazon Bedrock provider should be built in") + .is_amazon_bedrock_runtime() + ); +} + +#[test] +fn test_merge_configured_model_providers_adds_custom_provider() { + let custom_provider = ModelProviderInfo { + name: "Custom".to_string(), + base_url: Some("https://example.com/v1".to_string()), + ..ModelProviderInfo::default() + }; + let configured_model_providers = + std::collections::HashMap::from([("custom".to_string(), custom_provider.clone())]); + + let mut expected = built_in_model_providers(/*openai_base_url*/ None); + expected.insert("custom".to_string(), custom_provider); + + assert_eq!( + merge_configured_model_providers( + built_in_model_providers(/*openai_base_url*/ None), + configured_model_providers, + ), + Ok(expected) + ); +} + +#[test] +fn test_merge_configured_model_providers_applies_amazon_bedrock_aws_override() { + let credential_export = AwsCredentialExportConfig { + command: "aws-vault".to_string(), + args: vec!["export".into(), "codex-bedrock".into()], + timeout_ms: NonZeroU64::new(30_000).expect("timeout should be non-zero"), + }; + let auth_refresh = AwsAuthRefreshConfig { + command: "aws".to_string(), + args: vec!["login".into(), "--profile".into(), "codex-bedrock".into()], + timeout_ms: NonZeroU64::new(10_000).expect("timeout should be non-zero"), + }; + let configured_model_providers = std::collections::HashMap::from([( + AMAZON_BEDROCK_PROVIDER_ID.to_string(), + ModelProviderInfo { + aws: Some(ModelProviderAwsAuthInfo { + profile: None, + region: Some("us-west-2".to_string()), + credential_export: Some(credential_export.clone()), + auth_refresh: Some(auth_refresh.clone()), + }), + ..ModelProviderInfo::default() + }, + )]); + + let mut expected = built_in_model_providers(/*openai_base_url*/ None); + expected + .get_mut(AMAZON_BEDROCK_PROVIDER_ID) + .expect("Amazon Bedrock provider should be built in") + .aws = Some(ModelProviderAwsAuthInfo { + profile: None, + region: Some("us-west-2".to_string()), + credential_export: Some(credential_export), + auth_refresh: Some(auth_refresh), + }); + + assert_eq!( + merge_configured_model_providers( + built_in_model_providers(/*openai_base_url*/ None), + configured_model_providers, + ), + Ok(expected) + ); +} + +#[test] +fn test_merge_configured_model_providers_applies_runtime_overrides_independently() { + let runtime_aws = ModelProviderAwsAuthInfo { + profile: Some("runtime-profile".to_string()), + region: Some("eu-west-1".to_string()), + credential_export: None, + auth_refresh: None, + }; + let configured_model_providers = std::collections::HashMap::from([( + AMAZON_BEDROCK_RUNTIME_PROVIDER_ID.to_string(), + ModelProviderInfo { + base_url: Some("https://runtime.example.com/openai/v1".to_string()), + aws: Some(runtime_aws.clone()), + ..ModelProviderInfo::default() + }, + )]); + let mut expected = built_in_model_providers(/*openai_base_url*/ None); + let expected_runtime = expected + .get_mut(AMAZON_BEDROCK_RUNTIME_PROVIDER_ID) + .expect("Amazon Bedrock Runtime provider should be built in"); + expected_runtime.base_url = Some("https://runtime.example.com/openai/v1".to_string()); + expected_runtime.aws = Some(runtime_aws); + + assert_eq!( + merge_configured_model_providers( + built_in_model_providers(/*openai_base_url*/ None), + configured_model_providers, + ), + Ok(expected) + ); +} + +#[test] +fn test_merge_configured_model_providers_applies_amazon_bedrock_transport_overrides() { + let auth = provider_auth_for_test(); + let configured_model_providers = std::collections::HashMap::from([( + AMAZON_BEDROCK_PROVIDER_ID.to_string(), + ModelProviderInfo { + base_url: Some("https://proxy.example.com/v1".to_string()), + auth: Some(auth.clone()), + aws: Some(ModelProviderAwsAuthInfo { + profile: Some("codex-bedrock".to_string()), + region: Some("us-west-2".to_string()), + credential_export: None, + auth_refresh: None, + }), + http_headers: Some(maplit::hashmap! { + "x-example-header".to_string() => "value".into(), + }), + ..ModelProviderInfo::default() + }, + )]); + + let mut expected = built_in_model_providers(/*openai_base_url*/ None); + let expected_provider = expected + .get_mut(AMAZON_BEDROCK_PROVIDER_ID) + .expect("Amazon Bedrock provider should be built in"); + expected_provider.base_url = Some("https://proxy.example.com/v1".to_string()); + expected_provider.auth = Some(auth); + expected_provider.aws = Some(ModelProviderAwsAuthInfo { + profile: Some("codex-bedrock".to_string()), + region: Some("us-west-2".to_string()), + credential_export: None, + auth_refresh: None, + }); + expected_provider + .http_headers + .get_or_insert_default() + .insert("x-example-header".to_string(), "value".into()); + + assert_eq!( + merge_configured_model_providers( + built_in_model_providers(/*openai_base_url*/ None), + configured_model_providers, + ), + Ok(expected) + ); +} + +#[test] +fn test_merge_configured_model_providers_rejects_amazon_bedrock_non_default_fields() { + let configured_model_providers = std::collections::HashMap::from([( + AMAZON_BEDROCK_PROVIDER_ID.to_string(), + ModelProviderInfo { + name: "Custom Bedrock".to_string(), + aws: Some(ModelProviderAwsAuthInfo { + profile: Some("codex-bedrock".to_string()), + region: None, + credential_export: None, + auth_refresh: None, + }), + ..ModelProviderInfo::default() + }, + )]); + + assert_eq!( + merge_configured_model_providers( + built_in_model_providers(/*openai_base_url*/ None), + configured_model_providers, + ), + Err( + "model_providers.amazon-bedrock only supports changing `base_url`, `auth`, `http_headers`, `aws.profile`, `aws.region`, `aws.credential_export`, and `aws.auth_refresh`; other non-default provider fields are not supported" + .to_string() + ) + ); +} + +#[test] +fn test_merge_configured_model_providers_allows_amazon_bedrock_default_fields() { + let configured_model_providers = std::collections::HashMap::from([( + AMAZON_BEDROCK_PROVIDER_ID.to_string(), + ModelProviderInfo { + aws: Some(ModelProviderAwsAuthInfo { + profile: None, + region: None, + credential_export: None, + auth_refresh: None, + }), + wire_api: WireApi::Responses, + ..ModelProviderInfo::default() + }, + )]); + + assert_eq!( + merge_configured_model_providers( + built_in_model_providers(/*openai_base_url*/ None), + configured_model_providers, + ), + Ok(built_in_model_providers(/*openai_base_url*/ None)) + ); +} + +#[test] +fn test_validate_provider_aws_rejects_conflicting_auth() { + let provider = ModelProviderInfo { + aws: Some(ModelProviderAwsAuthInfo { + profile: None, + region: None, + credential_export: None, + auth_refresh: None, + }), + env_key: Some("AWS_BEARER_TOKEN_BEDROCK".to_string()), + supports_websockets: false, + ..ModelProviderInfo::create_openai_provider(/*base_url*/ None) + }; + + assert_eq!( + provider.validate(), + Err("provider aws cannot be combined with env_key, requires_openai_auth".to_string()) + ); +} + +#[test] +fn test_validate_provider_aws_rejects_websockets() { + let provider = ModelProviderInfo { + aws: Some(ModelProviderAwsAuthInfo { + profile: None, + region: None, + credential_export: None, + auth_refresh: None, + }), + requires_openai_auth: false, + supports_websockets: true, + ..ModelProviderInfo::create_openai_provider(/*base_url*/ None) + }; + + assert_eq!( + provider.validate(), + Err("provider aws cannot be combined with supports_websockets".to_string()) + ); +} + +#[test] +fn test_validate_provider_aws_auth_refresh_command() { + for (command, expected) in [ + ( + " ", + Err("provider aws.auth_refresh.command must not be empty".to_string()), + ), + ( + "other-command", + Err("provider aws.auth_refresh.command must be `aws`".to_string()), + ), + ("aws", Ok(())), + ] { + let provider = + ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { + profile: None, + region: None, + credential_export: None, + auth_refresh: Some(AwsAuthRefreshConfig { + command: command.to_string(), + args: Vec::new(), + timeout_ms: NonZeroU64::new(300_000).expect("timeout should be non-zero"), + }), + })); + + assert_eq!(provider.validate(), expected); + } +} + +#[test] +fn test_validate_provider_aws_credential_export_command() { + let absolute_command = std::env::current_exe() + .expect("current executable should have a path") + .to_string_lossy() + .into_owned(); + for (command, expected) in [ + ( + String::new(), + Err("provider aws.credential_export.command must not be empty".to_string()), + ), + ( + " ".to_string(), + Err("provider aws.credential_export.command must not be empty".to_string()), + ), + ("aws-vault".to_string(), Ok(())), + (absolute_command, Ok(())), + ( + "./scripts/export-credentials".to_string(), + Err( + "provider aws.credential_export.command must be an absolute path or a bare executable name" + .to_string(), + ), + ), + ( + "scripts/export-credentials".to_string(), + Err( + "provider aws.credential_export.command must be an absolute path or a bare executable name" + .to_string(), + ), + ), + ( + "scripts/.".to_string(), + Err( + "provider aws.credential_export.command must be an absolute path or a bare executable name" + .to_string(), + ), + ), + ] { + let mut provider = + ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { + profile: None, + region: Some("us-west-2".to_string()), + credential_export: Some(AwsCredentialExportConfig { + command, + args: vec!["--secret-argument".into()], + timeout_ms: NonZeroU64::new(30_000).expect("timeout should be non-zero"), + }), + auth_refresh: Some(AwsAuthRefreshConfig { + command: "aws".to_string(), + args: Vec::new(), + timeout_ms: NonZeroU64::new(300_000).expect("timeout should be non-zero"), + }), + })); + + assert_eq!(provider.validate(), expected); + if expected.is_ok() { + provider.aws.as_mut().expect("AWS config").profile = Some("codex-bedrock".to_string()); + assert_eq!( + provider.validate(), + Err("provider aws.credential_export cannot be combined with aws.profile".to_string()), + ); + } + let export = provider + .aws + .as_ref() + .and_then(|aws| aws.credential_export.as_ref()) + .expect("credential export should be configured"); + assert!(!format!("{export:?}").contains("--secret-argument")); + } +} + +#[test] +fn test_deserialize_provider_auth_config_allows_zero_refresh_interval() { + let base_dir = tempdir().unwrap(); + let provider_toml = r#" +name = "Corp" + +[auth] +command = "./scripts/print-token" +refresh_interval_ms = 0 + "#; + + let provider: ModelProviderInfo = { + let _guard = AbsolutePathBufGuard::new(base_dir.path()); + toml::from_str(provider_toml).unwrap() + }; + + let auth = provider.auth.expect("auth config should deserialize"); + assert_eq!(auth.refresh_interval_ms, 0); + assert_eq!(auth.refresh_interval(), None); +} diff --git a/codex-rs/prompts/src/compact.rs b/codex-rs/prompts/src/compact.rs new file mode 100644 index 0000000000000000000000000000000000000000..2ddc146aec081652200c7a03aa0244c0766d09c9 --- /dev/null +++ b/codex-rs/prompts/src/compact.rs @@ -0,0 +1,2 @@ +pub const SUMMARIZATION_PROMPT: &str = include_str!("../templates/compact/prompt.md"); +pub const SUMMARY_PREFIX: &str = include_str!("../templates/compact/summary_prefix.md"); diff --git a/codex-rs/prompts/src/lib.rs b/codex-rs/prompts/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..b7b4e444e2bf5bb0ccb72df6ea98eb408977f135 --- /dev/null +++ b/codex-rs/prompts/src/lib.rs @@ -0,0 +1,20 @@ +mod compact; +mod permissions_instructions; +mod realtime; +mod review_exit; +mod review_request; + +pub use compact::SUMMARIZATION_PROMPT; +pub use compact::SUMMARY_PREFIX; +pub use permissions_instructions::ApprovalPromptContext; +pub use permissions_instructions::PermissionsInstructions; +pub use realtime::BACKEND_PROMPT; +pub use realtime::END_INSTRUCTIONS; +pub use realtime::START_INSTRUCTIONS; +pub use review_exit::render_review_exit_interrupted; +pub use review_exit::render_review_exit_success; +pub use review_request::REVIEW_PROMPT; +pub use review_request::ResolvedReviewRequest; +pub use review_request::resolve_review_request; +pub use review_request::review_prompt; +pub use review_request::user_facing_hint; diff --git a/codex-rs/prompts/src/permissions_instructions.rs b/codex-rs/prompts/src/permissions_instructions.rs new file mode 100644 index 0000000000000000000000000000000000000000..9966618f70dbe359637c5f14ee99c47abe47d79b --- /dev/null +++ b/codex-rs/prompts/src/permissions_instructions.rs @@ -0,0 +1,455 @@ +use codex_context_fragments::ContextualUserFragment; +use codex_execpolicy::Policy; +use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::config_types::SandboxMode; +use codex_protocol::models::ContentItemKind; +use codex_protocol::models::PermissionProfile; +use codex_protocol::models::format_allow_prefixes; +use codex_protocol::openai_models::ApprovalMessages; +use codex_protocol::openai_models::PermissionMessages; +use codex_protocol::permissions::FileSystemSandboxPolicy; +use codex_protocol::permissions::NetworkSandboxPolicy; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::GranularApprovalConfig; +use codex_protocol::protocol::NetworkAccess; +use codex_protocol::protocol::WritableRoot; +use codex_utils_template::Template; +use std::path::Path; +use std::sync::LazyLock; + +const APPROVAL_POLICY_NEVER: &str = + include_str!("../templates/permissions/approval_policy/never.md"); +const APPROVAL_POLICY_UNLESS_TRUSTED: &str = + include_str!("../templates/permissions/approval_policy/unless_trusted.md"); +const APPROVAL_POLICY_ON_REQUEST_RULE: &str = + include_str!("../templates/permissions/approval_policy/on_request.md"); +const APPROVAL_POLICY_ON_REQUEST_RULE_REQUEST_PERMISSION: &str = + include_str!("../templates/permissions/approval_policy/on_request_rule_request_permission.md"); +const AUTO_REVIEW_APPROVAL_SUFFIX: &str = "`approvals_reviewer` is `auto_review`: Sandbox escalations with require_escalated will be reviewed for compliance with the policy. If a rejection happens, you should proceed only with a materially safer alternative, or inform the user of the risk and send a final message to ask for approval."; +const NETWORK_ACCESS_PLACEHOLDER: &str = "{{ network_access }}"; + +const SANDBOX_MODE_DANGER_FULL_ACCESS: &str = + include_str!("../templates/permissions/sandbox_mode/danger_full_access.md"); +const SANDBOX_MODE_WORKSPACE_WRITE: &str = + include_str!("../templates/permissions/sandbox_mode/workspace_write.md"); +const SANDBOX_MODE_READ_ONLY: &str = + include_str!("../templates/permissions/sandbox_mode/read_only.md"); + +static SANDBOX_MODE_DANGER_FULL_ACCESS_TEMPLATE: LazyLock