From 63402cbc3bf74528e0151505a0686ee12a16dc9e Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 29 Apr 2026 13:09:30 +0100 Subject: [PATCH] feat: codex journal --- codex-rs/Cargo.lock | 19 + codex-rs/Cargo.toml | 2 + codex-rs/config/src/config_toml.rs | 2 +- codex-rs/core/Cargo.toml | 1 + codex-rs/core/config.schema.json | 4 +- codex-rs/core/src/agent/control_tests.rs | 24 +- codex-rs/core/src/arc_monitor.rs | 4 +- codex-rs/core/src/compact.rs | 18 +- codex-rs/core/src/compact_remote.rs | 20 +- codex-rs/core/src/config/mod.rs | 2 +- codex-rs/core/src/context_manager/history.rs | 726 ------------------ codex-rs/core/src/context_manager/mod.rs | 10 +- codex-rs/core/src/context_manager/updates.rs | 214 +++++- codex-rs/core/src/guardian/prompt.rs | 6 +- codex-rs/core/src/prompt_debug.rs | 8 +- codex-rs/core/src/realtime_context.rs | 4 +- codex-rs/core/src/session/handlers.rs | 5 +- codex-rs/core/src/session/mod.rs | 205 +++-- .../src/session/rollout_reconstruction.rs | 25 +- codex-rs/core/src/session/tests.rs | 61 +- codex-rs/core/src/session/turn.rs | 13 +- codex-rs/core/src/state/history.rs | 314 ++++++++ .../history_tests.rs | 125 ++- codex-rs/core/src/state/mod.rs | 2 + codex-rs/core/src/state/session.rs | 71 +- codex-rs/journal/BUILD.bazel | 6 + codex-rs/journal/Cargo.toml | 23 + codex-rs/journal/src/error.rs | 24 + codex-rs/journal/src/history/classify.rs | 133 ++++ codex-rs/journal/src/history/estimate.rs | 174 +++++ codex-rs/journal/src/history/mod.rs | 26 + .../src/history}/normalize.rs | 188 +++-- codex-rs/journal/src/history/tests.rs | 286 +++++++ codex-rs/journal/src/history/transform.rs | 92 +++ codex-rs/journal/src/journal.rs | 345 +++++++++ codex-rs/journal/src/lib.rs | 29 + codex-rs/journal/src/prompt_view.rs | 33 + codex-rs/journal/src/tests.rs | 349 +++++++++ codex-rs/protocol/src/journal.rs | 383 +++++++++ codex-rs/protocol/src/lib.rs | 1 + codex-rs/thread-store/Cargo.toml | 3 +- codex-rs/thread-store/src/error.rs | 2 +- codex-rs/thread-store/src/journal_writer.rs | 355 +++++++++ codex-rs/thread-store/src/lib.rs | 2 + 44 files changed, 3329 insertions(+), 1010 deletions(-) delete mode 100644 codex-rs/core/src/context_manager/history.rs create mode 100644 codex-rs/core/src/state/history.rs rename codex-rs/core/src/{context_manager => state}/history_tests.rs (94%) create mode 100644 codex-rs/journal/BUILD.bazel create mode 100644 codex-rs/journal/Cargo.toml create mode 100644 codex-rs/journal/src/error.rs create mode 100644 codex-rs/journal/src/history/classify.rs create mode 100644 codex-rs/journal/src/history/estimate.rs create mode 100644 codex-rs/journal/src/history/mod.rs rename codex-rs/{core/src/context_manager => journal/src/history}/normalize.rs (69%) create mode 100644 codex-rs/journal/src/history/tests.rs create mode 100644 codex-rs/journal/src/history/transform.rs create mode 100644 codex-rs/journal/src/journal.rs create mode 100644 codex-rs/journal/src/lib.rs create mode 100644 codex-rs/journal/src/prompt_view.rs create mode 100644 codex-rs/journal/src/tests.rs create mode 100644 codex-rs/protocol/src/journal.rs create mode 100644 codex-rs/thread-store/src/journal_writer.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index f219096e35..507fe208c2 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2374,6 +2374,7 @@ dependencies = [ "codex-feedback", "codex-git-utils", "codex-hooks", + "codex-journal", "codex-login", "codex-mcp", "codex-memories-read", @@ -2815,6 +2816,23 @@ dependencies = [ "tempfile", ] +[[package]] +name = "codex-journal" +version = "0.0.0" +dependencies = [ + "base64 0.22.1", + "codex-protocol", + "codex-utils-cache", + "codex-utils-output-truncation", + "image", + "indexmap 2.13.0", + "pretty_assertions", + "serde_json", + "tempfile", + "thiserror 2.0.18", + "tracing", +] + [[package]] name = "codex-keyring-store" version = "0.0.0" @@ -3484,6 +3502,7 @@ dependencies = [ "async-trait", "chrono", "codex-git-utils", + "codex-journal", "codex-protocol", "codex-rollout", "codex-state", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index a9d6ea6c23..602b7b158c 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -98,6 +98,7 @@ members = [ "terminal-detection", "test-binary-support", "thread-store", + "journal", "uds", "codex-experimental-api-macros", "plugin", @@ -190,6 +191,7 @@ codex-stdio-to-uds = { path = "stdio-to-uds" } codex-terminal-detection = { path = "terminal-detection" } codex-test-binary-support = { path = "test-binary-support" } codex-thread-store = { path = "thread-store" } +codex-journal = { path = "journal" } codex-tools = { path = "tools" } codex-tui = { path = "tui" } codex-uds = { path = "uds" } diff --git a/codex-rs/config/src/config_toml.rs b/codex-rs/config/src/config_toml.rs index 3bd19f5568..38e5f524ab 100644 --- a/codex-rs/config/src/config_toml.rs +++ b/codex-rs/config/src/config_toml.rs @@ -207,7 +207,7 @@ pub struct ConfigToml { /// Ordered list of fallback filenames to look for when AGENTS.md is missing. pub project_doc_fallback_filenames: Option>, - /// Token budget applied when storing tool/function outputs in the context manager. + /// Token budget applied when storing tool/function outputs in thread history. pub tool_output_token_limit: Option, /// Maximum poll window for background terminal output (`write_stdin`), in milliseconds. diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 44c6aacac5..002982f593 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -60,6 +60,7 @@ codex-sandboxing = { workspace = true } codex-state = { workspace = true } codex-terminal-detection = { workspace = true } codex-thread-store = { workspace = true } +codex-journal = { workspace = true } codex-tools = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-cache = { workspace = true } diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index a17426d508..f8101661b1 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -3885,7 +3885,7 @@ "type": "boolean" }, "tool_output_token_limit": { - "description": "Token budget applied when storing tool/function outputs in the context manager.", + "description": "Token budget applied when storing tool/function outputs in thread history.", "format": "uint", "minimum": 0.0, "type": "integer" @@ -3946,4 +3946,4 @@ }, "title": "ConfigToml", "type": "object" -} \ No newline at end of file +} diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 8ce7fa2ea2..2c237a49cb 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -7,8 +7,10 @@ use crate::config::Config; use crate::config::ConfigBuilder; use crate::context::ContextualUserFragment; use crate::context::SubagentNotification; +use crate::state::history as session_history; use assert_matches::assert_matches; use codex_features::Feature; +use codex_journal::Journal; use codex_login::CodexAuth; use codex_protocol::AgentPath; use codex_protocol::config_types::ModeKind; @@ -34,6 +36,16 @@ use tokio::time::sleep; use tokio::time::timeout; use toml::Value as TomlValue; +trait JournalTestExt { + fn raw_items(&self) -> Vec; +} + +impl JournalTestExt for Journal { + fn raw_items(&self) -> Vec { + session_history::raw_items(self) + } +} + async fn test_config_with_cli_overrides( cli_overrides: Vec<(String, TomlValue)>, ) -> (TempDir, Config) { @@ -796,8 +808,9 @@ async fn spawn_agent_fork_flushes_parent_rollout_before_loading_history() { .await .expect("child thread should be registered"); let history = child_thread.codex.session.clone_history().await; + let history_items = history.raw_items(); assert!( - history_contains_text(history.raw_items(), "unflushed final answer"), + history_contains_text(&history_items, "unflushed final answer"), "forked child history should include unflushed assistant final answers after flushing the parent rollout" ); @@ -906,21 +919,22 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { .await .expect("child thread should be registered"); let history = child_thread.codex.session.clone_history().await; + let history_items = history.raw_items(); assert!( - !history_contains_text(history.raw_items(), "old parent context"), + !history_contains_text(&history_items, "old parent context"), "forked child history should drop parent context outside the requested last-N turn window" ); assert!( - !history_contains_text(history.raw_items(), "queued message"), + !history_contains_text(&history_items, "queued message"), "forked child history should drop queued inter-agent messages outside the requested last-N turn window" ); assert!( - !history_contains_text(history.raw_items(), "triggered context"), + !history_contains_text(&history_items, "triggered context"), "forked child history should filter assistant inter-agent messages even when they fall inside the requested last-N turn window" ); assert!( - history_contains_text(history.raw_items(), "current parent task"), + history_contains_text(&history_items, "current parent task"), "forked child history should keep the parent user message from the requested last-N turn window" ); diff --git a/codex-rs/core/src/arc_monitor.rs b/codex-rs/core/src/arc_monitor.rs index c7f12e1024..7eb846d2e1 100644 --- a/codex-rs/core/src/arc_monitor.rs +++ b/codex-rs/core/src/arc_monitor.rs @@ -9,6 +9,7 @@ use crate::compact::content_items_to_text; use crate::event_mapping::is_contextual_user_message_content; use crate::session::session::Session; use crate::session::turn_context::TurnContext; +use crate::state::history as session_history; use codex_login::default_client::build_reqwest_client; use codex_protocol::models::MessagePhase; use codex_protocol::models::ResponseItem; @@ -224,7 +225,8 @@ async fn build_arc_monitor_request( protection_client_callsite: &'static str, ) -> ArcMonitorRequest { let history = sess.clone_history().await; - let mut messages = build_arc_monitor_messages(history.raw_items()); + let history_items = session_history::raw_items(&history); + let mut messages = build_arc_monitor_messages(&history_items); if messages.is_empty() { messages.push(build_arc_monitor_message( "user", diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index 1ebc307e33..58e72913b5 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -9,6 +9,7 @@ use crate::session::PreviousTurnSettings; use crate::session::session::Session; use crate::session::turn::get_last_assistant_message_from_turn; use crate::session::turn_context::TurnContext; +use crate::state::history as session_history; use crate::util::backoff; use codex_analytics::CodexCompactionEvent; use codex_analytics::CompactionImplementation; @@ -160,7 +161,8 @@ async fn run_compact_task_inner_impl( let initial_input_for_turn: ResponseInputItem = ResponseInputItem::from(input); let mut history = sess.clone_history().await; - history.record_items( + session_history::record_items( + &mut history, &[initial_input_for_turn.into()], turn_context.truncation_policy, ); @@ -176,9 +178,8 @@ async fn run_compact_task_inner_impl( loop { // Clone is required because of the loop - let turn_input = history - .clone() - .for_prompt(&turn_context.model_info.input_modalities); + let turn_input = + session_history::for_prompt(&history, &turn_context.model_info.input_modalities); let turn_input_len = turn_input.len(); let prompt = Prompt { input: turn_input, @@ -218,7 +219,7 @@ async fn run_compact_task_inner_impl( error!( "Context window exceeded while compacting; removing oldest history item. Error: {e}" ); - history.remove_first_item(); + session_history::remove_first_item(&mut history); truncated_count += 1; retries = 0; continue; @@ -250,10 +251,11 @@ async fn run_compact_task_inner_impl( } let history_snapshot = sess.clone_history().await; - let history_items = history_snapshot.raw_items(); - let summary_suffix = get_last_assistant_message_from_turn(history_items).unwrap_or_default(); + let history_items = session_history::raw_items(&history_snapshot); + let summary_suffix = + get_last_assistant_message_from_turn(history_items.as_slice()).unwrap_or_default(); let summary_text = format!("{SUMMARY_PREFIX}\n{summary_suffix}"); - let user_messages = collect_user_messages(history_items); + let user_messages = collect_user_messages(history_items.as_slice()); let mut new_history = build_compacted_history(Vec::new(), &user_messages, &summary_text); diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index d8adb20772..aa6b4cf333 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -6,17 +6,18 @@ use crate::compact::CompactionAnalyticsAttempt; use crate::compact::InitialContextInjection; use crate::compact::compaction_status_from_result; use crate::compact::insert_initial_context_before_last_real_user_or_summary; -use crate::context_manager::ContextManager; -use crate::context_manager::TotalTokenUsageBreakdown; use crate::context_manager::estimate_response_item_model_visible_bytes; use crate::context_manager::is_codex_generated_item; use crate::session::session::Session; use crate::session::turn::built_tools; use crate::session::turn_context::TurnContext; +use crate::state::TotalTokenUsageBreakdown; +use crate::state::history as session_history; use codex_analytics::CompactionImplementation; use codex_analytics::CompactionPhase; use codex_analytics::CompactionReason; use codex_analytics::CompactionTrigger; +use codex_journal::Journal; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; use codex_protocol::items::ContextCompactionItem; @@ -144,8 +145,9 @@ async fn run_remote_compact_task_inner_impl( // This is the history selected for remote compaction, after any trimming required to fit the // compact endpoint. The checkpoint below records it separately from the next sampling request, // whose prompt will repeat current developer/context prefix items. - let trace_input_history = history.raw_items().to_vec(); - let prompt_input = history.for_prompt(&turn_context.model_info.input_modalities); + let trace_input_history = session_history::raw_items(&history); + let prompt_input = + session_history::for_prompt(&history, &turn_context.model_info.input_modalities); let tool_router = built_tools( sess.as_ref(), turn_context.as_ref(), @@ -325,7 +327,7 @@ fn log_remote_compact_failure( } fn trim_function_call_history_to_fit_context_window( - history: &mut ContextManager, + history: &mut Journal, turn_context: &TurnContext, base_instructions: &BaseInstructions, ) -> usize { @@ -334,17 +336,17 @@ fn trim_function_call_history_to_fit_context_window( return deleted_items; }; - while history - .estimate_token_count_with_base_instructions(base_instructions) + while session_history::estimate_token_count_with_base_instructions(history, base_instructions) .is_some_and(|estimated_tokens| estimated_tokens > context_window) { - let Some(last_item) = history.raw_items().last() else { + let history_items = session_history::raw_items(history); + let Some(last_item) = history_items.last() else { break; }; if !is_codex_generated_item(last_item) { break; } - if !history.remove_last_item() { + if !session_history::remove_last_item(history) { break; } deleted_items += 1; diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index fe9a8e4334..58794e52d4 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -559,7 +559,7 @@ pub struct Config { /// Additional filenames to try when looking for project-level docs. pub project_doc_fallback_filenames: Vec, - /// Token budget applied when storing tool/function outputs in the context manager. + /// Token budget applied when storing tool/function outputs in thread history. pub tool_output_token_limit: Option, /// Maximum number of agent threads that can be open concurrently. diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs deleted file mode 100644 index 415a601980..0000000000 --- a/codex-rs/core/src/context_manager/history.rs +++ /dev/null @@ -1,726 +0,0 @@ -use crate::context_manager::normalize; -use crate::event_mapping::has_non_contextual_dev_message_content; -use crate::event_mapping::is_contextual_dev_message_content; -use crate::event_mapping::is_contextual_user_message_content; -use crate::session::turn_context::TurnContext; -use base64::Engine; -use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; -use codex_protocol::models::BaseInstructions; -use codex_protocol::models::ContentItem; -use codex_protocol::models::FunctionCallOutputBody; -use codex_protocol::models::FunctionCallOutputContentItem; -use codex_protocol::models::FunctionCallOutputPayload; -use codex_protocol::models::ImageDetail; -use codex_protocol::models::ResponseItem; -use codex_protocol::openai_models::InputModality; -use codex_protocol::protocol::InterAgentCommunication; -use codex_protocol::protocol::TokenUsage; -use codex_protocol::protocol::TokenUsageInfo; -use codex_protocol::protocol::TurnContextItem; -use codex_utils_cache::BlockingLruCache; -use codex_utils_cache::sha1_digest; -use codex_utils_output_truncation::TruncationPolicy; -use codex_utils_output_truncation::approx_bytes_for_tokens; -use codex_utils_output_truncation::approx_token_count; -use codex_utils_output_truncation::approx_tokens_from_byte_count_i64; -use codex_utils_output_truncation::truncate_function_output_items_with_policy; -use codex_utils_output_truncation::truncate_text; -use std::num::NonZeroUsize; -use std::ops::Deref; -use std::sync::LazyLock; - -/// Transcript of thread history -#[derive(Debug, Clone, Default)] -pub(crate) struct ContextManager { - /// The oldest items are at the beginning of the vector. - items: Vec, - /// Bumped whenever history is rewritten, such as compaction or rollback. - history_version: u64, - token_info: Option, - /// Reference context snapshot used for diffing and producing model-visible - /// settings update items. - /// - /// This is the baseline for the next regular model turn, and may already - /// match the current turn after context updates are persisted. - /// - /// When this is `None`, settings diffing treats the next turn as having no - /// baseline and emits a full reinjection of context state. Rollback may - /// also clear this when it trims a mixed initial-context developer bundle - /// whose non-diff fragments no longer exist in the surviving history. - reference_context_item: Option, -} - -#[derive(Debug, Clone, Copy, Default)] -pub(crate) struct TotalTokenUsageBreakdown { - pub last_api_response_total_tokens: i64, - pub all_history_items_model_visible_bytes: i64, - pub estimated_tokens_of_items_added_since_last_successful_api_response: i64, - pub estimated_bytes_of_items_added_since_last_successful_api_response: i64, -} - -impl ContextManager { - pub(crate) fn new() -> Self { - Self { - items: Vec::new(), - history_version: 0, - token_info: TokenUsageInfo::new_or_append( - &None, &None, /*model_context_window*/ None, - ), - reference_context_item: None, - } - } - - pub(crate) fn token_info(&self) -> Option { - self.token_info.clone() - } - - pub(crate) fn set_token_info(&mut self, info: Option) { - self.token_info = info; - } - - pub(crate) fn set_reference_context_item(&mut self, item: Option) { - self.reference_context_item = item; - } - - pub(crate) fn reference_context_item(&self) -> Option { - self.reference_context_item.clone() - } - - pub(crate) fn set_token_usage_full(&mut self, context_window: i64) { - match &mut self.token_info { - Some(info) => info.fill_to_context_window(context_window), - None => { - self.token_info = Some(TokenUsageInfo::full_context_window(context_window)); - } - } - } - - /// `items` is ordered from oldest to newest. - pub(crate) fn record_items(&mut self, items: I, policy: TruncationPolicy) - where - I: IntoIterator, - I::Item: std::ops::Deref, - { - for item in items { - let item_ref = item.deref(); - if !is_api_message(item_ref) { - continue; - } - - let processed = self.process_item(item_ref, policy); - self.items.push(processed); - } - } - - /// Returns the history prepared for sending to the model. This applies a proper - /// normalization and drops un-suited items. When `input_modalities` does not - /// include `InputModality::Image`, images are stripped from messages and tool - /// outputs. - pub(crate) fn for_prompt(mut self, input_modalities: &[InputModality]) -> Vec { - self.normalize_history(input_modalities); - self.items - } - - /// Returns raw items in the history. - pub(crate) fn raw_items(&self) -> &[ResponseItem] { - &self.items - } - - pub(crate) fn history_version(&self) -> u64 { - self.history_version - } - - // Estimate token usage using byte-based heuristics from the truncation helpers. - // This is a coarse lower bound, not a tokenizer-accurate count. - pub(crate) fn estimate_token_count(&self, turn_context: &TurnContext) -> Option { - let model_info = &turn_context.model_info; - let personality = turn_context.personality.or(turn_context.config.personality); - let base_instructions = BaseInstructions { - text: model_info.get_model_instructions(personality), - }; - self.estimate_token_count_with_base_instructions(&base_instructions) - } - - pub(crate) fn estimate_token_count_with_base_instructions( - &self, - base_instructions: &BaseInstructions, - ) -> Option { - let base_tokens = - i64::try_from(approx_token_count(&base_instructions.text)).unwrap_or(i64::MAX); - - let items_tokens = self - .items - .iter() - .map(estimate_item_token_count) - .fold(0i64, i64::saturating_add); - - Some(base_tokens.saturating_add(items_tokens)) - } - - pub(crate) fn remove_first_item(&mut self) { - if !self.items.is_empty() { - // Remove the oldest item (front of the list). Items are ordered from - // oldest → newest, so index 0 is the first entry recorded. - let removed = self.items.remove(0); - // If the removed item participates in a call/output pair, also remove - // its corresponding counterpart to keep the invariants intact without - // running a full normalization pass. - normalize::remove_corresponding_for(&mut self.items, &removed); - } - } - - pub(crate) fn remove_last_item(&mut self) -> bool { - if let Some(removed) = self.items.pop() { - normalize::remove_corresponding_for(&mut self.items, &removed); - self.history_version = self.history_version.saturating_add(1); - true - } else { - false - } - } - - pub(crate) fn replace(&mut self, items: Vec) { - self.items = items; - self.history_version = self.history_version.saturating_add(1); - } - - /// Replace image content in the last turn if it originated from a tool output. - /// Returns true when a tool image was replaced, false otherwise. - pub(crate) fn replace_last_turn_images(&mut self, placeholder: &str) -> bool { - let Some(index) = self.items.iter().rposition(|item| { - matches!(item, ResponseItem::FunctionCallOutput { .. }) || is_user_turn_boundary(item) - }) else { - return false; - }; - - match &mut self.items[index] { - ResponseItem::FunctionCallOutput { output, .. } => { - let Some(content_items) = output.content_items_mut() else { - return false; - }; - let mut replaced = false; - let placeholder = placeholder.to_string(); - for item in content_items.iter_mut() { - if matches!(item, FunctionCallOutputContentItem::InputImage { .. }) { - *item = FunctionCallOutputContentItem::InputText { - text: placeholder.clone(), - }; - replaced = true; - } - } - if replaced { - self.history_version = self.history_version.saturating_add(1); - } - replaced - } - ResponseItem::Message { .. } => false, - _ => false, - } - } - - /// Drop the last `num_turns` instruction turns from this history. - /// - /// Instruction turns are history messages that should behave like a new prompt boundary: - /// ordinary user messages and structured assistant inter-agent instructions. - /// - /// This mirrors thread-rollback semantics: - /// - `num_turns == 0` is a no-op - /// - if there are no user turns, this is a no-op - /// - if `num_turns` exceeds the number of user turns, all user turns are dropped while - /// preserving any items that occurred before the first user message. - /// - /// If rollback trims a pre-turn developer message that mixes contextual fragments with - /// persistent developer text from `build_initial_context`, this also clears - /// `reference_context_item`. The surviving history no longer contains the full bundle that - /// established the prior baseline, so future turns must fall back to full reinjection instead - /// of diffing against stale state. - pub(crate) fn drop_last_n_user_turns(&mut self, num_turns: u32) { - if num_turns == 0 { - return; - } - - let snapshot = self.items.clone(); - let user_positions = user_message_positions(&snapshot); - let Some(&first_instruction_turn_idx) = user_positions.first() else { - self.replace(snapshot); - return; - }; - - let n_from_end = usize::try_from(num_turns).unwrap_or(usize::MAX); - let mut cut_idx = if n_from_end >= user_positions.len() { - first_instruction_turn_idx - } else { - user_positions[user_positions.len() - n_from_end] - }; - - cut_idx = - self.trim_pre_turn_context_updates(&snapshot, first_instruction_turn_idx, cut_idx); - - self.replace(snapshot[..cut_idx].to_vec()); - } - - pub(crate) fn update_token_info( - &mut self, - usage: &TokenUsage, - model_context_window: Option, - ) { - self.token_info = TokenUsageInfo::new_or_append( - &self.token_info, - &Some(usage.clone()), - model_context_window, - ); - } - - fn get_non_last_reasoning_items_tokens(&self) -> i64 { - // Get reasoning items excluding all the ones after the last instruction boundary. - let Some(last_user_index) = self.items.iter().rposition(is_user_turn_boundary) else { - return 0; - }; - - self.items - .iter() - .take(last_user_index) - .filter(|item| { - matches!( - item, - ResponseItem::Reasoning { - encrypted_content: Some(_), - .. - } - ) - }) - .map(estimate_item_token_count) - .fold(0i64, i64::saturating_add) - } - - // These are local items added after the most recent model-emitted item. - // They are not reflected in `last_token_usage.total_tokens`. - fn items_after_last_model_generated_item(&self) -> &[ResponseItem] { - let start = self - .items - .iter() - .rposition(is_model_generated_item) - .map_or(self.items.len(), |index| index.saturating_add(1)); - &self.items[start..] - } - - /// When true, the server already accounted for past reasoning tokens and - /// the client should not re-estimate them. - pub(crate) fn get_total_token_usage(&self, server_reasoning_included: bool) -> i64 { - let last_tokens = self - .token_info - .as_ref() - .map(|info| info.last_token_usage.total_tokens) - .unwrap_or(0); - let items_after_last_model_generated_tokens = self - .items_after_last_model_generated_item() - .iter() - .map(estimate_item_token_count) - .fold(0i64, i64::saturating_add); - if server_reasoning_included { - last_tokens.saturating_add(items_after_last_model_generated_tokens) - } else { - last_tokens - .saturating_add(self.get_non_last_reasoning_items_tokens()) - .saturating_add(items_after_last_model_generated_tokens) - } - } - - pub(crate) fn get_total_token_usage_breakdown(&self) -> TotalTokenUsageBreakdown { - let last_usage = self - .token_info - .as_ref() - .map(|info| info.last_token_usage.clone()) - .unwrap_or_default(); - let items_after_last_model_generated = self.items_after_last_model_generated_item(); - - TotalTokenUsageBreakdown { - last_api_response_total_tokens: last_usage.total_tokens, - all_history_items_model_visible_bytes: self - .items - .iter() - .map(estimate_response_item_model_visible_bytes) - .fold(0i64, i64::saturating_add), - estimated_tokens_of_items_added_since_last_successful_api_response: - items_after_last_model_generated - .iter() - .map(estimate_item_token_count) - .fold(0i64, i64::saturating_add), - estimated_bytes_of_items_added_since_last_successful_api_response: - items_after_last_model_generated - .iter() - .map(estimate_response_item_model_visible_bytes) - .fold(0i64, i64::saturating_add), - } - } - - /// This function enforces a couple of invariants on the in-memory history: - /// 1. every call (function/custom) has a corresponding output entry - /// 2. every output has a corresponding call entry - /// 3. when images are unsupported, image content is stripped from messages and tool outputs - fn normalize_history(&mut self, input_modalities: &[InputModality]) { - // all function/tool calls must have a corresponding output - normalize::ensure_call_outputs_present(&mut self.items); - - // all outputs must have a corresponding function/tool call - normalize::remove_orphan_outputs(&mut self.items); - - // strip images when model does not support them - normalize::strip_images_when_unsupported(input_modalities, &mut self.items); - } - - fn process_item(&self, item: &ResponseItem, policy: TruncationPolicy) -> ResponseItem { - let policy_with_serialization_budget = policy * 1.2; - match item { - ResponseItem::FunctionCallOutput { call_id, output } => { - ResponseItem::FunctionCallOutput { - call_id: call_id.clone(), - output: truncate_function_output_payload( - output, - policy_with_serialization_budget, - ), - } - } - ResponseItem::CustomToolCallOutput { - call_id, - name, - output, - } => ResponseItem::CustomToolCallOutput { - call_id: call_id.clone(), - name: name.clone(), - output: truncate_function_output_payload(output, policy_with_serialization_budget), - }, - ResponseItem::Message { .. } - | ResponseItem::Reasoning { .. } - | ResponseItem::LocalShellCall { .. } - | ResponseItem::FunctionCall { .. } - | ResponseItem::ToolSearchCall { .. } - | ResponseItem::ToolSearchOutput { .. } - | ResponseItem::WebSearchCall { .. } - | ResponseItem::ImageGenerationCall { .. } - | ResponseItem::CustomToolCall { .. } - | ResponseItem::Compaction { .. } - | ResponseItem::Other => item.clone(), - } - } - - /// Walk backward from a rollback cut and trim contiguous pre-turn context-update items. - /// - /// Returns the adjusted cut index after removing contextual developer/user items immediately - /// above the rolled-back turn boundary. - /// - /// `first_instruction_turn_idx` is the earliest rollback-eligible instruction-turn boundary - /// in `snapshot`; the trim walk never crosses it so any session-prefix items that predate the - /// first real turn survive rollback. - /// - /// `cut_idx` is the tentative slice boundary after dropping the requested number of - /// instruction turns, before stripping contextual pre-turn items that sit immediately above - /// that boundary. - /// - /// If any trimmed developer message was a mixed `build_initial_context` bundle containing both - /// rollback-trimmable contextual fragments and persistent developer text, this also clears the - /// stored `reference_context_item` baseline so the next real turn falls back to full - /// reinjection. - fn trim_pre_turn_context_updates( - &mut self, - snapshot: &[ResponseItem], - first_instruction_turn_idx: usize, - mut cut_idx: usize, - ) -> usize { - while cut_idx > first_instruction_turn_idx { - match &snapshot[cut_idx - 1] { - ResponseItem::Message { role, content, .. } - if role == "developer" && is_contextual_dev_message_content(content) => - { - if has_non_contextual_dev_message_content(content) { - // Mixed `build_initial_context` bundles are not reconstructible from - // steady-state diffs once trimmed, so the next real turn must fully - // reinject context instead of diffing against a stale baseline. - self.reference_context_item = None; - } - cut_idx -= 1; - } - ResponseItem::Message { role, content, .. } - if role == "user" && is_contextual_user_message_content(content) => - { - cut_idx -= 1; - } - _ => break, - } - } - cut_idx - } -} - -fn truncate_function_output_payload( - output: &FunctionCallOutputPayload, - policy: TruncationPolicy, -) -> FunctionCallOutputPayload { - let body = match &output.body { - FunctionCallOutputBody::Text(content) => { - FunctionCallOutputBody::Text(truncate_text(content, policy)) - } - FunctionCallOutputBody::ContentItems(items) => FunctionCallOutputBody::ContentItems( - truncate_function_output_items_with_policy(items, policy), - ), - }; - - FunctionCallOutputPayload { - body, - success: output.success, - } -} - -/// API messages include every non-system item (user/assistant messages, reasoning, -/// tool calls, tool outputs, shell calls, web-search calls, and image-generation -/// calls). -fn is_api_message(message: &ResponseItem) -> bool { - match message { - ResponseItem::Message { role, .. } => role.as_str() != "system", - ResponseItem::FunctionCallOutput { .. } - | ResponseItem::FunctionCall { .. } - | ResponseItem::ToolSearchCall { .. } - | ResponseItem::ToolSearchOutput { .. } - | ResponseItem::CustomToolCall { .. } - | ResponseItem::CustomToolCallOutput { .. } - | ResponseItem::LocalShellCall { .. } - | ResponseItem::Reasoning { .. } - | ResponseItem::WebSearchCall { .. } - | ResponseItem::ImageGenerationCall { .. } - | ResponseItem::Compaction { .. } => true, - ResponseItem::Other => false, - } -} - -fn estimate_reasoning_length(encoded_len: usize) -> usize { - encoded_len - .saturating_mul(3) - .checked_div(4) - .unwrap_or(0) - .saturating_sub(650) -} - -fn estimate_item_token_count(item: &ResponseItem) -> i64 { - let model_visible_bytes = estimate_response_item_model_visible_bytes(item); - approx_tokens_from_byte_count_i64(model_visible_bytes) -} - -/// Approximate model-visible byte cost for one image input. -/// -/// The estimator later converts bytes to tokens using a 4-bytes/token heuristic -/// with ceiling division, so 7,373 bytes maps to approximately 1,844 tokens. -const RESIZED_IMAGE_BYTES_ESTIMATE: i64 = 7373; -// See https://platform.openai.com/docs/guides/images-vision#calculating-costs. -// Use a direct 32px patch count only for `detail: "original"`; -// all other image inputs continue to use `RESIZED_IMAGE_BYTES_ESTIMATE`. -const ORIGINAL_IMAGE_PATCH_SIZE: u32 = 32; -// See https://platform.openai.com/docs/guides/images-vision#model-sizing-behavior. -// Keep this hard-coded for now; move it into model capabilities if the patch -// budget starts changing often across model releases. -const ORIGINAL_IMAGE_MAX_PATCHES: usize = 10_000; -const ORIGINAL_IMAGE_ESTIMATE_CACHE_SIZE: usize = 32; - -static ORIGINAL_IMAGE_ESTIMATE_CACHE: LazyLock>> = - LazyLock::new(|| { - BlockingLruCache::new( - NonZeroUsize::new(ORIGINAL_IMAGE_ESTIMATE_CACHE_SIZE).unwrap_or(NonZeroUsize::MIN), - ) - }); - -pub(crate) fn estimate_response_item_model_visible_bytes(item: &ResponseItem) -> i64 { - match item { - ResponseItem::Reasoning { - encrypted_content: Some(content), - .. - } - | ResponseItem::Compaction { - encrypted_content: content, - } => i64::try_from(estimate_reasoning_length(content.len())).unwrap_or(i64::MAX), - item => { - let raw = serde_json::to_string(item) - .map(|serialized| i64::try_from(serialized.len()).unwrap_or(i64::MAX)) - .unwrap_or_default(); - let (payload_bytes, replacement_bytes) = image_data_url_estimate_adjustment(item); - if payload_bytes == 0 || replacement_bytes == 0 { - raw - } else { - // Replace raw base64 payload bytes with a per-image estimate. - // We intentionally preserve the data URL prefix and JSON - // wrapper bytes already included in `raw`. - raw.saturating_sub(payload_bytes) - .saturating_add(replacement_bytes) - } - } - } -} - -/// Returns the base64 payload byte length for inline image data URLs that are -/// eligible for token-estimation discounting. -/// -/// We only discount payloads for `data:image/...;base64,...` URLs (case -/// insensitive markers) and leave everything else at raw serialized size. -fn parse_base64_image_data_url(url: &str) -> Option<&str> { - if !url - .get(.."data:".len()) - .is_some_and(|prefix| prefix.eq_ignore_ascii_case("data:")) - { - return None; - } - let comma_index = url.find(',')?; - let metadata = &url[..comma_index]; - let payload = &url[comma_index + 1..]; - // Parse the media type and parameters without decoding. This keeps the - // estimator cheap while ensuring we only apply the fixed-cost image - // heuristic to image-typed base64 data URLs. - let metadata_without_scheme = &metadata["data:".len()..]; - let mut metadata_parts = metadata_without_scheme.split(';'); - let mime_type = metadata_parts.next().unwrap_or_default(); - let has_base64_marker = metadata_parts.any(|part| part.eq_ignore_ascii_case("base64")); - if !mime_type - .get(.."image/".len()) - .is_some_and(|prefix| prefix.eq_ignore_ascii_case("image/")) - { - return None; - } - if !has_base64_marker { - return None; - } - Some(payload) -} - -fn estimate_original_image_bytes(image_url: &str) -> Option { - let key = sha1_digest(image_url.as_bytes()); - ORIGINAL_IMAGE_ESTIMATE_CACHE.get_or_insert_with(key, || { - let payload = match parse_base64_image_data_url(image_url) { - Some(payload) => payload, - None => { - tracing::trace!("skipping original-detail estimate for non-base64 image data URL"); - return None; - } - }; - let bytes = match BASE64_STANDARD.decode(payload) { - Ok(bytes) => bytes, - Err(error) => { - tracing::trace!("failed to decode original-detail image payload: {error}"); - return None; - } - }; - let dynamic = match image::load_from_memory(&bytes) { - Ok(dynamic) => dynamic, - Err(error) => { - tracing::trace!("failed to decode original-detail image bytes: {error}"); - return None; - } - }; - let width = i64::from(dynamic.width()); - let height = i64::from(dynamic.height()); - let patch_size = i64::from(ORIGINAL_IMAGE_PATCH_SIZE); - let patches_wide = width.saturating_add(patch_size.saturating_sub(1)) / patch_size; - let patches_high = height.saturating_add(patch_size.saturating_sub(1)) / patch_size; - let patch_count = patches_wide.saturating_mul(patches_high); - let patch_count = usize::try_from(patch_count).unwrap_or(usize::MAX); - let patch_count = patch_count.min(ORIGINAL_IMAGE_MAX_PATCHES); - Some(i64::try_from(approx_bytes_for_tokens(patch_count)).unwrap_or(i64::MAX)) - }) -} - -/// Scans one response item for discount-eligible inline image data URLs and -/// returns: -/// - total base64 payload bytes to subtract from raw serialized size -/// - total replacement byte estimate for those images -fn image_data_url_estimate_adjustment(item: &ResponseItem) -> (i64, i64) { - let mut payload_bytes = 0i64; - let mut replacement_bytes = 0i64; - - let mut accumulate = |image_url: &str, detail: Option| { - if let Some(payload_len) = parse_base64_image_data_url(image_url).map(str::len) { - payload_bytes = - payload_bytes.saturating_add(i64::try_from(payload_len).unwrap_or(i64::MAX)); - replacement_bytes = replacement_bytes.saturating_add(match detail { - Some(ImageDetail::Original) => { - estimate_original_image_bytes(image_url).unwrap_or(RESIZED_IMAGE_BYTES_ESTIMATE) - } - _ => RESIZED_IMAGE_BYTES_ESTIMATE, - }); - } - }; - - match item { - ResponseItem::Message { content, .. } => { - for content_item in content { - if let ContentItem::InputImage { image_url, detail } = content_item { - accumulate(image_url, *detail); - } - } - } - ResponseItem::FunctionCallOutput { output, .. } - | ResponseItem::CustomToolCallOutput { output, .. } => { - if let FunctionCallOutputBody::ContentItems(items) = &output.body { - for content_item in items { - if let FunctionCallOutputContentItem::InputImage { image_url, detail } = - content_item - { - accumulate(image_url, *detail); - } - } - } - } - _ => {} - } - - (payload_bytes, replacement_bytes) -} - -fn is_model_generated_item(item: &ResponseItem) -> bool { - match item { - ResponseItem::Message { role, .. } => role == "assistant", - ResponseItem::Reasoning { .. } - | ResponseItem::FunctionCall { .. } - | ResponseItem::ToolSearchCall { .. } - | ResponseItem::WebSearchCall { .. } - | ResponseItem::ImageGenerationCall { .. } - | ResponseItem::CustomToolCall { .. } - | ResponseItem::LocalShellCall { .. } - | ResponseItem::Compaction { .. } => true, - ResponseItem::FunctionCallOutput { .. } - | ResponseItem::ToolSearchOutput { .. } - | ResponseItem::CustomToolCallOutput { .. } - | ResponseItem::Other => false, - } -} - -pub(crate) fn is_codex_generated_item(item: &ResponseItem) -> bool { - matches!( - item, - ResponseItem::FunctionCallOutput { .. } - | ResponseItem::ToolSearchOutput { .. } - | ResponseItem::CustomToolCallOutput { .. } - ) || matches!(item, ResponseItem::Message { role, .. } if role == "developer") -} - -pub(crate) fn is_user_turn_boundary(item: &ResponseItem) -> bool { - let ResponseItem::Message { role, content, .. } = item else { - return false; - }; - - (role == "user" && !is_contextual_user_message_content(content)) - || (role == "assistant" && is_inter_agent_instruction_content(content)) -} - -fn is_inter_agent_instruction_content(content: &[ContentItem]) -> bool { - InterAgentCommunication::is_message_content(content) -} - -fn user_message_positions(items: &[ResponseItem]) -> Vec { - let mut positions = Vec::new(); - for (idx, item) in items.iter().enumerate() { - if is_user_turn_boundary(item) { - positions.push(idx); - } - } - positions -} - -#[cfg(test)] -#[path = "history_tests.rs"] -mod tests; diff --git a/codex-rs/core/src/context_manager/mod.rs b/codex-rs/core/src/context_manager/mod.rs index 853f8af5ac..783fe8d3a1 100644 --- a/codex-rs/core/src/context_manager/mod.rs +++ b/codex-rs/core/src/context_manager/mod.rs @@ -1,9 +1,5 @@ -mod history; -mod normalize; pub(crate) mod updates; -pub(crate) use history::ContextManager; -pub(crate) use history::TotalTokenUsageBreakdown; -pub(crate) use history::estimate_response_item_model_visible_bytes; -pub(crate) use history::is_codex_generated_item; -pub(crate) use history::is_user_turn_boundary; +pub(crate) use codex_journal::history::estimate_response_item_model_visible_bytes; +pub(crate) use codex_journal::history::is_codex_generated_item; +pub(crate) use codex_journal::history::is_user_turn_boundary; diff --git a/codex-rs/core/src/context_manager/updates.rs b/codex-rs/core/src/context_manager/updates.rs index 1bc2cb0895..ab1bf02458 100644 --- a/codex-rs/core/src/context_manager/updates.rs +++ b/codex-rs/core/src/context_manager/updates.rs @@ -12,17 +12,30 @@ use crate::session::turn_context::TurnContext; use crate::shell::Shell; use codex_execpolicy::Policy; use codex_features::Feature; +use codex_journal::Journal; +use codex_journal::JournalContextItem; +use codex_journal::JournalContextKey; +use codex_journal::JournalEntry; +use codex_journal::JournalItem; +use codex_journal::PromptMessage; +use codex_journal::PromptMessageRole; use codex_protocol::config_types::Personality; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelInfo; use codex_protocol::protocol::TurnContextItem; +const PROMPT_BUNDLE_KEY_PREFIX: &str = "prompt"; +const DEVELOPER_BUNDLE: &str = "developer"; +const USAGE_HINT_BUNDLE: &str = "usage_hint"; +const CONTEXTUAL_USER_BUNDLE: &str = "contextual_user"; +const GUARDIAN_BUNDLE: &str = "guardian"; + fn build_environment_update_item( previous: Option<&TurnContextItem>, next: &TurnContext, shell: &Shell, -) -> Option { +) -> Option { if !next.config.include_environment_context { return None; } @@ -34,9 +47,7 @@ fn build_environment_update_item( return None; } - Some(ContextualUserFragment::into( - EnvironmentContext::diff_from_turn_context_item(prev, &next_context), - )) + Some(EnvironmentContext::diff_from_turn_context_item(prev, &next_context).render()) } fn build_permissions_update_item( @@ -175,64 +186,183 @@ pub(crate) fn build_model_instructions_update_item( Some(ModelSwitchInstructions::new(model_instructions).render()) } -pub(crate) fn build_developer_update_item(text_sections: Vec) -> Option { - build_text_message("developer", text_sections) +pub(crate) fn developer_context_entry( + name: &str, + prompt_order: i64, + text: String, +) -> Option { + context_entry( + DEVELOPER_BUNDLE, + name, + prompt_order, + PromptMessage::developer_text(text), + ) } -pub(crate) fn build_contextual_user_message(text_sections: Vec) -> Option { - build_text_message("user", text_sections) +pub(crate) fn usage_hint_context_entry( + name: &str, + prompt_order: i64, + text: String, +) -> Option { + context_entry( + USAGE_HINT_BUNDLE, + name, + prompt_order, + PromptMessage::developer_text(text), + ) } -fn build_text_message(role: &str, text_sections: Vec) -> Option { - if text_sections.is_empty() { +pub(crate) fn contextual_user_context_entry( + name: &str, + prompt_order: i64, + text: String, +) -> Option { + context_entry( + CONTEXTUAL_USER_BUNDLE, + name, + prompt_order, + PromptMessage::user_text(text), + ) +} + +pub(crate) fn guardian_context_entry( + name: &str, + prompt_order: i64, + text: String, +) -> Option { + context_entry( + GUARDIAN_BUNDLE, + name, + prompt_order, + PromptMessage::developer_text(text), + ) +} + +fn context_entry( + bundle: &str, + name: &str, + prompt_order: i64, + message: PromptMessage, +) -> Option { + if message.content.is_empty() + || message.content.iter().all(|item| match item { + ContentItem::InputText { text } | ContentItem::OutputText { text } => { + text.trim().is_empty() + } + ContentItem::InputImage { .. } => false, + }) + { return None; } - let content = text_sections - .into_iter() - .map(|text| ContentItem::InputText { text }) - .collect(); - - Some(ResponseItem::Message { - id: None, - role: role.to_string(), - content, - phase: None, - }) + Some(JournalEntry::new( + [PROMPT_BUNDLE_KEY_PREFIX, bundle, name], + JournalContextItem::new(JournalContextKey::new(bundle, name, None), message) + .with_prompt_order(prompt_order), + )) } -pub(crate) fn build_settings_update_items( +pub(crate) fn render_context_entries(entries: Vec) -> Vec { + let flattened = match Journal::from_entries(entries).flatten() { + Ok(flattened) => flattened, + Err(error) => unreachable!("context-only journal entries should flatten: {error}"), + }; + + let mut rendered = Vec::new(); + let mut current_bundle: Option> = None; + let mut current_role: Option = None; + let mut current_content = Vec::new(); + + for entry in flattened.entries() { + let JournalItem::Context(item) = &entry.item else { + continue; + }; + let bundle = entry + .key + .parts() + .iter() + .take(2) + .cloned() + .collect::>(); + + if current_bundle.as_ref() != Some(&bundle) || current_role != Some(item.message.role) { + flush_prompt_message(&mut rendered, &mut current_role, &mut current_content); + current_bundle = Some(bundle); + current_role = Some(item.message.role); + } + + current_content.extend(item.message.content.clone()); + } + + flush_prompt_message(&mut rendered, &mut current_role, &mut current_content); + rendered +} + +fn flush_prompt_message( + rendered: &mut Vec, + role: &mut Option, + content: &mut Vec, +) { + let Some(role) = role.take() else { + return; + }; + if content.is_empty() { + return; + } + rendered.push(ResponseItem::from(PromptMessage::new( + role, + std::mem::take(content), + ))); +} + +pub(crate) fn build_settings_update_entries( previous: Option<&TurnContextItem>, previous_turn_settings: Option<&PreviousTurnSettings>, next: &TurnContext, shell: &Shell, exec_policy: &Policy, personality_feature_enabled: bool, -) -> Vec { +) -> Vec { // TODO(ccunningham): build_settings_update_items still does not cover every // model-visible item emitted by build_initial_context. Persist the remaining // inputs or add explicit replay events so fork/resume can diff everything // deterministically. - let contextual_user_message = build_environment_update_item(previous, next, shell); - let developer_update_sections = [ - // Keep model-switch instructions first so model-specific guidance is read before - // any other context diffs on this turn. - build_model_instructions_update_item(previous_turn_settings, next), - build_permissions_update_item(previous, next, exec_policy), - build_collaboration_mode_update_item(previous, next), - build_realtime_update_item(previous, previous_turn_settings, next), - build_personality_update_item(previous, next, personality_feature_enabled), - ] - .into_iter() - .flatten() - .collect(); + let mut entries = Vec::with_capacity(6); - let mut items = Vec::with_capacity(2); - if let Some(developer_message) = build_developer_update_item(developer_update_sections) { - items.push(developer_message); + if let Some(item) = + build_model_instructions_update_item(previous_turn_settings, next).and_then(|text| { + // Keep model-switch instructions first so model-specific guidance is read before + // any other context diffs on this turn. + developer_context_entry("model_switch", 10, text) + }) + { + entries.push(item); } - if let Some(contextual_user_message) = contextual_user_message { - items.push(contextual_user_message); + if let Some(item) = build_permissions_update_item(previous, next, exec_policy) + .and_then(|text| developer_context_entry("permissions", 20, text)) + { + entries.push(item); } - items + if let Some(item) = build_collaboration_mode_update_item(previous, next) + .and_then(|text| developer_context_entry("collaboration_mode", 30, text)) + { + entries.push(item); + } + if let Some(item) = build_realtime_update_item(previous, previous_turn_settings, next) + .and_then(|text| developer_context_entry("realtime", 40, text)) + { + entries.push(item); + } + if let Some(item) = build_personality_update_item(previous, next, personality_feature_enabled) + .and_then(|text| developer_context_entry("personality", 50, text)) + { + entries.push(item); + } + if let Some(item) = build_environment_update_item(previous, next, shell) + .and_then(|text| contextual_user_context_entry("environment", 60, text)) + { + entries.push(item); + } + + entries } diff --git a/codex-rs/core/src/guardian/prompt.rs b/codex-rs/core/src/guardian/prompt.rs index b1b132a984..d278b6e186 100644 --- a/codex-rs/core/src/guardian/prompt.rs +++ b/codex-rs/core/src/guardian/prompt.rs @@ -10,6 +10,7 @@ use serde_json::Value; use crate::compact::content_items_to_text; use crate::event_mapping::is_contextual_user_message_content; use crate::session::session::Session; +use crate::state::history as session_history; use codex_utils_output_truncation::approx_bytes_for_tokens; use codex_utils_output_truncation::approx_token_count; use codex_utils_output_truncation::approx_tokens_from_byte_count; @@ -93,9 +94,10 @@ pub(crate) async fn build_guardian_prompt_items( mode: GuardianPromptMode, ) -> serde_json::Result { let history = session.clone_history().await; - let transcript_entries = collect_guardian_transcript_entries(history.raw_items()); + let history_items = session_history::raw_items(&history); + let transcript_entries = collect_guardian_transcript_entries(&history_items); let transcript_cursor = GuardianTranscriptCursor { - parent_history_version: history.history_version(), + parent_history_version: session.history_version().await, transcript_entry_count: transcript_entries.len(), }; let planned_action_json = format_guardian_action_pretty(&request)?; diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index cc1203c2db..1e1d755d16 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -18,6 +18,7 @@ use crate::config::Config; use crate::session::session::Session; use crate::session::turn::build_prompt; use crate::session::turn::built_tools; +use crate::state::history as session_history; use crate::thread_manager::ThreadManager; /// Build the model-visible `input` list for a single debug turn. @@ -73,10 +74,9 @@ pub(crate) async fn build_prompt_input_from_session( .await; } - let prompt_input = sess - .clone_history() - .await - .for_prompt(&turn_context.model_info.input_modalities); + let history = sess.clone_history().await; + let prompt_input = + session_history::for_prompt(&history, &turn_context.model_info.input_modalities); let router = built_tools( sess, turn_context.as_ref(), diff --git a/codex-rs/core/src/realtime_context.rs b/codex-rs/core/src/realtime_context.rs index f9bb49a425..cb7558ec10 100644 --- a/codex-rs/core/src/realtime_context.rs +++ b/codex-rs/core/src/realtime_context.rs @@ -1,6 +1,7 @@ use crate::compact::content_items_to_text; use crate::event_mapping::is_contextual_user_message_content; use crate::session::session::Session; +use crate::state::history as session_history; use chrono::Utc; use codex_exec_server::LOCAL_FS; use codex_git_utils::resolve_root_git_project_for_trust; @@ -62,7 +63,8 @@ pub(crate) async fn build_realtime_startup_context( let config = sess.get_config().await; let cwd = config.cwd.clone(); let history = sess.clone_history().await; - let current_thread_section = build_current_thread_section(history.raw_items()); + let history_items = session_history::raw_items(&history); + let current_thread_section = build_current_thread_section(&history_items); let recent_threads = load_recent_threads(sess).await; let recent_work_section = build_recent_work_section(&cwd, &recent_threads).await; let workspace_section = build_workspace_section_with_user_root(&cwd, home_dir()).await; diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 4965cdb389..310e452358 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -19,6 +19,7 @@ use crate::realtime_context::truncate_realtime_text_to_token_budget; use crate::realtime_conversation::REALTIME_USER_TEXT_PREFIX; use crate::realtime_conversation::prefix_realtime_v2_text; use crate::session::spawn_review_thread; +use crate::state::history as session_history; use codex_config::CloudRequirementsLoader; use codex_config::LoaderOverrides; use codex_config::loader::load_config_layers_state; @@ -891,8 +892,8 @@ pub async fn shutdown(sess: &Arc, sub_id: String) -> bool { sess.guardian_review_session.shutdown().await; info!("Shutting down Codex instance"); let history = sess.clone_history().await; - let turn_count = history - .raw_items() + let history_items = session_history::raw_items(&history); + let turn_count = history_items .iter() .filter(|item| is_user_turn_boundary(item)) .count(); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 415b89c0de..f08c2ed36e 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -171,8 +171,7 @@ use crate::config::Constrained; use crate::config::ConstraintResult; use crate::config::StartedNetworkProxy; use crate::config::resolve_web_search_mode_for_turn; -use crate::context_manager::ContextManager; -use crate::context_manager::TotalTokenUsageBreakdown; +use crate::state::TotalTokenUsageBreakdown; use crate::thread_rollout_truncation::initial_history_has_prior_user_turns; use codex_config::CONFIG_TOML_FILE; use codex_config::types::McpServerConfig; @@ -1084,7 +1083,7 @@ impl Session { pub(crate) async fn get_total_token_usage_breakdown(&self) -> TotalTokenUsageBreakdown { let state = self.state.lock().await; - state.history.get_total_token_usage_breakdown() + state.get_total_token_usage_breakdown() } pub(crate) async fn total_token_usage(&self) -> Option { @@ -1108,7 +1107,7 @@ impl Session { turn_context: &TurnContext, ) -> Option { let state = self.state.lock().await; - state.history.estimate_token_count(turn_context) + state.estimate_token_count(turn_context) } pub(crate) async fn get_base_instructions(&self) -> BaseInstructions { @@ -1430,14 +1429,15 @@ impl Session { }; let shell = self.user_shell(); let exec_policy = self.services.exec_policy.current(); - crate::context_manager::updates::build_settings_update_items( + let entries = crate::context_manager::updates::build_settings_update_entries( reference_context_item, previous_turn_settings.as_ref(), current_context, shell.as_ref(), exec_policy.as_ref(), self.features.enabled(Feature::Personality), - ) + ); + crate::context_manager::updates::render_context_entries(entries) } /// Persist the event to rollout and send it to clients. @@ -2499,16 +2499,24 @@ impl Session { } } - #[expect( - clippy::await_holding_invalid_type, - reason = "MCP app context rendering reads through the session-owned manager guard" - )] pub(crate) async fn build_initial_context( &self, turn_context: &TurnContext, ) -> Vec { - let mut developer_sections = Vec::::with_capacity(8); - let mut contextual_user_sections = Vec::::with_capacity(2); + crate::context_manager::updates::render_context_entries( + self.build_initial_context_entries(turn_context).await, + ) + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "MCP app context rendering reads through the session-owned manager guard" + )] + async fn build_initial_context_entries( + &self, + turn_context: &TurnContext, + ) -> Vec { + let mut entries = Vec::::with_capacity(12); let shell = self.user_shell(); let ( reference_context_item, @@ -2531,26 +2539,36 @@ impl Session { previous_turn_settings.as_ref(), turn_context, ) + && let Some(entry) = crate::context_manager::updates::developer_context_entry( + "model_switch", + 10, + model_switch_message, + ) { - developer_sections.push(model_switch_message); + entries.push(entry); } if turn_context.config.include_permissions_instructions { - developer_sections.push( - PermissionsInstructions::from_permission_profile( - &turn_context.permission_profile, - turn_context.approval_policy.value(), - turn_context.config.approvals_reviewer, - self.services.exec_policy.current().as_ref(), - &turn_context.cwd, - turn_context - .features - .enabled(Feature::ExecPermissionApprovals), - turn_context - .features - .enabled(Feature::RequestPermissionsTool), - ) - .render(), - ); + let permissions_instructions = PermissionsInstructions::from_permission_profile( + &turn_context.permission_profile, + turn_context.approval_policy.value(), + turn_context.config.approvals_reviewer, + self.services.exec_policy.current().as_ref(), + &turn_context.cwd, + turn_context + .features + .enabled(Feature::ExecPermissionApprovals), + turn_context + .features + .enabled(Feature::RequestPermissionsTool), + ) + .render(); + if let Some(entry) = crate::context_manager::updates::developer_context_entry( + "permissions", + 20, + permissions_instructions, + ) { + entries.push(entry); + } } let separate_guardian_developer_message = crate::guardian::is_guardian_reviewer_source(&session_source); @@ -2559,29 +2577,48 @@ impl Session { if !separate_guardian_developer_message && let Some(developer_instructions) = turn_context.developer_instructions.as_deref() && !developer_instructions.is_empty() + && let Some(entry) = crate::context_manager::updates::developer_context_entry( + "developer_instructions", + 30, + developer_instructions.to_string(), + ) { - developer_sections.push(developer_instructions.to_string()); + entries.push(entry); } // Add developer instructions for memories. if turn_context.features.enabled(Feature::MemoryTool) && turn_context.config.memories.use_memories && let Some(memory_prompt) = build_memory_tool_developer_instructions(&turn_context.config.codex_home).await + && let Some(entry) = crate::context_manager::updates::developer_context_entry( + "memory_tool", + 40, + memory_prompt, + ) { - developer_sections.push(memory_prompt); + entries.push(entry); } // Add developer instructions from collaboration_mode if they exist and are non-empty if let Some(collab_instructions) = CollaborationModeInstructions::from_collaboration_mode(&collaboration_mode) + && let Some(entry) = crate::context_manager::updates::developer_context_entry( + "collaboration_mode", + 50, + collab_instructions.render(), + ) { - developer_sections.push(collab_instructions.render()); + entries.push(entry); } if let Some(realtime_update) = crate::context_manager::updates::build_initial_realtime_item( reference_context_item.as_ref(), previous_turn_settings.as_ref(), turn_context, + ) && let Some(entry) = crate::context_manager::updates::developer_context_entry( + "realtime", + 60, + realtime_update, ) { - developer_sections.push(realtime_update); + entries.push(entry); } if self.features.enabled(Feature::Personality) && let Some(personality) = turn_context.personality @@ -2595,9 +2632,13 @@ impl Session { &model_info, personality, ) + && let Some(entry) = crate::context_manager::updates::developer_context_entry( + "personality", + 70, + PersonalitySpecInstructions::new(personality_message).render(), + ) { - developer_sections - .push(PersonalitySpecInstructions::new(personality_message).render()); + entries.push(entry); } } if turn_context.config.include_apps_instructions && turn_context.apps_enabled() { @@ -2610,8 +2651,13 @@ impl Session { .await; if let Some(apps_instructions) = AppsInstructions::from_connectors(&accessible_and_enabled_connectors) + && let Some(entry) = crate::context_manager::updates::developer_context_entry( + "apps", + 80, + apps_instructions.render(), + ) { - developer_sections.push(apps_instructions.render()); + entries.push(entry); } } if turn_context.config.include_skill_instructions { @@ -2634,7 +2680,13 @@ impl Session { }) .await; } - developer_sections.push(skills_instructions.render()); + if let Some(entry) = crate::context_manager::updates::developer_context_entry( + "skills", + 90, + skills_instructions.render(), + ) { + entries.push(entry); + } } } let loaded_plugins = self @@ -2644,24 +2696,38 @@ impl Session { .await; if let Some(plugin_instructions) = AvailablePluginsInstructions::from_plugins(loaded_plugins.capability_summaries()) + && let Some(entry) = crate::context_manager::updates::developer_context_entry( + "plugins", + 100, + plugin_instructions.render(), + ) { - developer_sections.push(plugin_instructions.render()); + entries.push(entry); } if turn_context.features.enabled(Feature::CodexGitCommit) && let Some(commit_message_instruction) = commit_message_trailer_instruction( turn_context.config.commit_attribution.as_deref(), ) + && let Some(entry) = crate::context_manager::updates::developer_context_entry( + "commit_message_trailer", + 110, + commit_message_instruction, + ) { - developer_sections.push(commit_message_instruction); + entries.push(entry); } - if let Some(user_instructions) = turn_context.user_instructions.as_deref() { - contextual_user_sections.push( + if let Some(user_instructions) = turn_context.user_instructions.as_deref() + && let Some(entry) = crate::context_manager::updates::contextual_user_context_entry( + "user_instructions", + 130, UserInstructions { text: user_instructions.to_string(), directory: turn_context.cwd.to_string_lossy().into_owned(), } .render(), - ); + ) + { + entries.push(entry); } if turn_context.config.include_environment_context { let subagents = self @@ -2669,48 +2735,43 @@ impl Session { .agent_control .format_environment_context_subagents(self.conversation_id) .await; - contextual_user_sections.push( + if let Some(entry) = crate::context_manager::updates::contextual_user_context_entry( + "environment", + 140, crate::context::EnvironmentContext::from_turn_context(turn_context, shell.as_ref()) .with_subagents(subagents) .render(), - ); + ) { + entries.push(entry); + } } let multi_agent_v2_usage_hint_text = multi_agents::usage_hint_text(turn_context, &session_source); - let mut items = Vec::with_capacity(4); - if let Some(developer_message) = - crate::context_manager::updates::build_developer_update_item(developer_sections) - { - items.push(developer_message); - } if let Some(usage_hint_text) = multi_agent_v2_usage_hint_text - && let Some(usage_hint_message) = - crate::context_manager::updates::build_developer_update_item(vec![ - usage_hint_text.to_string(), - ]) + && let Some(entry) = crate::context_manager::updates::usage_hint_context_entry( + "multi_agent_v2", + 120, + usage_hint_text.to_string(), + ) { - items.push(usage_hint_message); - } - if let Some(contextual_user_message) = - crate::context_manager::updates::build_contextual_user_message(contextual_user_sections) - { - items.push(contextual_user_message); + entries.push(entry); } // Emit the guardian policy prompt as a separate developer item so the guardian // subagent sees a distinct, easy-to-audit instruction block. if separate_guardian_developer_message && let Some(developer_instructions) = turn_context.developer_instructions.as_deref() && !developer_instructions.is_empty() - && let Some(guardian_developer_message) = - crate::context_manager::updates::build_developer_update_item(vec![ - developer_instructions.to_string(), - ]) + && let Some(entry) = crate::context_manager::updates::guardian_context_entry( + "developer_instructions", + 150, + developer_instructions.to_string(), + ) { - items.push(guardian_developer_message); + entries.push(entry); } - items + entries } pub(crate) async fn persist_rollout_items(&self, items: &[RolloutItem]) { @@ -2721,11 +2782,16 @@ impl Session { } } - pub(crate) async fn clone_history(&self) -> ContextManager { + pub(crate) async fn clone_history(&self) -> codex_journal::Journal { let state = self.state.lock().await; state.clone_history() } + pub(crate) async fn history_version(&self) -> u64 { + let state = self.state.lock().await; + state.history_version() + } + pub(crate) async fn reference_context_item(&self) -> Option { let state = self.state.lock().await; state.reference_context_item() @@ -2792,7 +2858,10 @@ impl Session { let history = self.clone_history().await; let base_instructions = self.get_base_instructions().await; let Some(estimated_total_tokens) = - history.estimate_token_count_with_base_instructions(&base_instructions) + crate::state::history::estimate_token_count_with_base_instructions( + &history, + &base_instructions, + ) else { return; }; diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index a4c042af0c..fd9ae9cf94 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -1,5 +1,7 @@ use super::*; use crate::context_manager::is_user_turn_boundary; +use crate::state::history as session_history; +use codex_journal::Journal; // Return value of `Session::reconstruct_history_from_rollout`, bundling the rebuilt history with // the resume/fork hydration metadata derived from the same replay. @@ -231,10 +233,11 @@ impl Session { ); } - let mut history = ContextManager::new(); + let mut history = Journal::new(); + let mut rebuilt_history_reference_context_item = None; let mut saw_legacy_compaction_without_replacement_history = false; if let Some(base_replacement_history) = base_replacement_history { - history.replace(base_replacement_history.to_vec()); + session_history::replace_history(&mut history, base_replacement_history.to_vec()); } // Materialize exact history semantics from the replay-derived suffix. The eventual lazy // design should keep this same replay shape, but drive it from a resumable reverse source @@ -242,7 +245,8 @@ impl Session { for item in rollout_suffix { match item { RolloutItem::ResponseItem(response_item) => { - history.record_items( + session_history::record_items( + &mut history, std::iter::once(response_item), turn_context.truncation_policy, ); @@ -251,7 +255,7 @@ impl Session { if let Some(replacement_history) = &compacted.replacement_history { // This should actually never happen, because the reverse loop above (to build rollout_suffix) // should stop before any compaction that has Some replacement_history - history.replace(replacement_history.clone()); + session_history::replace_history(&mut history, replacement_history.clone()); } else { saw_legacy_compaction_without_replacement_history = true; // Legacy rollouts without `replacement_history` should rebuild the @@ -262,17 +266,22 @@ impl Session { // prompt shape. // TODO(ccunningham): if we drop support for None replacement_history compaction items, // we can get rid of this second loop entirely and just build `history` directly in the first loop. - let user_messages = collect_user_messages(history.raw_items()); + let history_items = session_history::raw_items(&history); + let user_messages = collect_user_messages(history_items.as_slice()); let rebuilt = compact::build_compacted_history( Vec::new(), &user_messages, &compacted.message, ); - history.replace(rebuilt); + session_history::replace_history(&mut history, rebuilt); } } RolloutItem::EventMsg(EventMsg::ThreadRolledBack(rollback)) => { - history.drop_last_n_user_turns(rollback.num_turns); + session_history::drop_last_n_user_turns( + &mut history, + &mut rebuilt_history_reference_context_item, + rollback.num_turns, + ); } RolloutItem::EventMsg(_) | RolloutItem::TurnContext(_) @@ -293,7 +302,7 @@ impl Session { }; RolloutReconstruction { - history: history.raw_items().to_vec(), + history: session_history::raw_items(&history), previous_turn_settings, reference_context_item, } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 3062f8f6f4..401a6bfe5c 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -9,6 +9,7 @@ use crate::function_tool::FunctionCallError; use crate::shell::default_user_shell; use crate::skills::SkillRenderSideEffects; use crate::skills::render::SkillMetadataBudget; +use crate::state::history as session_history; use crate::test_support::models_manager_with_provider; use crate::tools::format_exec_output_str; use codex_config::ConfigLayerStack; @@ -20,6 +21,7 @@ use codex_config::RequirementSource; use codex_config::Sourced; use codex_config::loader::project_trust_key; use codex_config::types::ToolSuggestDisabledTool; +use codex_journal::Journal; use codex_features::Feature; use codex_features::Features; @@ -78,6 +80,63 @@ use codex_execpolicy::Policy; use codex_network_proxy::NetworkProxyConfig; use codex_otel::MetricsClient; use codex_otel::MetricsConfig; + +trait JournalTestExt { + fn raw_items(&self) -> Vec; + fn for_prompt( + &self, + input_modalities: &[codex_protocol::openai_models::InputModality], + ) -> Vec; + fn estimate_token_count(&self, turn_context: &TurnContext) -> Option; + fn estimate_token_count_with_base_instructions( + &self, + base_instructions: &BaseInstructions, + ) -> Option; + fn record_items( + &mut self, + items: I, + policy: codex_utils_output_truncation::TruncationPolicy, + ) where + I: IntoIterator, + I::Item: std::ops::Deref; + fn replace(&mut self, items: Vec); +} + +impl JournalTestExt for Journal { + fn raw_items(&self) -> Vec { + session_history::raw_items(self) + } + + fn for_prompt( + &self, + input_modalities: &[codex_protocol::openai_models::InputModality], + ) -> Vec { + session_history::for_prompt(self, input_modalities) + } + + fn estimate_token_count(&self, turn_context: &TurnContext) -> Option { + session_history::estimate_token_count(self, turn_context) + } + + fn estimate_token_count_with_base_instructions( + &self, + base_instructions: &BaseInstructions, + ) -> Option { + session_history::estimate_token_count_with_base_instructions(self, base_instructions) + } + + fn record_items(&mut self, items: I, policy: codex_utils_output_truncation::TruncationPolicy) + where + I: IntoIterator, + I::Item: std::ops::Deref, + { + session_history::record_items(self, items, policy); + } + + fn replace(&mut self, items: Vec) { + session_history::replace_history(self, items); + } +} use codex_otel::THREAD_SKILLS_DESCRIPTION_TRUNCATED_CHARS_METRIC; use codex_otel::THREAD_SKILLS_ENABLED_TOTAL_METRIC; use codex_otel::THREAD_SKILLS_KEPT_TOTAL_METRIC; @@ -7629,7 +7688,7 @@ async fn sample_rollout( _turn_context: &TurnContext, ) -> (Vec, Vec) { let mut rollout_items = Vec::new(); - let mut live_history = ContextManager::new(); + let mut live_history = Journal::new(); // Use the same turn_context source as record_initial_history so model_info (and thus // personality_spec) matches reconstruction. diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 823dcbfc3d..174b4d5f97 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -42,6 +42,7 @@ use crate::resolve_skill_dependencies_for_turn; use crate::session::PreviousTurnSettings; use crate::session::session::Session; use crate::session::turn_context::TurnContext; +use crate::state::history as session_history; use crate::stream_events_utils::HandleOutputCtx; use crate::stream_events_utils::handle_non_tool_response_item; use crate::stream_events_utils::handle_output_item_done; @@ -428,9 +429,8 @@ pub(crate) async fn run_turn( // Construct the input that we will send to the model. let sampling_request_input: Vec = { - sess.clone_history() - .await - .for_prompt(&turn_context.model_info.input_modalities) + let history = sess.clone_history().await; + session_history::for_prompt(&history, &turn_context.model_info.input_modalities) }; let sampling_request_input_messages = sampling_request_input @@ -629,7 +629,7 @@ pub(crate) async fn run_turn( error_or_panic( "Invalid image detected; sanitizing tool output to prevent poisoning", ); - if state.history.replace_last_turn_images("Invalid image") { + if state.replace_last_turn_images("Invalid image") { continue; } } @@ -1005,9 +1005,8 @@ async fn run_sampling_request( let prompt_input = if let Some(input) = initial_input.take() { input } else { - sess.clone_history() - .await - .for_prompt(&turn_context.model_info.input_modalities) + let history = sess.clone_history().await; + session_history::for_prompt(&history, &turn_context.model_info.input_modalities) }; let prompt = build_prompt( prompt_input, diff --git a/codex-rs/core/src/state/history.rs b/codex-rs/core/src/state/history.rs new file mode 100644 index 0000000000..0daa1a3e54 --- /dev/null +++ b/codex-rs/core/src/state/history.rs @@ -0,0 +1,314 @@ +use crate::event_mapping::has_non_contextual_dev_message_content; +use crate::event_mapping::is_contextual_dev_message_content; +use crate::event_mapping::is_contextual_user_message_content; +use crate::session::turn_context::TurnContext; +use codex_journal::Journal; +use codex_journal::JournalHistoryItem; +use codex_journal::JournalItem; +use codex_journal::JournalKey; +use codex_journal::history as thread_history; +use codex_journal::history::estimate_item_token_count; +use codex_journal::history::estimate_response_item_model_visible_bytes; +use codex_journal::history::is_api_message; +use codex_journal::history::is_model_generated_item; +use codex_journal::history::is_user_turn_boundary; +use codex_journal::history::user_turn_boundary_positions; +use codex_protocol::models::BaseInstructions; +use codex_protocol::models::ResponseItem; +use codex_protocol::openai_models::InputModality; +use codex_protocol::protocol::TokenUsage; +use codex_protocol::protocol::TokenUsageInfo; +use codex_protocol::protocol::TurnContextItem; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::approx_token_count; +use std::ops::Deref; + +#[derive(Debug, Clone, Copy, Default)] +pub(crate) struct TotalTokenUsageBreakdown { + pub(crate) last_api_response_total_tokens: i64, + pub(crate) all_history_items_model_visible_bytes: i64, + pub(crate) estimated_tokens_of_items_added_since_last_successful_api_response: i64, + pub(crate) estimated_bytes_of_items_added_since_last_successful_api_response: i64, +} + +pub(crate) fn record_items(journal: &mut Journal, items: I, policy: TruncationPolicy) +where + I: IntoIterator, + I::Item: Deref, +{ + for item in items { + let item_ref = item.deref(); + if !is_api_message(item_ref) { + continue; + } + + let processed = thread_history::truncate_history_item(item_ref, policy); + push_history_item(journal, processed); + } +} + +pub(crate) fn raw_items(journal: &Journal) -> Vec { + journal + .entries() + .iter() + .filter_map(|entry| match &entry.item { + JournalItem::History(item) => Some(item.item.clone()), + JournalItem::Context(_) | JournalItem::Checkpoint(_) => None, + }) + .collect() +} + +pub(crate) fn for_prompt( + journal: &Journal, + input_modalities: &[InputModality], +) -> Vec { + let mut items = raw_items(journal); + normalize_history(&mut items, input_modalities); + items +} + +pub(crate) fn estimate_token_count(journal: &Journal, turn_context: &TurnContext) -> Option { + let model_info = &turn_context.model_info; + let personality = turn_context.personality.or(turn_context.config.personality); + let base_instructions = BaseInstructions { + text: model_info.get_model_instructions(personality), + }; + estimate_token_count_with_base_instructions(journal, &base_instructions) +} + +pub(crate) fn estimate_token_count_with_base_instructions( + journal: &Journal, + base_instructions: &BaseInstructions, +) -> Option { + let base_tokens = + i64::try_from(approx_token_count(&base_instructions.text)).unwrap_or(i64::MAX); + + let items_tokens = raw_items(journal) + .iter() + .map(estimate_item_token_count) + .fold(0i64, i64::saturating_add); + + Some(base_tokens.saturating_add(items_tokens)) +} + +pub(crate) fn remove_first_item(journal: &mut Journal) -> bool { + let mut items = raw_items(journal); + if items.is_empty() { + return false; + } + + let removed = items.remove(0); + thread_history::remove_corresponding_for(&mut items, &removed); + replace_history(journal, items); + true +} + +pub(crate) fn remove_last_item(journal: &mut Journal) -> bool { + let mut items = raw_items(journal); + if let Some(removed) = items.pop() { + thread_history::remove_corresponding_for(&mut items, &removed); + replace_history(journal, items); + true + } else { + false + } +} + +pub(crate) fn replace_history(journal: &mut Journal, items: Vec) { + *journal = journal_from_items(items); +} + +pub(crate) fn replace_last_turn_images(journal: &mut Journal, placeholder: &str) -> bool { + let mut items = raw_items(journal); + let replaced = thread_history::replace_last_turn_images(&mut items, placeholder); + if replaced { + replace_history(journal, items); + } + replaced +} + +pub(crate) fn drop_last_n_user_turns( + journal: &mut Journal, + reference_context_item: &mut Option, + num_turns: u32, +) { + if num_turns == 0 { + return; + } + + let snapshot = raw_items(journal); + let user_positions = user_turn_boundary_positions(&snapshot); + let Some(&first_instruction_turn_idx) = user_positions.first() else { + replace_history(journal, snapshot); + return; + }; + + let n_from_end = usize::try_from(num_turns).unwrap_or(usize::MAX); + let mut cut_idx = if n_from_end >= user_positions.len() { + first_instruction_turn_idx + } else { + user_positions[user_positions.len() - n_from_end] + }; + + cut_idx = trim_pre_turn_context_updates( + reference_context_item, + &snapshot, + first_instruction_turn_idx, + cut_idx, + ); + + replace_history(journal, snapshot[..cut_idx].to_vec()); +} + +pub(crate) fn update_token_info( + token_info: &mut Option, + usage: &TokenUsage, + model_context_window: Option, +) { + *token_info = + TokenUsageInfo::new_or_append(token_info, &Some(usage.clone()), model_context_window); +} + +pub(crate) fn set_token_usage_full(token_info: &mut Option, context_window: i64) { + match token_info { + Some(info) => info.fill_to_context_window(context_window), + None => { + *token_info = Some(TokenUsageInfo::full_context_window(context_window)); + } + } +} + +pub(crate) fn get_total_token_usage( + journal: &Journal, + token_info: &Option, + server_reasoning_included: bool, +) -> i64 { + let items = raw_items(journal); + let last_tokens = token_info + .as_ref() + .map(|info| info.last_token_usage.total_tokens) + .unwrap_or(0); + let items_after_last_model_generated_tokens = items_after_last_model_generated_item(&items) + .iter() + .map(estimate_item_token_count) + .fold(0i64, i64::saturating_add); + if server_reasoning_included { + last_tokens.saturating_add(items_after_last_model_generated_tokens) + } else { + last_tokens + .saturating_add(get_non_last_reasoning_items_tokens(&items)) + .saturating_add(items_after_last_model_generated_tokens) + } +} + +pub(crate) fn get_total_token_usage_breakdown( + journal: &Journal, + token_info: &Option, +) -> TotalTokenUsageBreakdown { + let items = raw_items(journal); + let last_usage = token_info + .as_ref() + .map(|info| info.last_token_usage.clone()) + .unwrap_or_default(); + let items_after_last_model_generated = items_after_last_model_generated_item(&items); + + TotalTokenUsageBreakdown { + last_api_response_total_tokens: last_usage.total_tokens, + all_history_items_model_visible_bytes: items + .iter() + .map(estimate_response_item_model_visible_bytes) + .fold(0i64, i64::saturating_add), + estimated_tokens_of_items_added_since_last_successful_api_response: + items_after_last_model_generated + .iter() + .map(estimate_item_token_count) + .fold(0i64, i64::saturating_add), + estimated_bytes_of_items_added_since_last_successful_api_response: + items_after_last_model_generated + .iter() + .map(estimate_response_item_model_visible_bytes) + .fold(0i64, i64::saturating_add), + } +} + +fn push_history_item(journal: &mut Journal, item: ResponseItem) { + let history_item = JournalHistoryItem::new(item); + let key = JournalKey::new(vec!["history".to_string(), history_item.id.clone()]); + journal.add(key, history_item); +} + +fn journal_from_items(items: Vec) -> Journal { + let mut journal = Journal::new(); + for item in items { + if is_api_message(&item) { + push_history_item(&mut journal, item); + } + } + journal +} + +fn get_non_last_reasoning_items_tokens(items: &[ResponseItem]) -> i64 { + let Some(last_user_index) = items.iter().rposition(is_user_turn_boundary) else { + return 0; + }; + + items + .iter() + .take(last_user_index) + .filter(|item| { + matches!( + item, + ResponseItem::Reasoning { + encrypted_content: Some(_), + .. + } + ) + }) + .map(estimate_item_token_count) + .fold(0i64, i64::saturating_add) +} + +fn items_after_last_model_generated_item(items: &[ResponseItem]) -> &[ResponseItem] { + let start = items + .iter() + .rposition(is_model_generated_item) + .map_or(items.len(), |index| index.saturating_add(1)); + &items[start..] +} + +fn normalize_history(items: &mut Vec, input_modalities: &[InputModality]) { + thread_history::ensure_call_outputs_present(items); + thread_history::remove_orphan_outputs(items); + thread_history::strip_images_when_unsupported(input_modalities, items); +} + +fn trim_pre_turn_context_updates( + reference_context_item: &mut Option, + snapshot: &[ResponseItem], + first_instruction_turn_idx: usize, + mut cut_idx: usize, +) -> usize { + while cut_idx > first_instruction_turn_idx { + match &snapshot[cut_idx - 1] { + ResponseItem::Message { role, content, .. } + if role == "developer" && is_contextual_dev_message_content(content) => + { + if has_non_contextual_dev_message_content(content) { + *reference_context_item = None; + } + cut_idx -= 1; + } + ResponseItem::Message { role, content, .. } + if role == "user" && is_contextual_user_message_content(content) => + { + cut_idx -= 1; + } + _ => break, + } + } + cut_idx +} + +#[cfg(test)] +#[path = "history_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/state/history_tests.rs similarity index 94% rename from codex-rs/core/src/context_manager/history_tests.rs rename to codex-rs/core/src/state/history_tests.rs index 74f4d29bfb..9d511cb361 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/state/history_tests.rs @@ -1,6 +1,9 @@ use super::*; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use codex_journal::Journal; +use codex_journal::history::ORIGINAL_IMAGE_MAX_PATCHES; +use codex_journal::history::RESIZED_IMAGE_BYTES_ESTIMATE; use codex_protocol::AgentPath; use codex_protocol::config_types::ReasoningSummary; use codex_protocol::models::BaseInstructions; @@ -20,8 +23,11 @@ use codex_protocol::openai_models::default_input_modalities; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::SandboxPolicy; +use codex_protocol::protocol::TokenUsage; +use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::TurnContextItem; use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::approx_bytes_for_tokens; use codex_utils_output_truncation::truncate_text; use image::ImageBuffer; use image::ImageFormat; @@ -34,6 +40,98 @@ use std::path::PathBuf; const EXEC_FORMAT_MAX_BYTES: usize = 10_000; const EXEC_FORMAT_MAX_TOKENS: usize = 2_500; +#[derive(Clone, Default)] +struct TestHistory { + journal: Journal, + token_info: Option, + reference_context_item: Option, +} + +impl TestHistory { + fn new() -> Self { + Self { + journal: Journal::new(), + token_info: TokenUsageInfo::new_or_append( + &None, &None, /*model_context_window*/ None, + ), + reference_context_item: None, + } + } + + fn record_items(&mut self, items: I, policy: TruncationPolicy) + where + I: IntoIterator, + I::Item: std::ops::Deref, + { + super::record_items(&mut self.journal, items, policy); + } + + fn raw_items(&self) -> Vec { + super::raw_items(&self.journal) + } + + fn for_prompt(&self, input_modalities: &[InputModality]) -> Vec { + super::for_prompt(&self.journal, input_modalities) + } + + fn estimate_token_count_with_base_instructions( + &self, + base_instructions: &BaseInstructions, + ) -> Option { + super::estimate_token_count_with_base_instructions(&self.journal, base_instructions) + } + + fn remove_first_item(&mut self) { + super::remove_first_item(&mut self.journal); + } + + fn remove_last_item(&mut self) -> bool { + super::remove_last_item(&mut self.journal) + } + + fn replace_last_turn_images(&mut self, placeholder: &str) -> bool { + super::replace_last_turn_images(&mut self.journal, placeholder) + } + + fn drop_last_n_user_turns(&mut self, num_turns: u32) { + super::drop_last_n_user_turns( + &mut self.journal, + &mut self.reference_context_item, + num_turns, + ); + } + + fn update_token_info(&mut self, usage: &TokenUsage, model_context_window: Option) { + super::update_token_info(&mut self.token_info, usage, model_context_window); + } + + fn get_total_token_usage(&self, server_reasoning_included: bool) -> i64 { + super::get_total_token_usage(&self.journal, &self.token_info, server_reasoning_included) + } + + fn set_reference_context_item(&mut self, item: Option) { + self.reference_context_item = item; + } + + fn reference_context_item(&self) -> Option { + self.reference_context_item.clone() + } + + fn normalize_history(&mut self, input_modalities: &[InputModality]) { + let mut items = self.raw_items(); + super::normalize_history(&mut items, input_modalities); + super::replace_history(&mut self.journal, items); + } + + fn get_non_last_reasoning_items_tokens(&self) -> i64 { + super::get_non_last_reasoning_items_tokens(&self.raw_items()) + } + + fn items_after_last_model_generated_item(&self) -> Vec { + super::items_after_last_model_generated_item(&self.raw_items()).to_vec() + } +} + fn assistant_msg(text: &str) -> ResponseItem { ResponseItem::Message { id: None, @@ -63,8 +161,8 @@ fn inter_agent_assistant_msg(text: &str) -> ResponseItem { } } -fn create_history_with_items(items: Vec) -> ContextManager { - let mut h = ContextManager::new(); +fn create_history_with_items(items: Vec) -> TestHistory { + let mut h = TestHistory::new(); // Use a generous but fixed token budget; tests only rely on truncation // behavior, not on a specific model's token limit. h.record_items(items.iter(), TruncationPolicy::Tokens(10_000)); @@ -185,7 +283,7 @@ fn approx_token_count_for_text(text: &str) -> i64 { #[test] fn filters_non_api_messages() { - let mut h = ContextManager::default(); + let mut h = TestHistory::default(); let policy = TruncationPolicy::Tokens(10_000); // System message is not API messages; Other is ignored. let system = ResponseItem::Message { @@ -327,7 +425,7 @@ fn drop_last_n_user_turns_treats_inter_agent_assistant_messages_as_instruction_t history.drop_last_n_user_turns(/*num_turns*/ 1); - assert_eq!(history.raw_items(), &vec![first_turn, first_reply]); + assert_eq!(history.raw_items(), vec![first_turn, first_reply]); } #[test] @@ -1014,7 +1112,7 @@ fn normalization_retains_local_shell_outputs() { #[test] fn record_items_truncates_function_call_output_content() { - let mut history = ContextManager::new(); + let mut history = TestHistory::new(); // Any reasonably small token budget works; the test only cares that // truncation happens and the marker is present. let policy = TruncationPolicy::Tokens(1_000); @@ -1030,8 +1128,9 @@ fn record_items_truncates_function_call_output_content() { history.record_items([&item], policy); - assert_eq!(history.items.len(), 1); - match &history.items[0] { + let history_items = history.raw_items(); + assert_eq!(history_items.len(), 1); + match &history_items[0] { ResponseItem::FunctionCallOutput { output, .. } => { let content = output.text_content().unwrap_or_default(); assert_ne!(content, long_output); @@ -1050,7 +1149,7 @@ fn record_items_truncates_function_call_output_content() { #[test] fn record_items_truncates_custom_tool_call_output_content() { - let mut history = ContextManager::new(); + let mut history = TestHistory::new(); let policy = TruncationPolicy::Tokens(1_000); let line = "custom output that is very long\n"; let long_output = line.repeat(2_500); @@ -1062,8 +1161,9 @@ fn record_items_truncates_custom_tool_call_output_content() { history.record_items([&item], policy); - assert_eq!(history.items.len(), 1); - match &history.items[0] { + let history_items = history.raw_items(); + assert_eq!(history_items.len(), 1); + match &history_items[0] { ResponseItem::CustomToolCallOutput { output, .. } => { let output = output.text_content().unwrap_or_default(); assert_ne!(output, long_output); @@ -1082,7 +1182,7 @@ fn record_items_truncates_custom_tool_call_output_content() { #[test] fn record_items_respects_custom_token_limit() { - let mut history = ContextManager::new(); + let mut history = TestHistory::new(); let policy = TruncationPolicy::Tokens(10); let long_output = "tokenized content repeated many times ".repeat(200); let item = ResponseItem::FunctionCallOutput { @@ -1095,7 +1195,8 @@ fn record_items_respects_custom_token_limit() { history.record_items([&item], policy); - let stored = match &history.items[0] { + let history_items = history.raw_items(); + let stored = match &history_items[0] { ResponseItem::FunctionCallOutput { output, .. } => output, other => panic!("unexpected history item: {other:?}"), }; diff --git a/codex-rs/core/src/state/mod.rs b/codex-rs/core/src/state/mod.rs index 13f3bf6c86..c9dd947319 100644 --- a/codex-rs/core/src/state/mod.rs +++ b/codex-rs/core/src/state/mod.rs @@ -1,7 +1,9 @@ +pub(crate) mod history; mod service; mod session; mod turn; +pub(crate) use history::TotalTokenUsageBreakdown; pub(crate) use service::SessionServices; pub(crate) use session::SessionState; pub(crate) use turn::ActiveTurn; diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 3bd4b8a26e..a5b4fba765 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -1,15 +1,17 @@ //! Session-wide mutable state. +use codex_journal::Journal; use codex_protocol::models::AdditionalPermissionProfile; use codex_protocol::models::ResponseItem; use codex_sandboxing::policy_transforms::merge_permission_profiles; use std::collections::HashMap; use std::collections::HashSet; -use crate::context_manager::ContextManager; use crate::session::PreviousTurnSettings; use crate::session::session::SessionConfiguration; use crate::session_startup_prewarm::SessionStartupPrewarmHandle; +use crate::state::history as session_history; +use crate::state::history::TotalTokenUsageBreakdown; use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; @@ -19,7 +21,10 @@ use codex_utils_output_truncation::TruncationPolicy; /// Persistent, session-scoped state previously stored directly on `Session`. pub(crate) struct SessionState { pub(crate) session_configuration: SessionConfiguration, - pub(crate) history: ContextManager, + journal: Journal, + history_version: u64, + token_info: Option, + reference_context_item: Option, pub(crate) latest_rate_limits: Option, pub(crate) server_reasoning_included: bool, pub(crate) dependency_env: HashMap, @@ -39,10 +44,14 @@ pub(crate) struct SessionState { impl SessionState { /// Create a new session state mirroring previous `State::default()` semantics. pub(crate) fn new(session_configuration: SessionConfiguration) -> Self { - let history = ContextManager::new(); Self { session_configuration, - history, + journal: Journal::new(), + history_version: 0, + token_info: TokenUsageInfo::new_or_append( + &None, &None, /*model_context_window*/ None, + ), + reference_context_item: None, latest_rate_limits: None, server_reasoning_included: false, dependency_env: HashMap::new(), @@ -62,7 +71,7 @@ impl SessionState { I: IntoIterator, I::Item: std::ops::Deref, { - self.history.record_items(items, policy); + session_history::record_items(&mut self.journal, items, policy); } pub(crate) fn previous_turn_settings(&self) -> Option { @@ -85,8 +94,8 @@ impl SessionState { is_first_turn } - pub(crate) fn clone_history(&self) -> ContextManager { - self.history.clone() + pub(crate) fn clone_history(&self) -> Journal { + self.journal.clone() } pub(crate) fn replace_history( @@ -94,21 +103,21 @@ impl SessionState { items: Vec, reference_context_item: Option, ) { - self.history.replace(items); - self.history - .set_reference_context_item(reference_context_item); + session_history::replace_history(&mut self.journal, items); + self.history_version = self.history_version.saturating_add(1); + self.reference_context_item = reference_context_item; } pub(crate) fn set_token_info(&mut self, info: Option) { - self.history.set_token_info(info); + self.token_info = info; } pub(crate) fn set_reference_context_item(&mut self, item: Option) { - self.history.set_reference_context_item(item); + self.reference_context_item = item; } pub(crate) fn reference_context_item(&self) -> Option { - self.history.reference_context_item() + self.reference_context_item.clone() } // Token/rate limit helpers @@ -117,11 +126,11 @@ impl SessionState { usage: &TokenUsage, model_context_window: Option, ) { - self.history.update_token_info(usage, model_context_window); + session_history::update_token_info(&mut self.token_info, usage, model_context_window); } pub(crate) fn token_info(&self) -> Option { - self.history.token_info() + self.token_info.clone() } pub(crate) fn set_rate_limits(&mut self, snapshot: RateLimitSnapshot) { @@ -138,12 +147,38 @@ impl SessionState { } pub(crate) fn set_token_usage_full(&mut self, context_window: i64) { - self.history.set_token_usage_full(context_window); + session_history::set_token_usage_full(&mut self.token_info, context_window); } pub(crate) fn get_total_token_usage(&self, server_reasoning_included: bool) -> i64 { - self.history - .get_total_token_usage(server_reasoning_included) + session_history::get_total_token_usage( + &self.journal, + &self.token_info, + server_reasoning_included, + ) + } + + pub(crate) fn get_total_token_usage_breakdown(&self) -> TotalTokenUsageBreakdown { + session_history::get_total_token_usage_breakdown(&self.journal, &self.token_info) + } + + pub(crate) fn estimate_token_count( + &self, + turn_context: &crate::session::turn_context::TurnContext, + ) -> Option { + session_history::estimate_token_count(&self.journal, turn_context) + } + + pub(crate) fn history_version(&self) -> u64 { + self.history_version + } + + pub(crate) fn replace_last_turn_images(&mut self, placeholder: &str) -> bool { + let replaced = session_history::replace_last_turn_images(&mut self.journal, placeholder); + if replaced { + self.history_version = self.history_version.saturating_add(1); + } + replaced } pub(crate) fn set_server_reasoning_included(&mut self, included: bool) { diff --git a/codex-rs/journal/BUILD.bazel b/codex-rs/journal/BUILD.bazel new file mode 100644 index 0000000000..d12c8e200b --- /dev/null +++ b/codex-rs/journal/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "journal", + crate_name = "codex_journal", +) diff --git a/codex-rs/journal/Cargo.toml b/codex-rs/journal/Cargo.toml new file mode 100644 index 0000000000..14aab7d67a --- /dev/null +++ b/codex-rs/journal/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "codex-journal" +version.workspace = true +edition.workspace = true +license.workspace = true + +[lints] +workspace = true + +[dependencies] +base64 = { workspace = true } +codex-utils-cache = { workspace = true } +codex-utils-output-truncation = { workspace = true } +codex-protocol = { workspace = true } +image = { workspace = true, features = ["jpeg", "png", "webp"] } +indexmap = { workspace = true } +serde_json = { workspace = true } +thiserror = { workspace = true } +tracing = { workspace = true } + +[dev-dependencies] +pretty_assertions = { workspace = true } +tempfile = { workspace = true } diff --git a/codex-rs/journal/src/error.rs b/codex-rs/journal/src/error.rs new file mode 100644 index 0000000000..f907d66a04 --- /dev/null +++ b/codex-rs/journal/src/error.rs @@ -0,0 +1,24 @@ +use thiserror::Error; + +/// Convenience result type for journal operations. +pub type Result = std::result::Result; + +/// Errors produced while materializing or persisting a journal. +#[derive(Debug, Error)] +pub enum JournalError { + #[error("failed to read or write journal")] + Io(#[from] std::io::Error), + #[error("failed to serialize journal item")] + SerializeJson { + #[from] + source: serde_json::Error, + }, + #[error("failed to parse journal item at line {line_number}")] + ParseJson { + line_number: usize, + #[source] + source: serde_json::Error, + }, + #[error("history cursor referenced unknown history item id `{history_item_id}`")] + UnknownHistoryItemId { history_item_id: String }, +} diff --git a/codex-rs/journal/src/history/classify.rs b/codex-rs/journal/src/history/classify.rs new file mode 100644 index 0000000000..e3fd93bfe0 --- /dev/null +++ b/codex-rs/journal/src/history/classify.rs @@ -0,0 +1,133 @@ +use codex_protocol::items::parse_hook_prompt_fragment; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::ENVIRONMENT_CONTEXT_CLOSE_TAG; +use codex_protocol::protocol::ENVIRONMENT_CONTEXT_OPEN_TAG; +use codex_protocol::protocol::InterAgentCommunication; + +const USER_INSTRUCTIONS_START_MARKER: &str = "# AGENTS.md instructions for "; +const USER_INSTRUCTIONS_END_MARKER: &str = ""; +const SKILL_START_MARKER: &str = ""; +const SKILL_END_MARKER: &str = ""; +const USER_SHELL_COMMAND_START_MARKER: &str = ""; +const USER_SHELL_COMMAND_END_MARKER: &str = ""; +const TURN_ABORTED_START_MARKER: &str = ""; +const TURN_ABORTED_END_MARKER: &str = ""; +const SUBAGENT_NOTIFICATION_START_MARKER: &str = ""; +const SUBAGENT_NOTIFICATION_END_MARKER: &str = ""; + +/// Returns whether an item should be carried in API-visible conversation history. +pub fn is_api_message(message: &ResponseItem) -> bool { + match message { + ResponseItem::Message { role, .. } => role.as_str() != "system", + ResponseItem::FunctionCallOutput { .. } + | ResponseItem::FunctionCall { .. } + | ResponseItem::ToolSearchCall { .. } + | ResponseItem::ToolSearchOutput { .. } + | ResponseItem::CustomToolCall { .. } + | ResponseItem::CustomToolCallOutput { .. } + | ResponseItem::LocalShellCall { .. } + | ResponseItem::Reasoning { .. } + | ResponseItem::WebSearchCall { .. } + | ResponseItem::ImageGenerationCall { .. } + | ResponseItem::Compaction { .. } => true, + ResponseItem::Other => false, + } +} + +/// Returns whether an item originated from model generation rather than client-side bookkeeping. +pub fn is_model_generated_item(item: &ResponseItem) -> bool { + match item { + ResponseItem::Message { role, .. } => role == "assistant", + ResponseItem::Reasoning { .. } + | ResponseItem::FunctionCall { .. } + | ResponseItem::ToolSearchCall { .. } + | ResponseItem::WebSearchCall { .. } + | ResponseItem::ImageGenerationCall { .. } + | ResponseItem::CustomToolCall { .. } + | ResponseItem::LocalShellCall { .. } + | ResponseItem::Compaction { .. } => true, + ResponseItem::FunctionCallOutput { .. } + | ResponseItem::ToolSearchOutput { .. } + | ResponseItem::CustomToolCallOutput { .. } + | ResponseItem::Other => false, + } +} + +/// Returns whether an item was injected by Codex rather than supplied by the user or model. +pub fn is_codex_generated_item(item: &ResponseItem) -> bool { + matches!( + item, + ResponseItem::FunctionCallOutput { .. } + | ResponseItem::ToolSearchOutput { .. } + | ResponseItem::CustomToolCallOutput { .. } + ) || matches!(item, ResponseItem::Message { role, .. } if role == "developer") +} + +/// Returns whether an item should count as an instruction-turn boundary for history rewrites. +pub fn is_user_turn_boundary(item: &ResponseItem) -> bool { + let ResponseItem::Message { role, content, .. } = item else { + return false; + }; + + (role == "user" && !is_contextual_user_message_content(content)) + || (role == "assistant" && is_inter_agent_instruction_content(content)) +} + +/// Returns the indexes of all instruction-turn boundaries in order. +pub fn user_turn_boundary_positions(items: &[ResponseItem]) -> Vec { + let mut positions = Vec::new(); + for (index, item) in items.iter().enumerate() { + if is_user_turn_boundary(item) { + positions.push(index); + } + } + positions +} + +fn is_inter_agent_instruction_content(content: &[ContentItem]) -> bool { + InterAgentCommunication::is_message_content(content) +} + +fn is_contextual_user_message_content(content: &[ContentItem]) -> bool { + content.iter().any(|item| match item { + ContentItem::InputText { text } => { + parse_hook_prompt_fragment(text).is_some() + || matches_fragment( + text, + USER_INSTRUCTIONS_START_MARKER, + USER_INSTRUCTIONS_END_MARKER, + ) + || matches_fragment( + text, + ENVIRONMENT_CONTEXT_OPEN_TAG, + ENVIRONMENT_CONTEXT_CLOSE_TAG, + ) + || matches_fragment(text, SKILL_START_MARKER, SKILL_END_MARKER) + || matches_fragment( + text, + USER_SHELL_COMMAND_START_MARKER, + USER_SHELL_COMMAND_END_MARKER, + ) + || matches_fragment(text, TURN_ABORTED_START_MARKER, TURN_ABORTED_END_MARKER) + || matches_fragment( + text, + SUBAGENT_NOTIFICATION_START_MARKER, + SUBAGENT_NOTIFICATION_END_MARKER, + ) + } + ContentItem::InputImage { .. } | ContentItem::OutputText { .. } => false, + }) +} + +fn matches_fragment(text: &str, start_marker: &str, end_marker: &str) -> bool { + let trimmed = text.trim_start(); + let starts_with_marker = trimmed + .get(..start_marker.len()) + .is_some_and(|candidate| candidate.eq_ignore_ascii_case(start_marker)); + let trimmed = trimmed.trim_end(); + let ends_with_marker = trimmed + .get(trimmed.len().saturating_sub(end_marker.len())..) + .is_some_and(|candidate| candidate.eq_ignore_ascii_case(end_marker)); + starts_with_marker && ends_with_marker +} diff --git a/codex-rs/journal/src/history/estimate.rs b/codex-rs/journal/src/history/estimate.rs new file mode 100644 index 0000000000..c1c0f6ee6b --- /dev/null +++ b/codex-rs/journal/src/history/estimate.rs @@ -0,0 +1,174 @@ +use base64::Engine; +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use codex_protocol::models::ContentItem; +use codex_protocol::models::FunctionCallOutputBody; +use codex_protocol::models::FunctionCallOutputContentItem; +use codex_protocol::models::ImageDetail; +use codex_protocol::models::ResponseItem; +use codex_utils_cache::BlockingLruCache; +use codex_utils_cache::sha1_digest; +use codex_utils_output_truncation::approx_bytes_for_tokens; +use codex_utils_output_truncation::approx_tokens_from_byte_count_i64; +use std::num::NonZeroUsize; +use std::sync::LazyLock; + +fn estimate_reasoning_length(encoded_len: usize) -> usize { + encoded_len + .saturating_mul(3) + .checked_div(4) + .unwrap_or(0) + .saturating_sub(650) +} + +/// Approximates the token cost of one history item using model-visible byte heuristics. +pub fn estimate_item_token_count(item: &ResponseItem) -> i64 { + let model_visible_bytes = estimate_response_item_model_visible_bytes(item); + approx_tokens_from_byte_count_i64(model_visible_bytes) +} + +/// Approximates the model-visible byte size of one history item. +/// +/// Inline base64 image payloads are discounted to a fixed vision-token estimate instead of their +/// raw serialized size, while encrypted reasoning content uses a coarse decoded-length heuristic. +pub fn estimate_response_item_model_visible_bytes(item: &ResponseItem) -> i64 { + match item { + ResponseItem::Reasoning { + encrypted_content: Some(content), + .. + } + | ResponseItem::Compaction { + encrypted_content: content, + } => i64::try_from(estimate_reasoning_length(content.len())).unwrap_or(i64::MAX), + item => { + let raw = serde_json::to_string(item) + .map(|serialized| i64::try_from(serialized.len()).unwrap_or(i64::MAX)) + .unwrap_or_default(); + let (payload_bytes, replacement_bytes) = image_data_url_estimate_adjustment(item); + if payload_bytes == 0 || replacement_bytes == 0 { + raw + } else { + raw.saturating_sub(payload_bytes) + .saturating_add(replacement_bytes) + } + } + } +} + +/// Approximate model-visible byte cost for one resized image input. +pub const RESIZED_IMAGE_BYTES_ESTIMATE: i64 = 7373; +const ORIGINAL_IMAGE_PATCH_SIZE: u32 = 32; +/// Maximum patch budget used when estimating `detail: "original"` image cost. +pub const ORIGINAL_IMAGE_MAX_PATCHES: usize = 10_000; +const ORIGINAL_IMAGE_ESTIMATE_CACHE_SIZE: usize = 32; + +static ORIGINAL_IMAGE_ESTIMATE_CACHE: LazyLock>> = + LazyLock::new(|| { + BlockingLruCache::new( + NonZeroUsize::new(ORIGINAL_IMAGE_ESTIMATE_CACHE_SIZE).unwrap_or(NonZeroUsize::MIN), + ) + }); + +fn parse_base64_image_data_url(url: &str) -> Option<&str> { + if !url + .get(.."data:".len()) + .is_some_and(|prefix| prefix.eq_ignore_ascii_case("data:")) + { + return None; + } + let comma_index = url.find(',')?; + let metadata = &url[..comma_index]; + let payload = &url[comma_index + 1..]; + let metadata_without_scheme = &metadata["data:".len()..]; + let mut metadata_parts = metadata_without_scheme.split(';'); + let mime_type = metadata_parts.next().unwrap_or_default(); + let has_base64_marker = metadata_parts.any(|part| part.eq_ignore_ascii_case("base64")); + if !mime_type + .get(.."image/".len()) + .is_some_and(|prefix| prefix.eq_ignore_ascii_case("image/")) + { + return None; + } + if !has_base64_marker { + return None; + } + Some(payload) +} + +fn estimate_original_image_bytes(image_url: &str) -> Option { + let key = sha1_digest(image_url.as_bytes()); + ORIGINAL_IMAGE_ESTIMATE_CACHE.get_or_insert_with(key, || { + let payload = match parse_base64_image_data_url(image_url) { + Some(payload) => payload, + None => { + tracing::trace!("skipping original-detail estimate for non-base64 image data URL"); + return None; + } + }; + let bytes = match BASE64_STANDARD.decode(payload) { + Ok(bytes) => bytes, + Err(error) => { + tracing::trace!("failed to decode original-detail image payload: {error}"); + return None; + } + }; + let dynamic = match image::load_from_memory(&bytes) { + Ok(dynamic) => dynamic, + Err(error) => { + tracing::trace!("failed to decode original-detail image bytes: {error}"); + return None; + } + }; + let width = i64::from(dynamic.width()); + let height = i64::from(dynamic.height()); + let patch_size = i64::from(ORIGINAL_IMAGE_PATCH_SIZE); + let patches_wide = width.saturating_add(patch_size.saturating_sub(1)) / patch_size; + let patches_high = height.saturating_add(patch_size.saturating_sub(1)) / patch_size; + let patch_count = patches_wide.saturating_mul(patches_high); + let patch_count = usize::try_from(patch_count).unwrap_or(usize::MAX); + let patch_count = patch_count.min(ORIGINAL_IMAGE_MAX_PATCHES); + Some(i64::try_from(approx_bytes_for_tokens(patch_count)).unwrap_or(i64::MAX)) + }) +} + +fn image_data_url_estimate_adjustment(item: &ResponseItem) -> (i64, i64) { + let mut payload_bytes = 0i64; + let mut replacement_bytes = 0i64; + + let mut accumulate = |image_url: &str, detail: Option| { + if let Some(payload_len) = parse_base64_image_data_url(image_url).map(str::len) { + payload_bytes = + payload_bytes.saturating_add(i64::try_from(payload_len).unwrap_or(i64::MAX)); + replacement_bytes = replacement_bytes.saturating_add(match detail { + Some(ImageDetail::Original) => { + estimate_original_image_bytes(image_url).unwrap_or(RESIZED_IMAGE_BYTES_ESTIMATE) + } + _ => RESIZED_IMAGE_BYTES_ESTIMATE, + }); + } + }; + + match item { + ResponseItem::Message { content, .. } => { + for content_item in content { + if let ContentItem::InputImage { image_url, detail } = content_item { + accumulate(image_url, *detail); + } + } + } + ResponseItem::FunctionCallOutput { output, .. } + | ResponseItem::CustomToolCallOutput { output, .. } => { + if let FunctionCallOutputBody::ContentItems(items) = &output.body { + for content_item in items { + if let FunctionCallOutputContentItem::InputImage { image_url, detail } = + content_item + { + accumulate(image_url, *detail); + } + } + } + } + _ => {} + } + + (payload_bytes, replacement_bytes) +} diff --git a/codex-rs/journal/src/history/mod.rs b/codex-rs/journal/src/history/mod.rs new file mode 100644 index 0000000000..2640782077 --- /dev/null +++ b/codex-rs/journal/src/history/mod.rs @@ -0,0 +1,26 @@ +//! Pure history utilities for normalizing, classifying, rewriting, and estimating prompt-ready +//! history items. + +mod classify; +mod estimate; +mod normalize; +mod transform; + +#[cfg(test)] +mod tests; + +pub use classify::is_api_message; +pub use classify::is_codex_generated_item; +pub use classify::is_model_generated_item; +pub use classify::is_user_turn_boundary; +pub use classify::user_turn_boundary_positions; +pub use estimate::ORIGINAL_IMAGE_MAX_PATCHES; +pub use estimate::RESIZED_IMAGE_BYTES_ESTIMATE; +pub use estimate::estimate_item_token_count; +pub use estimate::estimate_response_item_model_visible_bytes; +pub use normalize::ensure_call_outputs_present; +pub use normalize::remove_corresponding_for; +pub use normalize::remove_orphan_outputs; +pub use normalize::strip_images_when_unsupported; +pub use transform::replace_last_turn_images; +pub use transform::truncate_history_item; diff --git a/codex-rs/core/src/context_manager/normalize.rs b/codex-rs/journal/src/history/normalize.rs similarity index 69% rename from codex-rs/core/src/context_manager/normalize.rs rename to codex-rs/journal/src/history/normalize.rs index 839bae331e..d0dbda0eba 100644 --- a/codex-rs/core/src/context_manager/normalize.rs +++ b/codex-rs/journal/src/history/normalize.rs @@ -3,24 +3,21 @@ use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::InputModality; -use std::collections::HashSet; - -use crate::util::error_or_panic; +use tracing::error; use tracing::info; const IMAGE_CONTENT_OMITTED_PLACEHOLDER: &str = "image content omitted because you do not support image input"; -pub(crate) fn ensure_call_outputs_present(items: &mut Vec) { - // Collect synthetic outputs to insert immediately after their calls. - // Store the insertion position (index of call) alongside the item so - // we can insert in reverse order and avoid index shifting. +/// Ensures every tool-call item has a matching output item, inserting synthetic aborted outputs +/// immediately after calls that are missing one. +pub fn ensure_call_outputs_present(items: &mut Vec) { let mut missing_outputs_to_insert: Vec<(usize, ResponseItem)> = Vec::new(); - for (idx, item) in items.iter().enumerate() { + for (index, item) in items.iter().enumerate() { match item { ResponseItem::FunctionCall { call_id, .. } => { - let has_output = items.iter().any(|i| match i { + let has_output = items.iter().any(|other| match other { ResponseItem::FunctionCallOutput { call_id: existing, .. } => existing == call_id, @@ -30,7 +27,7 @@ pub(crate) fn ensure_call_outputs_present(items: &mut Vec) { if !has_output { info!("Function call output is missing for call id: {call_id}"); missing_outputs_to_insert.push(( - idx, + index, ResponseItem::FunctionCallOutput { call_id: call_id.clone(), output: FunctionCallOutputPayload::from_text("aborted".to_string()), @@ -42,7 +39,7 @@ pub(crate) fn ensure_call_outputs_present(items: &mut Vec) { call_id: Some(call_id), .. } => { - let has_output = items.iter().any(|i| match i { + let has_output = items.iter().any(|other| match other { ResponseItem::ToolSearchOutput { call_id: Some(existing), .. @@ -53,7 +50,7 @@ pub(crate) fn ensure_call_outputs_present(items: &mut Vec) { if !has_output { info!("Tool search output is missing for call id: {call_id}"); missing_outputs_to_insert.push(( - idx, + index, ResponseItem::ToolSearchOutput { call_id: Some(call_id.clone()), status: "completed".to_string(), @@ -64,7 +61,7 @@ pub(crate) fn ensure_call_outputs_present(items: &mut Vec) { } } ResponseItem::CustomToolCall { call_id, .. } => { - let has_output = items.iter().any(|i| match i { + let has_output = items.iter().any(|other| match other { ResponseItem::CustomToolCallOutput { call_id: existing, .. } => existing == call_id, @@ -72,11 +69,11 @@ pub(crate) fn ensure_call_outputs_present(items: &mut Vec) { }); if !has_output { - error_or_panic(format!( + report_invariant_violation(format!( "Custom tool call output is missing for call id: {call_id}" )); missing_outputs_to_insert.push(( - idx, + index, ResponseItem::CustomToolCallOutput { call_id: call_id.clone(), name: None, @@ -85,10 +82,9 @@ pub(crate) fn ensure_call_outputs_present(items: &mut Vec) { )); } } - // LocalShellCall is represented in upstream streams by a FunctionCallOutput ResponseItem::LocalShellCall { call_id, .. } => { if let Some(call_id) = call_id.as_ref() { - let has_output = items.iter().any(|i| match i { + let has_output = items.iter().any(|other| match other { ResponseItem::FunctionCallOutput { call_id: existing, .. } => existing == call_id, @@ -96,11 +92,11 @@ pub(crate) fn ensure_call_outputs_present(items: &mut Vec) { }); if !has_output { - error_or_panic(format!( + report_invariant_violation(format!( "Local shell call output is missing for call id: {call_id}" )); missing_outputs_to_insert.push(( - idx, + index, ResponseItem::FunctionCallOutput { call_id: call_id.clone(), output: FunctionCallOutputPayload::from_text("aborted".to_string()), @@ -113,57 +109,57 @@ pub(crate) fn ensure_call_outputs_present(items: &mut Vec) { } } - // Insert synthetic outputs in reverse index order to avoid re-indexing. - for (idx, output_item) in missing_outputs_to_insert.into_iter().rev() { - items.insert(idx + 1, output_item); + for (index, output_item) in missing_outputs_to_insert.into_iter().rev() { + items.insert(index + 1, output_item); } } -pub(crate) fn remove_orphan_outputs(items: &mut Vec) { - let function_call_ids: HashSet = items +/// Removes output items whose corresponding tool-call items no longer exist. +pub fn remove_orphan_outputs(items: &mut Vec) { + let function_call_ids = items .iter() - .filter_map(|i| match i { + .filter_map(|item| match item { ResponseItem::FunctionCall { call_id, .. } => Some(call_id.clone()), _ => None, }) - .collect(); + .collect::>(); - let tool_search_call_ids: HashSet = items + let tool_search_call_ids = items .iter() - .filter_map(|i| match i { + .filter_map(|item| match item { ResponseItem::ToolSearchCall { call_id: Some(call_id), .. } => Some(call_id.clone()), _ => None, }) - .collect(); + .collect::>(); - let local_shell_call_ids: HashSet = items + let local_shell_call_ids = items .iter() - .filter_map(|i| match i { + .filter_map(|item| match item { ResponseItem::LocalShellCall { call_id: Some(call_id), .. } => Some(call_id.clone()), _ => None, }) - .collect(); + .collect::>(); - let custom_tool_call_ids: HashSet = items + let custom_tool_call_ids = items .iter() - .filter_map(|i| match i { + .filter_map(|item| match item { ResponseItem::CustomToolCall { call_id, .. } => Some(call_id.clone()), _ => None, }) - .collect(); + .collect::>(); items.retain(|item| match item { ResponseItem::FunctionCallOutput { call_id, .. } => { let has_match = function_call_ids.contains(call_id) || local_shell_call_ids.contains(call_id); if !has_match { - error_or_panic(format!( + report_invariant_violation(format!( "Orphan function call output for call id: {call_id}" )); } @@ -172,7 +168,7 @@ pub(crate) fn remove_orphan_outputs(items: &mut Vec) { ResponseItem::CustomToolCallOutput { call_id, .. } => { let has_match = custom_tool_call_ids.contains(call_id); if !has_match { - error_or_panic(format!( + report_invariant_violation(format!( "Orphan custom tool call output for call id: {call_id}" )); } @@ -185,7 +181,9 @@ pub(crate) fn remove_orphan_outputs(items: &mut Vec) { } => { let has_match = tool_search_call_ids.contains(call_id); if !has_match { - error_or_panic(format!("Orphan tool search output for call id: {call_id}")); + report_invariant_violation(format!( + "Orphan tool search output for call id: {call_id}" + )); } has_match } @@ -194,12 +192,13 @@ pub(crate) fn remove_orphan_outputs(items: &mut Vec) { }); } -pub(crate) fn remove_corresponding_for(items: &mut Vec, item: &ResponseItem) { +/// Removes the matching paired item for `item` from the history, if one exists. +pub fn remove_corresponding_for(items: &mut Vec, item: &ResponseItem) { match item { ResponseItem::FunctionCall { call_id, .. } => { - remove_first_matching(items, |i| { + remove_first_matching(items, |other| { matches!( - i, + other, ResponseItem::FunctionCallOutput { call_id: existing, .. } if existing == call_id @@ -207,23 +206,35 @@ pub(crate) fn remove_corresponding_for(items: &mut Vec, item: &Res }); } ResponseItem::FunctionCallOutput { call_id, .. } => { - if let Some(pos) = items.iter().position(|i| { - matches!(i, ResponseItem::FunctionCall { call_id: existing, .. } if existing == call_id) + if let Some(position) = items.iter().position(|other| { + matches!( + other, + ResponseItem::FunctionCall { + call_id: existing, + .. + } if existing == call_id + ) }) { - items.remove(pos); - } else if let Some(pos) = items.iter().position(|i| { - matches!(i, ResponseItem::LocalShellCall { call_id: Some(existing), .. } if existing == call_id) + items.remove(position); + } else if let Some(position) = items.iter().position(|other| { + matches!( + other, + ResponseItem::LocalShellCall { + call_id: Some(existing), + .. + } if existing == call_id + ) }) { - items.remove(pos); + items.remove(position); } } ResponseItem::ToolSearchCall { call_id: Some(call_id), .. } => { - remove_first_matching(items, |i| { + remove_first_matching(items, |other| { matches!( - i, + other, ResponseItem::ToolSearchOutput { call_id: Some(existing), .. @@ -235,23 +246,20 @@ pub(crate) fn remove_corresponding_for(items: &mut Vec, item: &Res call_id: Some(call_id), .. } => { - remove_first_matching( - items, - |i| { - matches!( - i, - ResponseItem::ToolSearchCall { - call_id: Some(existing), - .. - } if existing == call_id - ) - }, - ); + remove_first_matching(items, |other| { + matches!( + other, + ResponseItem::ToolSearchCall { + call_id: Some(existing), + .. + } if existing == call_id + ) + }); } ResponseItem::CustomToolCall { call_id, .. } => { - remove_first_matching(items, |i| { + remove_first_matching(items, |other| { matches!( - i, + other, ResponseItem::CustomToolCallOutput { call_id: existing, .. } if existing == call_id @@ -259,18 +267,23 @@ pub(crate) fn remove_corresponding_for(items: &mut Vec, item: &Res }); } ResponseItem::CustomToolCallOutput { call_id, .. } => { - remove_first_matching( - items, - |i| matches!(i, ResponseItem::CustomToolCall { call_id: existing, .. } if existing == call_id), - ); + remove_first_matching(items, |other| { + matches!( + other, + ResponseItem::CustomToolCall { + call_id: existing, + .. + } if existing == call_id + ) + }); } ResponseItem::LocalShellCall { call_id: Some(call_id), .. } => { - remove_first_matching(items, |i| { + remove_first_matching(items, |other| { matches!( - i, + other, ResponseItem::FunctionCallOutput { call_id: existing, .. } if existing == call_id @@ -281,23 +294,15 @@ pub(crate) fn remove_corresponding_for(items: &mut Vec, item: &Res } } -fn remove_first_matching(items: &mut Vec, predicate: F) -where - F: Fn(&ResponseItem) -> bool, -{ - if let Some(pos) = items.iter().position(predicate) { - items.remove(pos); - } -} - -/// Strip image content from messages and tool outputs when the model does not support images. -/// When `input_modalities` contains `InputModality::Image`, no stripping is performed. -pub(crate) fn strip_images_when_unsupported( +/// Strips image inputs from messages and tool outputs when the model does not support images. +/// +/// Image-generation call results are cleared in the same mode so text-only prompts do not carry +/// image payloads forward. +pub fn strip_images_when_unsupported( input_modalities: &[InputModality], items: &mut [ResponseItem], ) { - let supports_images = input_modalities.contains(&InputModality::Image); - if supports_images { + if input_modalities.contains(&InputModality::Image) { return; } @@ -343,3 +348,20 @@ pub(crate) fn strip_images_when_unsupported( } } } + +fn remove_first_matching(items: &mut Vec, predicate: F) +where + F: Fn(&ResponseItem) -> bool, +{ + if let Some(position) = items.iter().position(predicate) { + items.remove(position); + } +} + +fn report_invariant_violation(message: String) { + if cfg!(debug_assertions) { + panic!("{message}"); + } else { + error!("{message}"); + } +} diff --git a/codex-rs/journal/src/history/tests.rs b/codex-rs/journal/src/history/tests.rs new file mode 100644 index 0000000000..4daf6dad7a --- /dev/null +++ b/codex-rs/journal/src/history/tests.rs @@ -0,0 +1,286 @@ +use super::ensure_call_outputs_present; +use super::estimate_response_item_model_visible_bytes; +use super::is_user_turn_boundary; +use super::remove_corresponding_for; +use super::replace_last_turn_images; +use super::strip_images_when_unsupported; +use super::truncate_history_item; +use codex_protocol::AgentPath; +use codex_protocol::models::ContentItem; +use codex_protocol::models::DEFAULT_IMAGE_DETAIL; +use codex_protocol::models::FunctionCallOutputContentItem; +use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::ResponseItem; +use codex_protocol::openai_models::InputModality; +use codex_protocol::protocol::InterAgentCommunication; +use codex_utils_output_truncation::TruncationPolicy; +use pretty_assertions::assert_eq; + +fn user_message(text: &str) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: text.to_string(), + }], + phase: None, + } +} + +fn assistant_message(text: &str) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: text.to_string(), + }], + phase: None, + } +} + +#[test] +fn ensure_call_outputs_present_inserts_missing_function_output_after_call() { + let mut items = vec![ResponseItem::FunctionCall { + id: None, + name: "shell".to_string(), + namespace: None, + arguments: "{}".to_string(), + call_id: "call-1".to_string(), + }]; + + ensure_call_outputs_present(&mut items); + + assert_eq!( + items, + vec![ + ResponseItem::FunctionCall { + id: None, + name: "shell".to_string(), + namespace: None, + arguments: "{}".to_string(), + call_id: "call-1".to_string(), + }, + ResponseItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload::from_text("aborted".to_string()), + }, + ] + ); +} + +#[test] +fn remove_corresponding_for_removes_matching_tool_output() { + let removed = ResponseItem::FunctionCall { + id: None, + name: "shell".to_string(), + namespace: None, + arguments: "{}".to_string(), + call_id: "call-1".to_string(), + }; + let mut items = vec![ + removed.clone(), + ResponseItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload::from_text("done".to_string()), + }, + ]; + + items.remove(0); + remove_corresponding_for(&mut items, &removed); + + assert_eq!(items, Vec::new()); +} + +#[test] +fn strip_images_when_unsupported_rewrites_messages_and_tool_outputs() { + let mut items = vec![ + ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ + ContentItem::InputText { + text: "look".to_string(), + }, + ContentItem::InputImage { + image_url: "https://example.com/img.png".to_string(), + detail: Some(DEFAULT_IMAGE_DETAIL), + }, + ], + phase: None, + }, + ResponseItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload::from_content_items(vec![ + FunctionCallOutputContentItem::InputImage { + image_url: "https://example.com/tool.png".to_string(), + detail: Some(DEFAULT_IMAGE_DETAIL), + }, + ]), + }, + ResponseItem::ImageGenerationCall { + id: "ig-1".to_string(), + status: "completed".to_string(), + revised_prompt: None, + result: "Zm9v".to_string(), + }, + ]; + + strip_images_when_unsupported(&[InputModality::Text], &mut items); + + assert_eq!( + items, + vec![ + ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ + ContentItem::InputText { + text: "look".to_string(), + }, + ContentItem::InputText { + text: "image content omitted because you do not support image input" + .to_string(), + }, + ], + phase: None, + }, + ResponseItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload::from_content_items(vec![ + FunctionCallOutputContentItem::InputText { + text: "image content omitted because you do not support image input" + .to_string(), + }, + ]), + }, + ResponseItem::ImageGenerationCall { + id: "ig-1".to_string(), + status: "completed".to_string(), + revised_prompt: None, + result: String::new(), + }, + ] + ); +} + +#[test] +fn inter_agent_assistant_messages_are_turn_boundaries() { + let communication = InterAgentCommunication::new( + AgentPath::root(), + AgentPath::root() + .join("worker") + .expect("sub-agent path should be valid"), + Vec::new(), + "continue".to_string(), + /*trigger_turn*/ true, + ); + let item = ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: serde_json::to_string(&communication).expect("message should serialize"), + }], + phase: None, + }; + + assert!(is_user_turn_boundary(&item)); +} + +#[test] +fn estimate_response_item_model_visible_bytes_discounts_inline_image_data_urls() { + let payload = "a".repeat(20_000); + let image_item = ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputImage { + image_url: format!("data:image/png;base64,{payload}"), + detail: Some(DEFAULT_IMAGE_DETAIL), + }], + phase: None, + }; + + let estimated = estimate_response_item_model_visible_bytes(&image_item); + let raw = i64::try_from( + serde_json::to_string(&image_item) + .expect("item should serialize") + .len(), + ) + .expect("raw length should fit"); + + assert!(estimated > 0); + assert!(estimated < raw); +} + +#[test] +fn truncate_history_item_truncates_tool_outputs_but_not_messages() { + let output = "x".repeat(20_000); + let function_output = ResponseItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload::from_text(output), + }; + let user = user_message("hello"); + + let truncated = truncate_history_item( + &function_output, + TruncationPolicy::Tokens(/*max_tokens*/ 64), + ); + + assert_ne!(truncated, function_output); + assert_eq!( + truncate_history_item(&user, TruncationPolicy::Tokens(/*max_tokens*/ 64)), + user + ); +} + +#[test] +fn replace_last_turn_images_only_rewrites_latest_tool_output() { + let mut items = vec![ + assistant_message("already done"), + ResponseItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload::from_content_items(vec![ + FunctionCallOutputContentItem::InputImage { + image_url: "https://example.com/older.png".to_string(), + detail: Some(DEFAULT_IMAGE_DETAIL), + }, + ]), + }, + user_message("new turn"), + ResponseItem::FunctionCallOutput { + call_id: "call-2".to_string(), + output: FunctionCallOutputPayload::from_content_items(vec![ + FunctionCallOutputContentItem::InputImage { + image_url: "https://example.com/newer.png".to_string(), + detail: Some(DEFAULT_IMAGE_DETAIL), + }, + ]), + }, + ]; + + let replaced = replace_last_turn_images(&mut items, "omitted"); + + assert!(replaced); + assert_eq!( + items[1], + ResponseItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload::from_content_items(vec![ + FunctionCallOutputContentItem::InputImage { + image_url: "https://example.com/older.png".to_string(), + detail: Some(DEFAULT_IMAGE_DETAIL), + }, + ]), + } + ); + assert_eq!( + items[3], + ResponseItem::FunctionCallOutput { + call_id: "call-2".to_string(), + output: FunctionCallOutputPayload::from_content_items(vec![ + FunctionCallOutputContentItem::InputText { + text: "omitted".to_string(), + }, + ]), + } + ); +} diff --git a/codex-rs/journal/src/history/transform.rs b/codex-rs/journal/src/history/transform.rs new file mode 100644 index 0000000000..47a5a10191 --- /dev/null +++ b/codex-rs/journal/src/history/transform.rs @@ -0,0 +1,92 @@ +use crate::history::is_user_turn_boundary; +use codex_protocol::models::FunctionCallOutputBody; +use codex_protocol::models::FunctionCallOutputContentItem; +use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::ResponseItem; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::truncate_function_output_items_with_policy; +use codex_utils_output_truncation::truncate_text; + +/// Truncates one history item before it is recorded in the in-memory journal. +/// +/// Tool outputs receive a small serialization headroom multiplier so the JSON wrapper bytes do not +/// cause unexpected overages after truncation. +pub fn truncate_history_item(item: &ResponseItem, policy: TruncationPolicy) -> ResponseItem { + let policy_with_serialization_budget = policy * 1.2; + match item { + ResponseItem::FunctionCallOutput { call_id, output } => ResponseItem::FunctionCallOutput { + call_id: call_id.clone(), + output: truncate_function_output_payload(output, policy_with_serialization_budget), + }, + ResponseItem::CustomToolCallOutput { + call_id, + name, + output, + } => ResponseItem::CustomToolCallOutput { + call_id: call_id.clone(), + name: name.clone(), + output: truncate_function_output_payload(output, policy_with_serialization_budget), + }, + ResponseItem::Message { .. } + | ResponseItem::Reasoning { .. } + | ResponseItem::LocalShellCall { .. } + | ResponseItem::FunctionCall { .. } + | ResponseItem::ToolSearchCall { .. } + | ResponseItem::ToolSearchOutput { .. } + | ResponseItem::WebSearchCall { .. } + | ResponseItem::ImageGenerationCall { .. } + | ResponseItem::CustomToolCall { .. } + | ResponseItem::Compaction { .. } + | ResponseItem::Other => item.clone(), + } +} + +/// Replaces image content in the last tool-output item of the current turn. +pub fn replace_last_turn_images(items: &mut [ResponseItem], placeholder: &str) -> bool { + let Some(index) = items.iter().rposition(|item| { + matches!(item, ResponseItem::FunctionCallOutput { .. }) || is_user_turn_boundary(item) + }) else { + return false; + }; + + match &mut items[index] { + ResponseItem::FunctionCallOutput { output, .. } => { + let Some(content_items) = output.content_items_mut() else { + return false; + }; + + let mut replaced = false; + let placeholder = placeholder.to_string(); + for item in content_items.iter_mut() { + if matches!(item, FunctionCallOutputContentItem::InputImage { .. }) { + *item = FunctionCallOutputContentItem::InputText { + text: placeholder.clone(), + }; + replaced = true; + } + } + replaced + } + ResponseItem::Message { .. } => false, + _ => false, + } +} + +fn truncate_function_output_payload( + output: &FunctionCallOutputPayload, + policy: TruncationPolicy, +) -> FunctionCallOutputPayload { + let body = match &output.body { + FunctionCallOutputBody::Text(content) => { + FunctionCallOutputBody::Text(truncate_text(content, policy)) + } + FunctionCallOutputBody::ContentItems(items) => FunctionCallOutputBody::ContentItems( + truncate_function_output_items_with_policy(items, policy), + ), + }; + + FunctionCallOutputPayload { + body, + success: output.success, + } +} diff --git a/codex-rs/journal/src/journal.rs b/codex-rs/journal/src/journal.rs new file mode 100644 index 0000000000..d4f70e29bf --- /dev/null +++ b/codex-rs/journal/src/journal.rs @@ -0,0 +1,345 @@ +use crate::JournalEntry; +use crate::JournalError; +use crate::JournalKey; +use crate::KeyFilter; +use crate::PromptView; +use crate::Result; +use codex_protocol::journal::JournalCheckpointItem; +use codex_protocol::journal::JournalContextAudience; +use codex_protocol::journal::JournalContextForkBehavior; +use codex_protocol::journal::JournalContextItem; +use codex_protocol::journal::JournalContextKey; +use codex_protocol::journal::JournalHistoryCursor; +use codex_protocol::journal::JournalHistoryItem; +use codex_protocol::journal::JournalItem; +use codex_protocol::journal::JournalReplacePrefixCheckpoint; +use codex_protocol::journal::JournalTruncateHistoryCheckpoint; +use codex_protocol::models::ResponseItem; +use indexmap::IndexMap; +use std::fs::File; +use std::io::BufRead; +use std::io::BufReader; +use std::io::BufWriter; +use std::io::Write; +use std::path::Path; + +/// Canonical typed journal. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct Journal { + entries: Vec, +} + +impl Journal { + /// Creates an empty journal. + pub fn new() -> Self { + Self::default() + } + + /// Creates a journal from an explicit list of entries. + pub fn from_entries(entries: Vec) -> Self { + Self { entries } + } + + /// Appends one keyed item to the journal. + pub fn add(&mut self, key: K, item: T) + where + K: Into, + T: Into, + { + self.entries.push(JournalEntry::new(key, item)); + } + + /// Appends several keyed journal entries. + pub fn extend(&mut self, entries: I) + where + I: IntoIterator, + T: Into, + { + self.entries.extend(entries.into_iter().map(Into::into)); + } + + /// Returns the raw append-only journal entries. + pub fn entries(&self) -> &[JournalEntry] { + &self.entries + } + + /// Returns a journal containing only journal entries whose keys match the filter. + pub fn filter(&self, filter: &KeyFilter) -> Self { + let entries = self + .entries + .iter() + .filter(|entry| filter.matches(&entry.key)) + .cloned() + .collect(); + Self::from_entries(entries) + } + + /// Returns the number of journal entries. + pub fn len(&self) -> usize { + self.entries.len() + } + + /// Returns whether the journal is empty. + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + /// Renders the current effective journal view into model prompt items. + pub fn to_prompt(&self, view: &PromptView) -> Result> { + self.to_prompt_matching_filter(view, None) + } + + /// Renders the current effective journal view into model prompt items after selecting + /// only journal entries whose keys match the filter. + pub fn to_prompt_with_filter( + &self, + view: &PromptView, + filter: &KeyFilter, + ) -> Result> { + self.to_prompt_matching_filter(view, Some(filter)) + } + + /// Produces a flattened child state for the provided view. + pub fn fork(&self, view: &PromptView) -> Result { + self.fork_matching_filter(view, None) + } + + /// Produces a flattened child state for the provided view after selecting only + /// journal entries whose keys match the filter. + pub fn fork_with_filter(&self, view: &PromptView, filter: &KeyFilter) -> Result { + self.fork_matching_filter(view, Some(filter)) + } + + /// Drops obsolete journal entries and keeps only the current effective journal view. + /// + /// This is the first building block for a rolling in-memory window: callers can + /// persist the full journal elsewhere, then keep only the flattened journal hot. + pub fn flatten(&self) -> Result { + let resolved = self.resolve(None)?; + Ok(Self::from_entries( + resolved + .contexts + .into_iter() + .chain(resolved.history) + .collect(), + )) + } + + /// Keeps only the current effective journal view plus the history suffix that starts + /// at the resolved cursor. + /// + /// This is a lightweight rolling-window helper: callers can persist the full + /// journal on disk, then keep only a recent hot suffix in memory. + pub fn with_history_window(&self, start: &JournalHistoryCursor) -> Result { + let resolved = self.resolve(None)?; + let start_index = resolve_cursor(resolved.history.as_slice(), start)?; + Ok(Self::from_entries( + resolved + .contexts + .into_iter() + .chain(resolved.history.into_iter().skip(start_index)) + .collect(), + )) + } + + /// Persists the raw journal to a JSONL file, one `JournalEntry` per line. + pub fn persist_jsonl(&self, path: &Path) -> Result<()> { + let file = File::create(path)?; + let mut writer = BufWriter::new(file); + for entry in &self.entries { + serde_json::to_writer(&mut writer, entry)?; + writer.write_all(b"\n")?; + } + writer.flush()?; + Ok(()) + } + + /// Loads a raw journal from a JSONL file written by [`Self::persist_jsonl`]. + pub fn load_jsonl(path: &Path) -> Result { + let file = File::open(path)?; + let reader = BufReader::new(file); + let mut entries = Vec::new(); + for (line_index, line) in reader.lines().enumerate() { + let line = line?; + if line.trim().is_empty() { + continue; + } + let entry = serde_json::from_str::(&line).map_err(|source| { + JournalError::ParseJson { + line_number: line_index + 1, + source, + } + })?; + entries.push(entry); + } + Ok(Self::from_entries(entries)) + } + + fn to_prompt_matching_filter( + &self, + view: &PromptView, + filter: Option<&KeyFilter>, + ) -> Result> { + let resolved = self.resolve(filter)?; + let mut prompt: Vec = resolved + .contexts + .into_iter() + .filter(|entry| context_visible_in_view(context_item(entry), view)) + .map(|entry| match entry.item { + JournalItem::Context(item) => ResponseItem::from(item), + _ => unreachable!("resolved context entries must be context items"), + }) + .collect(); + prompt.extend(resolved.history.into_iter().map(|entry| match entry.item { + JournalItem::History(item) => ResponseItem::from(item), + _ => unreachable!("resolved history entries must be history items"), + })); + Ok(prompt) + } + + fn fork_matching_filter(&self, view: &PromptView, filter: Option<&KeyFilter>) -> Result { + let resolved = self.resolve(filter)?; + Ok(Self::from_entries( + resolved + .contexts + .into_iter() + .filter(|entry| context_item(entry).on_fork == JournalContextForkBehavior::Keep) + .filter(|entry| context_visible_in_view(context_item(entry), view)) + .chain(resolved.history) + .collect(), + )) + } + + fn resolve(&self, filter: Option<&KeyFilter>) -> Result { + let mut history = Vec::::new(); + let mut latest_context_by_key = IndexMap::::new(); + + for (index, entry) in self.entries.iter().enumerate() { + if let Some(filter) = filter + && !filter.matches(&entry.key) + { + continue; + } + match &entry.item { + JournalItem::History(_) => history.push(entry.clone()), + JournalItem::Context(context_item) => { + latest_context_by_key.insert(context_item.key.clone(), (index, entry.clone())); + } + JournalItem::Checkpoint(checkpoint) => { + apply_checkpoint(&mut history, &entry.key, checkpoint)?; + } + } + } + + let mut contexts = latest_context_by_key.into_values().collect::>(); + contexts.sort_by(|(left_index, left_entry), (right_index, right_entry)| { + context_item(left_entry) + .prompt_order + .cmp(&context_item(right_entry).prompt_order) + .then_with(|| left_index.cmp(right_index)) + }); + + Ok(ResolvedJournal { + contexts: contexts.into_iter().map(|(_, entry)| entry).collect(), + history, + }) + } +} + +#[derive(Debug)] +struct ResolvedJournal { + contexts: Vec, + history: Vec, +} + +fn apply_checkpoint( + history: &mut Vec, + checkpoint_key: &JournalKey, + checkpoint: &JournalCheckpointItem, +) -> Result<()> { + match checkpoint { + JournalCheckpointItem::ReplacePrefix(JournalReplacePrefixCheckpoint { + through, + replacement, + }) => { + let keep_from = resolve_cursor(history.as_slice(), through)?; + let mut next_history = + Vec::with_capacity(replacement.len() + history.len().saturating_sub(keep_from)); + next_history.extend( + replacement + .iter() + .cloned() + .enumerate() + .map(|(index, item)| { + JournalEntry::new( + replacement_history_key(checkpoint_key, index, &item), + item, + ) + }), + ); + next_history.extend(history[keep_from..].iter().cloned()); + *history = next_history; + Ok(()) + } + JournalCheckpointItem::TruncateHistory(JournalTruncateHistoryCheckpoint { through }) => { + let keep_len = resolve_cursor(history.as_slice(), through)?; + history.truncate(keep_len); + Ok(()) + } + } +} + +fn resolve_cursor(history: &[JournalEntry], cursor: &JournalHistoryCursor) -> Result { + match cursor { + JournalHistoryCursor::Start => Ok(0), + JournalHistoryCursor::End => Ok(history.len()), + JournalHistoryCursor::AfterItem(history_item_id) => history + .iter() + .position(|entry| history_item(entry).id == *history_item_id) + .map(|index| index + 1) + .ok_or_else(|| JournalError::UnknownHistoryItemId { + history_item_id: history_item_id.clone(), + }), + } +} + +fn replacement_history_key( + checkpoint_key: &JournalKey, + index: usize, + history_item: &JournalHistoryItem, +) -> JournalKey { + checkpoint_key + .child("replacement") + .child(index.to_string()) + .child(history_item.id.clone()) +} + +fn context_item(entry: &JournalEntry) -> &JournalContextItem { + match &entry.item { + JournalItem::Context(item) => item, + _ => unreachable!("resolved context entries must be context items"), + } +} + +fn history_item(entry: &JournalEntry) -> &JournalHistoryItem { + match &entry.item { + JournalItem::History(item) => item, + _ => unreachable!("resolved history entries must be history items"), + } +} + +fn context_visible_in_view(item: &JournalContextItem, view: &PromptView) -> bool { + match &item.audience { + JournalContextAudience::All => true, + JournalContextAudience::RootOnly => view.is_root, + JournalContextAudience::SubAgentsOnly => !view.is_root, + JournalContextAudience::AgentPathPrefix(prefix) => view + .agent_path + .as_deref() + .is_some_and(|agent_path| agent_path.starts_with(prefix)), + JournalContextAudience::AgentRole(role) => view + .agent_role + .as_deref() + .is_some_and(|agent_role| agent_role == role), + } +} diff --git a/codex-rs/journal/src/lib.rs b/codex-rs/journal/src/lib.rs new file mode 100644 index 0000000000..5d020bffd3 --- /dev/null +++ b/codex-rs/journal/src/lib.rs @@ -0,0 +1,29 @@ +//! Typed journal model for prompt rendering, filtering, forking, and persistence. + +mod error; +pub mod history; +mod journal; +mod prompt_view; + +#[cfg(test)] +mod tests; + +pub use codex_protocol::journal::JournalCheckpointItem; +pub use codex_protocol::journal::JournalContextAudience; +pub use codex_protocol::journal::JournalContextForkBehavior; +pub use codex_protocol::journal::JournalContextItem; +pub use codex_protocol::journal::JournalContextKey; +pub use codex_protocol::journal::JournalEntry; +pub use codex_protocol::journal::JournalHistoryCursor; +pub use codex_protocol::journal::JournalHistoryItem; +pub use codex_protocol::journal::JournalItem; +pub use codex_protocol::journal::JournalKey; +pub use codex_protocol::journal::JournalReplacePrefixCheckpoint; +pub use codex_protocol::journal::JournalTruncateHistoryCheckpoint; +pub use codex_protocol::journal::KeyFilter; +pub use codex_protocol::journal::PromptMessage; +pub use codex_protocol::journal::PromptMessageRole; +pub use error::JournalError; +pub use error::Result; +pub use journal::Journal; +pub use prompt_view::PromptView; diff --git a/codex-rs/journal/src/prompt_view.rs b/codex-rs/journal/src/prompt_view.rs new file mode 100644 index 0000000000..3e7fb15064 --- /dev/null +++ b/codex-rs/journal/src/prompt_view.rs @@ -0,0 +1,33 @@ +/// View configuration used when rendering or forking a journal. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct PromptView { + pub(crate) is_root: bool, + pub(crate) agent_path: Option, + pub(crate) agent_role: Option, +} + +impl PromptView { + /// Returns the view for the root agent. + pub fn root() -> Self { + Self { + is_root: true, + agent_path: None, + agent_role: None, + } + } + + /// Returns the view for a spawned subagent. + pub fn subagent(agent_path: impl Into, agent_role: Option) -> Self { + Self { + is_root: false, + agent_path: Some(agent_path.into()), + agent_role, + } + } + + /// Sets the agent role used for audience matching. + pub fn with_agent_role(mut self, agent_role: impl Into) -> Self { + self.agent_role = Some(agent_role.into()); + self + } +} diff --git a/codex-rs/journal/src/tests.rs b/codex-rs/journal/src/tests.rs new file mode 100644 index 0000000000..81bac17474 --- /dev/null +++ b/codex-rs/journal/src/tests.rs @@ -0,0 +1,349 @@ +use crate::Journal; +use crate::JournalCheckpointItem; +use crate::JournalContextAudience; +use crate::JournalContextForkBehavior; +use crate::JournalContextItem; +use crate::JournalContextKey; +use crate::JournalEntry; +use crate::JournalHistoryCursor; +use crate::JournalHistoryItem; +use crate::JournalReplacePrefixCheckpoint; +use crate::JournalTruncateHistoryCheckpoint; +use crate::KeyFilter; +use crate::PromptMessage; +use crate::PromptView; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; +use pretty_assertions::assert_eq; +use tempfile::tempdir; + +fn user_message(text: &str) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: text.to_string(), + }], + phase: None, + } +} + +fn assistant_message(text: &str) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: text.to_string(), + }], + phase: None, + } +} + +fn developer_context( + namespace: &str, + name: &str, + text: &str, + prompt_order: i64, +) -> JournalContextItem { + JournalContextItem::new( + JournalContextKey::new(namespace, name, None), + PromptMessage::developer_text(text), + ) + .with_prompt_order(prompt_order) +} + +#[test] +fn to_prompt_uses_latest_context_for_key() { + let mut state = Journal::new(); + state.add( + ["prompt", "permissions", "older"], + developer_context("context", "permissions", "older permissions", 10), + ); + state.add(["history", "hello"], user_message("hello")); + state.add( + ["prompt", "permissions", "newer"], + developer_context("context", "permissions", "newer permissions", 10), + ); + + let prompt = state + .to_prompt(&PromptView::root()) + .expect("prompt should render"); + + assert_eq!( + prompt, + vec![ + ResponseItem::from(PromptMessage::developer_text("newer permissions")), + user_message("hello"), + ] + ); +} + +#[test] +fn to_prompt_filters_context_by_audience() { + let mut state = Journal::new(); + state.add( + ["prompt", "root", "hint"], + developer_context("context", "root", "root-only", 0) + .with_audience(JournalContextAudience::RootOnly), + ); + state.add( + ["prompt", "child", "hint"], + developer_context("context", "child", "child-only", 1) + .with_audience(JournalContextAudience::SubAgentsOnly), + ); + + let root_prompt = state + .to_prompt(&PromptView::root()) + .expect("root prompt should render"); + let child_prompt = state + .to_prompt(&PromptView::subagent( + "/root/worker", + Option::::None, + )) + .expect("child prompt should render"); + + assert_eq!( + root_prompt, + vec![ResponseItem::from(PromptMessage::developer_text( + "root-only" + ))] + ); + assert_eq!( + child_prompt, + vec![ResponseItem::from(PromptMessage::developer_text( + "child-only" + ))] + ); +} + +#[test] +fn to_prompt_with_filter_matches_key_prefix() { + let mut state = Journal::new(); + state.add( + ["prompt", "root", "keep"], + developer_context("context", "keep", "keep me", 0), + ); + state.add( + ["prompt", "child", "drop"], + developer_context("context", "drop", "drop me", 1), + ); + + let prompt = state + .to_prompt_with_filter(&PromptView::root(), &KeyFilter::prefix(["prompt", "root"])) + .expect("prompt should render"); + + assert_eq!( + prompt, + vec![ResponseItem::from(PromptMessage::developer_text("keep me"))] + ); +} + +#[test] +fn checkpoints_replace_prefix_and_then_truncate_history() { + let first = JournalHistoryItem::new(user_message("turn 1")); + let second = JournalHistoryItem::new(assistant_message("turn 1 answer")); + let third = JournalHistoryItem::new(user_message("turn 2")); + let summary = JournalHistoryItem { + id: "summary".to_string(), + turn_id: None, + item: assistant_message("summary"), + }; + + let state = Journal::from_entries(vec![ + JournalEntry::new(["history", "1"], first), + JournalEntry::new(["history", "2"], second.clone()), + JournalEntry::new(["history", "3"], third), + JournalEntry::new( + ["checkpoint", "replace"], + JournalCheckpointItem::ReplacePrefix(JournalReplacePrefixCheckpoint { + through: JournalHistoryCursor::AfterItem(second.id), + replacement: vec![summary.clone()], + }), + ), + JournalEntry::new( + ["checkpoint", "truncate"], + JournalCheckpointItem::TruncateHistory(JournalTruncateHistoryCheckpoint { + through: JournalHistoryCursor::AfterItem(summary.id), + }), + ), + ]); + + let prompt = state + .to_prompt(&PromptView::root()) + .expect("prompt should render"); + + assert_eq!(prompt, vec![assistant_message("summary")]); +} + +#[test] +fn flatten_preserves_prompt_and_drops_obsolete_items() { + let first = JournalHistoryItem::new(user_message("turn 1")); + let answer = JournalHistoryItem::new(assistant_message("turn 1 answer")); + let summary = JournalHistoryItem { + id: "summary".to_string(), + turn_id: None, + item: assistant_message("summary"), + }; + let state = Journal::from_entries(vec![ + JournalEntry::new( + ["prompt", "permissions", "old"], + developer_context("context", "permissions", "old", 0), + ), + JournalEntry::new( + ["prompt", "permissions", "new"], + developer_context("context", "permissions", "new", 0), + ), + JournalEntry::new(["history", "1"], first), + JournalEntry::new(["history", "2"], answer.clone()), + JournalEntry::new( + ["checkpoint", "replace"], + JournalCheckpointItem::ReplacePrefix(JournalReplacePrefixCheckpoint { + through: JournalHistoryCursor::AfterItem(answer.id), + replacement: vec![summary.clone()], + }), + ), + ]); + + let before = state + .to_prompt(&PromptView::root()) + .expect("prompt should render"); + let flattened = state.flatten().expect("flatten should succeed"); + let after = flattened + .to_prompt(&PromptView::root()) + .expect("flattened prompt should render"); + + assert_eq!(before, after); + assert_eq!( + flattened.entries(), + vec![ + JournalEntry::new( + ["prompt", "permissions", "new"], + developer_context("context", "permissions", "new", 0), + ), + JournalEntry::new( + ["checkpoint", "replace", "replacement", "0", "summary"], + summary, + ), + ] + ); +} + +#[test] +fn with_history_window_keeps_only_recent_effective_history() { + let first = JournalHistoryItem::new(user_message("turn 1")); + let second = JournalHistoryItem::new(assistant_message("turn 1 answer")); + let third = JournalHistoryItem::new(user_message("turn 2")); + let state = Journal::from_entries(vec![ + JournalEntry::new( + ["prompt", "permissions", "current"], + developer_context("context", "permissions", "p", 0), + ), + JournalEntry::new(["history", "1"], first), + JournalEntry::new(["history", "2"], second.clone()), + JournalEntry::new(["history", "3"], third.clone()), + ]); + + let windowed = state + .with_history_window(&JournalHistoryCursor::AfterItem(second.id)) + .expect("windowing should succeed"); + + assert_eq!( + windowed.entries(), + vec![ + JournalEntry::new( + ["prompt", "permissions", "current"], + developer_context("context", "permissions", "p", 0), + ), + JournalEntry::new(["history", "3"], third), + ] + ); +} + +#[test] +fn fork_drops_non_keep_context_and_respects_audience() { + let history = JournalHistoryItem::new(user_message("hello")); + let state = Journal::from_entries(vec![ + JournalEntry::new( + ["prompt", "child", "shared"], + developer_context("context", "shared", "shared child context", 0) + .with_audience(JournalContextAudience::SubAgentsOnly), + ), + JournalEntry::new( + ["prompt", "child", "regenerate"], + developer_context("context", "regenerate", "usage hint", 1) + .with_audience(JournalContextAudience::SubAgentsOnly) + .with_on_fork(JournalContextForkBehavior::Regenerate), + ), + JournalEntry::new( + ["prompt", "root", "only"], + developer_context("context", "root", "root only", 2) + .with_audience(JournalContextAudience::RootOnly), + ), + JournalEntry::new(["history", "hello"], history.clone()), + ]); + + let forked = state + .fork(&PromptView::subagent( + "/root/worker", + Option::::None, + )) + .expect("fork should succeed"); + + assert_eq!( + forked.entries(), + vec![ + JournalEntry::new( + ["prompt", "child", "shared"], + developer_context("context", "shared", "shared child context", 0) + .with_audience(JournalContextAudience::SubAgentsOnly), + ), + JournalEntry::new(["history", "hello"], history), + ] + ); +} + +#[test] +fn persist_and_load_jsonl_round_trip() { + let dir = tempdir().expect("tempdir"); + let path = dir.path().join("journal.jsonl"); + let history = JournalHistoryItem::new(user_message("hello")); + let state = Journal::from_entries(vec![ + JournalEntry::new( + ["prompt", "permissions", "current"], + developer_context("context", "permissions", "p", 0), + ), + JournalEntry::new(["history", "hello"], history), + ]); + + state + .persist_jsonl(path.as_path()) + .expect("journal should persist"); + let loaded = Journal::load_jsonl(path.as_path()).expect("journal should load"); + + assert_eq!(loaded, state); +} + +#[test] +fn filter_returns_matching_raw_entries() { + let state = Journal::from_entries(vec![ + JournalEntry::new( + ["prompt", "root", "keep"], + developer_context("context", "keep", "keep me", 0), + ), + JournalEntry::new( + ["prompt", "child", "drop"], + developer_context("context", "drop", "drop me", 1), + ), + JournalEntry::new(["history", "hello"], user_message("hello")), + ]); + + let filtered = state.filter(&KeyFilter::prefix(["prompt", "root"])); + + assert_eq!( + filtered.entries(), + vec![JournalEntry::new( + ["prompt", "root", "keep"], + developer_context("context", "keep", "keep me", 0), + )] + ); +} diff --git a/codex-rs/protocol/src/journal.rs b/codex-rs/protocol/src/journal.rs new file mode 100644 index 0000000000..88388f2601 --- /dev/null +++ b/codex-rs/protocol/src/journal.rs @@ -0,0 +1,383 @@ +use crate::models::ContentItem; +use crate::models::ResponseItem; +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; +use ts_rs::TS; +use uuid::Uuid; + +/// Stable key attached to a journal entry. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema, TS, Default)] +pub struct JournalKey { + pub parts: Vec, +} + +impl JournalKey { + /// Creates a key from an ordered sequence of parts. + pub fn new(parts: I) -> Self + where + I: IntoIterator, + S: Into, + { + Self { + parts: parts.into_iter().map(Into::into).collect(), + } + } + + /// Returns the ordered key parts. + pub fn parts(&self) -> &[String] { + &self.parts + } + + /// Returns a new key with one child part appended. + pub fn child(&self, part: impl Into) -> Self { + let mut parts = self.parts.clone(); + parts.push(part.into()); + Self { parts } + } + + /// Returns whether this key starts with the provided prefix. + pub fn starts_with(&self, prefix: &JournalKey) -> bool { + self.parts.starts_with(prefix.parts.as_slice()) + } +} + +impl From<[S; N]> for JournalKey +where + S: Into, +{ + fn from(value: [S; N]) -> Self { + Self::new(value) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] +#[serde(tag = "type", content = "value", rename_all = "snake_case")] +#[ts(tag = "type", content = "value", rename_all = "snake_case")] +pub enum KeyFilter { + Exact(JournalKey), + Prefix(JournalKey), +} + +impl KeyFilter { + /// Matches one exact key. + pub fn exact(key: impl Into) -> Self { + Self::Exact(key.into()) + } + + /// Matches any key with the provided prefix. + pub fn prefix(prefix: impl Into) -> Self { + Self::Prefix(prefix.into()) + } + + /// Returns whether the filter matches the key. + pub fn matches(&self, key: &JournalKey) -> bool { + match self { + Self::Exact(expected) => key == expected, + Self::Prefix(prefix) => key.starts_with(prefix), + } + } +} + +/// Stable identity for a prompt-context entry. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema, TS)] +pub struct JournalContextKey { + pub namespace: String, + pub name: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub instance: Option, +} + +impl JournalContextKey { + pub fn new( + namespace: impl Into, + name: impl Into, + instance: Option, + ) -> Self { + Self { + namespace: namespace.into(), + name: name.into(), + instance, + } + } +} + +/// Minimal message block that can be projected into a model-visible prompt item. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, TS)] +pub struct PromptMessage { + pub role: PromptMessageRole, + pub content: Vec, +} + +impl PromptMessage { + pub fn new(role: PromptMessageRole, content: Vec) -> Self { + Self { role, content } + } + + pub fn developer_text(text: impl Into) -> Self { + Self::text(PromptMessageRole::Developer, text) + } + + pub fn user_text(text: impl Into) -> Self { + Self::text(PromptMessageRole::User, text) + } + + pub fn assistant_text(text: impl Into) -> Self { + Self::text(PromptMessageRole::Assistant, text) + } + + fn text(role: PromptMessageRole, text: impl Into) -> Self { + let text = text.into(); + let content = match role { + PromptMessageRole::Developer | PromptMessageRole::User => { + vec![ContentItem::InputText { text }] + } + PromptMessageRole::Assistant => vec![ContentItem::OutputText { text }], + }; + Self::new(role, content) + } +} + +impl From for ResponseItem { + fn from(value: PromptMessage) -> Self { + ResponseItem::Message { + id: None, + role: value.role.as_str().to_string(), + content: value.content, + phase: None, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS, Default)] +#[serde(rename_all = "snake_case")] +#[ts(rename_all = "snake_case")] +pub enum PromptMessageRole { + #[default] + Developer, + User, + Assistant, +} + +impl PromptMessageRole { + pub fn as_str(self) -> &'static str { + match self { + Self::Developer => "developer", + Self::User => "user", + Self::Assistant => "assistant", + } + } +} + +/// Durable history item. Unlike prompt context, history keeps original ordering. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, TS)] +pub struct JournalHistoryItem { + pub id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub turn_id: Option, + pub item: ResponseItem, +} + +impl JournalHistoryItem { + pub fn new(item: ResponseItem) -> Self { + Self { + id: Uuid::new_v4().to_string(), + turn_id: None, + item, + } + } + + pub fn with_turn_id(mut self, turn_id: impl Into) -> Self { + self.turn_id = Some(turn_id.into()); + self + } +} + +impl From for JournalHistoryItem { + fn from(value: ResponseItem) -> Self { + Self::new(value) + } +} + +impl From for ResponseItem { + fn from(value: JournalHistoryItem) -> Self { + value.item + } +} + +impl From for ResponseItem { + fn from(value: JournalContextItem) -> Self { + value.message.into() + } +} + +/// Prompt-context entry with stable identity and filtering metadata. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, TS)] +pub struct JournalContextItem { + pub key: JournalContextKey, + pub message: PromptMessage, + #[serde(default)] + pub prompt_order: i64, + #[serde(default)] + pub audience: JournalContextAudience, + #[serde(default)] + pub on_fork: JournalContextForkBehavior, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub tags: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub source: Option, +} + +impl JournalContextItem { + pub fn new(key: JournalContextKey, message: PromptMessage) -> Self { + Self { + key, + message, + prompt_order: 0, + audience: JournalContextAudience::default(), + on_fork: JournalContextForkBehavior::default(), + tags: Vec::new(), + source: None, + } + } + + pub fn with_prompt_order(mut self, prompt_order: i64) -> Self { + self.prompt_order = prompt_order; + self + } + + pub fn with_audience(mut self, audience: JournalContextAudience) -> Self { + self.audience = audience; + self + } + + pub fn with_on_fork(mut self, on_fork: JournalContextForkBehavior) -> Self { + self.on_fork = on_fork; + self + } + + pub fn with_tags(mut self, tags: Vec) -> Self { + self.tags = tags; + self + } + + pub fn with_source(mut self, source: impl Into) -> Self { + self.source = Some(source.into()); + self + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS, Default)] +#[serde(tag = "type", content = "value", rename_all = "snake_case")] +#[ts(tag = "type", content = "value", rename_all = "snake_case")] +pub enum JournalContextAudience { + #[default] + All, + RootOnly, + SubAgentsOnly, + AgentPathPrefix(String), + AgentRole(String), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS, Default)] +#[serde(rename_all = "snake_case")] +#[ts(rename_all = "snake_case")] +pub enum JournalContextForkBehavior { + #[default] + Keep, + Drop, + Regenerate, +} + +/// Cursor into the effective history at the point a checkpoint is applied. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema, TS)] +#[serde(tag = "type", content = "value", rename_all = "snake_case")] +#[ts(tag = "type", content = "value", rename_all = "snake_case")] +pub enum JournalHistoryCursor { + Start, + AfterItem(String), + End, +} + +/// Replace the current history prefix through the resolved cursor. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, TS)] +pub struct JournalReplacePrefixCheckpoint { + pub through: JournalHistoryCursor, + pub replacement: Vec, +} + +/// Keep only the current history prefix through the resolved cursor. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, TS)] +pub struct JournalTruncateHistoryCheckpoint { + pub through: JournalHistoryCursor, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, TS)] +#[serde(tag = "type", content = "payload", rename_all = "snake_case")] +#[ts(tag = "type", content = "payload", rename_all = "snake_case")] +pub enum JournalCheckpointItem { + ReplacePrefix(JournalReplacePrefixCheckpoint), + TruncateHistory(JournalTruncateHistoryCheckpoint), +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, TS)] +#[serde(tag = "type", content = "payload", rename_all = "snake_case")] +#[ts(tag = "type", content = "payload", rename_all = "snake_case")] +pub enum JournalItem { + History(JournalHistoryItem), + Context(JournalContextItem), + Checkpoint(JournalCheckpointItem), +} + +impl From for JournalItem { + fn from(value: ResponseItem) -> Self { + Self::History(value.into()) + } +} + +impl From for JournalItem { + fn from(value: JournalHistoryItem) -> Self { + Self::History(value) + } +} + +impl From for JournalItem { + fn from(value: JournalContextItem) -> Self { + Self::Context(value) + } +} + +impl From for JournalItem { + fn from(value: JournalCheckpointItem) -> Self { + Self::Checkpoint(value) + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, TS)] +pub struct JournalEntry { + pub key: JournalKey, + pub item: JournalItem, +} + +impl JournalEntry { + /// Creates a keyed journal entry. + pub fn new(key: impl Into, item: impl Into) -> Self { + Self { + key: key.into(), + item: item.into(), + } + } +} + +impl From<(K, T)> for JournalEntry +where + K: Into, + T: Into, +{ + fn from(value: (K, T)) -> Self { + Self::new(value.0, value.1) + } +} diff --git a/codex-rs/protocol/src/lib.rs b/codex-rs/protocol/src/lib.rs index 175c92331f..fcd58a4d72 100644 --- a/codex-rs/protocol/src/lib.rs +++ b/codex-rs/protocol/src/lib.rs @@ -12,6 +12,7 @@ pub mod dynamic_tools; pub mod error; pub mod exec_output; pub mod items; +pub mod journal; pub mod mcp; pub mod memory_citation; pub mod message_history; diff --git a/codex-rs/thread-store/Cargo.toml b/codex-rs/thread-store/Cargo.toml index 3a0428f20e..7e665c0a42 100644 --- a/codex-rs/thread-store/Cargo.toml +++ b/codex-rs/thread-store/Cargo.toml @@ -22,11 +22,12 @@ codex-git-utils = { workspace = true } codex-protocol = { workspace = true } codex-rollout = { workspace = true } codex-state = { workspace = true } +codex-journal = { workspace = true } prost = "0.14.3" serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } thiserror = { workspace = true } -tokio = { workspace = true } +tokio = { workspace = true, features = ["fs", "io-util", "rt", "sync"] } tonic = { workspace = true } tonic-prost = { workspace = true } tracing = { workspace = true } diff --git a/codex-rs/thread-store/src/error.rs b/codex-rs/thread-store/src/error.rs index c5cee9a8b8..e6aab717d1 100644 --- a/codex-rs/thread-store/src/error.rs +++ b/codex-rs/thread-store/src/error.rs @@ -4,7 +4,7 @@ use codex_protocol::ThreadId; pub type ThreadStoreResult = Result; /// Error type shared by thread-store implementations. -#[derive(Debug, thiserror::Error)] +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum ThreadStoreError { /// The requested thread does not exist in this store. #[error("thread {thread_id} not found")] diff --git a/codex-rs/thread-store/src/journal_writer.rs b/codex-rs/thread-store/src/journal_writer.rs new file mode 100644 index 0000000000..8eb0546d62 --- /dev/null +++ b/codex-rs/thread-store/src/journal_writer.rs @@ -0,0 +1,355 @@ +use std::path::Path; +use std::path::PathBuf; + +use codex_journal::JournalEntry; +use tokio::fs::OpenOptions; +use tokio::io::AsyncWriteExt; +use tokio::io::BufWriter; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio::task::JoinHandle; + +use crate::ThreadStoreError; +use crate::ThreadStoreResult; + +const DEFAULT_QUEUE_CAPACITY: usize = 1024; + +/// Async append-only JSONL journal for [`JournalEntry`] items. +/// +/// Callers update the in-memory journal first, then enqueue the same keyed entries here for durable +/// persistence. Writes happen on a dedicated worker task; use [`Self::flush`] or +/// [`Self::shutdown`] when durability matters. +pub struct JournalWriter { + path: PathBuf, + tx: mpsc::Sender, + worker: Option>>, +} + +impl JournalWriter { + /// Opens or creates a journal at `path`. + pub async fn open(path: impl Into) -> ThreadStoreResult { + Self::open_with_capacity(path, DEFAULT_QUEUE_CAPACITY).await + } + + /// Opens or creates a journal at `path` with a custom queue capacity. + pub async fn open_with_capacity( + path: impl Into, + queue_capacity: usize, + ) -> ThreadStoreResult { + if queue_capacity == 0 { + return Err(ThreadStoreError::InvalidRequest { + message: "journal queue capacity must be greater than zero".to_string(), + }); + } + + let path = path.into(); + if let Some(parent) = path.parent() { + tokio::fs::create_dir_all(parent).await.map_err(io_error)?; + } + + let file = OpenOptions::new() + .create(true) + .append(true) + .open(&path) + .await + .map_err(io_error)?; + let writer = BufWriter::new(file); + let (tx, rx) = mpsc::channel(queue_capacity); + let worker_path = path.clone(); + let worker = tokio::spawn(async move { run_worker(worker_path, writer, rx).await }); + + Ok(Self { + path, + tx, + worker: Some(worker), + }) + } + + /// Returns the backing journal path. + pub fn path(&self) -> &Path { + self.path.as_path() + } + + /// Enqueues one entry for async persistence. + pub async fn enqueue(&self, entry: JournalEntry) -> ThreadStoreResult<()> { + self.tx + .send(Command::Append(vec![entry])) + .await + .map_err(send_error) + } + + /// Enqueues several entries for async persistence. + pub async fn enqueue_all(&self, entries: I) -> ThreadStoreResult<()> + where + I: IntoIterator, + { + let entries = entries.into_iter().collect::>(); + if entries.is_empty() { + return Ok(()); + } + + self.tx + .send(Command::Append(entries)) + .await + .map_err(send_error) + } + + /// Waits until all previously enqueued entries are durable. + pub async fn flush(&self) -> ThreadStoreResult<()> { + let (reply_tx, reply_rx) = oneshot::channel(); + self.tx + .send(Command::Flush { reply: reply_tx }) + .await + .map_err(send_error)?; + reply_rx.await.map_err(recv_error)? + } + + /// Flushes pending entries and stops the background worker. + pub async fn shutdown(mut self) -> ThreadStoreResult<()> { + let (reply_tx, reply_rx) = oneshot::channel(); + self.tx + .send(Command::Shutdown { reply: reply_tx }) + .await + .map_err(send_error)?; + let flush_result = reply_rx.await.map_err(recv_error)?; + + let Some(worker) = self.worker.take() else { + return flush_result; + }; + let worker_result = worker.await.map_err(join_error)?; + + flush_result?; + worker_result + } +} + +enum Command { + Append(Vec), + Flush { + reply: oneshot::Sender>, + }, + Shutdown { + reply: oneshot::Sender>, + }, +} + +async fn run_worker( + path: PathBuf, + mut writer: BufWriter, + mut rx: mpsc::Receiver, +) -> ThreadStoreResult<()> { + let mut failure: Option = None; + + while let Some(command) = rx.recv().await { + match command { + Command::Append(entries) => { + if failure.is_none() + && let Err(err) = append_entries(&mut writer, entries.as_slice()).await + { + failure = Some(err); + } + } + Command::Flush { reply } => { + let result = flush_writer(&mut writer, failure.clone(), path.as_path()).await; + if let Err(err) = &result { + failure = Some(err.clone()); + } + let _ = reply.send(result); + } + Command::Shutdown { reply } => { + let result = flush_writer(&mut writer, failure.clone(), path.as_path()).await; + let worker_result = result.clone(); + let _ = reply.send(result); + return worker_result; + } + } + } + + flush_writer(&mut writer, failure, path.as_path()).await +} + +async fn append_entries( + writer: &mut BufWriter, + entries: &[JournalEntry], +) -> ThreadStoreResult<()> { + for entry in entries { + let mut line = serde_json::to_vec(entry).map_err(|source| ThreadStoreError::Internal { + message: format!("failed to serialize journal entry: {source}"), + })?; + line.push(b'\n'); + writer.write_all(line.as_slice()).await.map_err(io_error)?; + } + Ok(()) +} + +async fn flush_writer( + writer: &mut BufWriter, + failure: Option, + path: &Path, +) -> ThreadStoreResult<()> { + if let Some(err) = failure { + return Err(err); + } + + writer.flush().await.map_err(io_error)?; + writer + .get_mut() + .sync_data() + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to sync journal at {}: {err}", path.display()), + })?; + Ok(()) +} + +fn io_error(err: std::io::Error) -> ThreadStoreError { + ThreadStoreError::Internal { + message: err.to_string(), + } +} + +fn send_error(_: mpsc::error::SendError) -> ThreadStoreError { + ThreadStoreError::Internal { + message: "journal worker is not available".to_string(), + } +} + +fn recv_error(_: oneshot::error::RecvError) -> ThreadStoreError { + ThreadStoreError::Internal { + message: "journal worker dropped its response channel".to_string(), + } +} + +fn join_error(err: tokio::task::JoinError) -> ThreadStoreError { + ThreadStoreError::Internal { + message: format!("journal worker task failed: {err}"), + } +} + +#[cfg(test)] +mod tests { + use codex_journal::Journal; + use codex_journal::JournalContextItem; + use codex_journal::JournalContextKey; + use codex_journal::JournalHistoryItem; + use codex_journal::PromptMessage; + use codex_protocol::models::ContentItem; + use codex_protocol::models::ResponseItem; + use pretty_assertions::assert_eq; + use tempfile::tempdir; + + use super::JournalWriter; + + fn developer_entry(text: &str) -> codex_journal::JournalEntry { + codex_journal::JournalEntry::new( + ["prompt", text], + JournalContextItem::new( + JournalContextKey::new("context", text, None), + PromptMessage::developer_text(text), + ), + ) + } + + fn user_entry(text: &str) -> codex_journal::JournalEntry { + codex_journal::JournalEntry::new( + ["history", text], + JournalHistoryItem { + id: format!("history-{text}"), + turn_id: None, + item: ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: text.to_string(), + }], + phase: None, + }, + }, + ) + } + + #[tokio::test] + async fn enqueue_and_flush_persists_entries_in_order() { + let dir = tempdir().expect("tempdir"); + let path = dir.path().join("journal.jsonl"); + let journal = JournalWriter::open(path.clone()) + .await + .expect("journal should open"); + + journal + .enqueue(developer_entry("first")) + .await + .expect("first entry should enqueue"); + journal + .enqueue(user_entry("hello")) + .await + .expect("second entry should enqueue"); + journal.flush().await.expect("journal should flush"); + + let loaded = Journal::load_jsonl(path.as_path()).expect("journal should load"); + + assert_eq!( + loaded.entries(), + vec![developer_entry("first"), user_entry("hello")] + ); + } + + #[tokio::test] + async fn enqueue_all_and_shutdown_persist_entries() { + let dir = tempdir().expect("tempdir"); + let path = dir.path().join("journal.jsonl"); + let journal = JournalWriter::open(path.clone()) + .await + .expect("journal should open"); + + journal + .enqueue_all(vec![developer_entry("one"), developer_entry("two")]) + .await + .expect("entries should enqueue"); + journal.shutdown().await.expect("journal should shut down"); + + let loaded = Journal::load_jsonl(path.as_path()).expect("journal should load"); + + assert_eq!( + loaded.entries(), + vec![developer_entry("one"), developer_entry("two")] + ); + } + + #[tokio::test] + async fn reopen_appends_to_existing_journal() { + let dir = tempdir().expect("tempdir"); + let path = dir.path().join("journal.jsonl"); + + let first = JournalWriter::open(path.clone()) + .await + .expect("first journal should open"); + first + .enqueue(developer_entry("first")) + .await + .expect("entry should enqueue"); + first + .shutdown() + .await + .expect("first journal should shut down"); + + let second = JournalWriter::open(path.clone()) + .await + .expect("second journal should open"); + second + .enqueue(user_entry("second")) + .await + .expect("entry should enqueue"); + second + .shutdown() + .await + .expect("second journal should shut down"); + + let loaded = Journal::load_jsonl(path.as_path()).expect("journal should load"); + + assert_eq!( + loaded.entries(), + vec![developer_entry("first"), user_entry("second")] + ); + } +} diff --git a/codex-rs/thread-store/src/lib.rs b/codex-rs/thread-store/src/lib.rs index 42b9297bca..b5cd7b97a2 100644 --- a/codex-rs/thread-store/src/lib.rs +++ b/codex-rs/thread-store/src/lib.rs @@ -7,6 +7,7 @@ mod error; #[cfg(debug_assertions)] mod in_memory; +mod journal_writer; mod live_thread; mod local; mod remote; @@ -19,6 +20,7 @@ pub use error::ThreadStoreResult; pub use in_memory::InMemoryThreadStore; #[cfg(debug_assertions)] pub use in_memory::InMemoryThreadStoreCalls; +pub use journal_writer::JournalWriter; pub use live_thread::LiveThread; pub use live_thread::LiveThreadInitGuard; pub use local::LocalThreadStore;