From 0fd69f1b37fb3320fbc1ffc8b57775b585168ff2 Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Sun, 8 Feb 2026 20:02:58 -0800 Subject: [PATCH] Extract ContextDiscoverable history handling --- .../context_manager/context_discoverable.rs | 42 +++++++++++++ codex-rs/core/src/context_manager/history.rs | 62 +++++-------------- codex-rs/core/src/context_manager/mod.rs | 1 + 3 files changed, 58 insertions(+), 47 deletions(-) create mode 100644 codex-rs/core/src/context_manager/context_discoverable.rs diff --git a/codex-rs/core/src/context_manager/context_discoverable.rs b/codex-rs/core/src/context_manager/context_discoverable.rs new file mode 100644 index 0000000000..3689cd614f --- /dev/null +++ b/codex-rs/core/src/context_manager/context_discoverable.rs @@ -0,0 +1,42 @@ +use crate::truncate::TruncationPolicy; +use crate::truncate::truncate_function_output_items_with_policy; +use crate::truncate::truncate_text; +use codex_protocol::models::CustomToolCallOutput; +use codex_protocol::models::FunctionCallOutput; +use codex_protocol::models::FunctionCallOutputBody; +use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::ResponseItem; + +pub(super) trait ContextDiscoverable { + fn discoverable_history_item(&self, policy: TruncationPolicy) -> ResponseItem; +} + +impl ContextDiscoverable for FunctionCallOutput { + fn discoverable_history_item(&self, policy: TruncationPolicy) -> ResponseItem { + let body = match &self.output.body { + FunctionCallOutputBody::Text(content) => { + FunctionCallOutputBody::Text(truncate_text(content, policy)) + } + FunctionCallOutputBody::ContentItems(items) => FunctionCallOutputBody::ContentItems( + truncate_function_output_items_with_policy(items, policy), + ), + }; + + ResponseItem::FunctionCallOutput(FunctionCallOutput { + call_id: self.call_id.clone(), + output: FunctionCallOutputPayload { + body, + success: self.output.success, + }, + }) + } +} + +impl ContextDiscoverable for CustomToolCallOutput { + fn discoverable_history_item(&self, policy: TruncationPolicy) -> ResponseItem { + ResponseItem::CustomToolCallOutput(CustomToolCallOutput { + call_id: self.call_id.clone(), + output: truncate_text(&self.output, policy), + }) + } +} diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 597b14ae45..0b637f1fd5 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -1,4 +1,5 @@ use crate::codex::TurnContext; +use crate::context_manager::context_discoverable::ContextDiscoverable; use crate::context_manager::normalize; use crate::instructions::SkillInstructions; use crate::instructions::UserInstructions; @@ -6,14 +7,10 @@ use crate::session_prefix::is_session_prefix; use crate::truncate::TruncationPolicy; use crate::truncate::approx_token_count; use crate::truncate::approx_tokens_from_byte_count; -use crate::truncate::truncate_function_output_items_with_policy; -use crate::truncate::truncate_text; use crate::user_shell_command::is_user_shell_command_text; 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::ResponseItem; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; @@ -293,50 +290,21 @@ impl ContextManager { fn process_item(&self, item: &ResponseItem, policy: TruncationPolicy) -> ResponseItem { let policy_with_serialization_budget = policy * 1.2; match item { - ResponseItem::FunctionCallOutput(codex_protocol::models::FunctionCallOutput { - call_id, - output, - }) => { - let body = match &output.body { - FunctionCallOutputBody::Text(content) => FunctionCallOutputBody::Text( - truncate_text(content, policy_with_serialization_budget), - ), - FunctionCallOutputBody::ContentItems(items) => { - FunctionCallOutputBody::ContentItems( - truncate_function_output_items_with_policy( - items, - policy_with_serialization_budget, - ), - ) - } - }; - ResponseItem::FunctionCallOutput(codex_protocol::models::FunctionCallOutput { - call_id: call_id.clone(), - output: FunctionCallOutputPayload { - body, - success: output.success, - }, - }) - } - ResponseItem::CustomToolCallOutput(codex_protocol::models::CustomToolCallOutput { - call_id, - output, - }) => { - let truncated = truncate_text(output, policy_with_serialization_budget); - ResponseItem::CustomToolCallOutput(codex_protocol::models::CustomToolCallOutput { - call_id: call_id.clone(), - output: truncated, - }) - } - ResponseItem::Message(codex_protocol::models::Message { .. }) - | ResponseItem::Reasoning(codex_protocol::models::Reasoning { .. }) - | ResponseItem::LocalShellCall(codex_protocol::models::LocalShellCall { .. }) - | ResponseItem::FunctionCall(codex_protocol::models::FunctionCall { .. }) - | ResponseItem::WebSearchCall(codex_protocol::models::WebSearchCall { .. }) - | ResponseItem::CustomToolCall(codex_protocol::models::CustomToolCall { .. }) - | ResponseItem::Compaction(codex_protocol::models::Compaction { .. }) - | ResponseItem::GhostSnapshot(codex_protocol::models::GhostSnapshot { .. }) + ResponseItem::Message(_) + | ResponseItem::Reasoning(_) + | ResponseItem::LocalShellCall(_) + | ResponseItem::FunctionCall(_) + | ResponseItem::CustomToolCall(_) + | ResponseItem::WebSearchCall(_) + | ResponseItem::Compaction(_) + | ResponseItem::GhostSnapshot(_) | ResponseItem::Other => item.clone(), + ResponseItem::FunctionCallOutput(item) => { + item.discoverable_history_item(policy_with_serialization_budget) + } + ResponseItem::CustomToolCallOutput(item) => { + item.discoverable_history_item(policy_with_serialization_budget) + } } } } diff --git a/codex-rs/core/src/context_manager/mod.rs b/codex-rs/core/src/context_manager/mod.rs index 22e9682fe3..cdaf58fc16 100644 --- a/codex-rs/core/src/context_manager/mod.rs +++ b/codex-rs/core/src/context_manager/mod.rs @@ -1,3 +1,4 @@ +mod context_discoverable; mod history; mod normalize;