diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 22b34b1cd0..10dcb27d27 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1,4 +1,3 @@ -use std::borrow::Cow; use std::fmt::Debug; use std::path::PathBuf; use std::sync::Arc; @@ -533,7 +532,8 @@ impl Session { InitialHistory::New => { // Build and record initial items (user instructions + environment context) let items = self.build_initial_context(turn_context); - self.record_conversation_items(&items).await; + self.record_conversation_items(&items, TaskKind::Regular) + .await; } InitialHistory::Resumed(_) | InitialHistory::Forked(_) => { let rollout_items = conversation_history.get_rollout_items(); @@ -543,7 +543,8 @@ impl Session { let reconstructed_history = self.reconstruct_history_from_rollout(turn_context, &rollout_items); if !reconstructed_history.is_empty() { - self.record_into_history(&reconstructed_history).await; + self.record_into_history(&reconstructed_history, TaskKind::Regular) + .await; } // If persisting, persist all rollout items as-is (recorder filters) @@ -668,9 +669,21 @@ impl Session { /// Records input items: always append to conversation history and /// persist these response items to rollout. - async fn record_conversation_items(&self, items: &[ResponseItem]) { - self.record_into_history(items).await; - self.persist_rollout_response_items(items).await; + async fn record_conversation_items(&self, items: &[ResponseItem], task_kind: TaskKind) { + match task_kind { + TaskKind::Regular | TaskKind::Compact => { + self.record_into_history(items, task_kind).await; + self.persist_rollout_response_items(items).await; + } + TaskKind::Review => { + self.record_into_history(items, task_kind).await; + } + } + } + + async fn clear_review_thread(&self) { + let mut state = self.state.lock().await; + state.clear_review_thread(); } fn reconstruct_history_from_rollout( @@ -682,7 +695,7 @@ impl Session { for item in rollout_items { match item { RolloutItem::ResponseItem(response_item) => { - history.record_items(std::iter::once(response_item)); + history.record_items(std::iter::once(response_item.clone()), TaskKind::Regular); } RolloutItem::Compacted(compacted) => { let snapshot = history.contents(); @@ -701,9 +714,9 @@ impl Session { } /// Append ResponseItems to the in-memory conversation history only. - async fn record_into_history(&self, items: &[ResponseItem]) { + async fn record_into_history(&self, items: &[ResponseItem], task_kind: TaskKind) { let mut state = self.state.lock().await; - state.record_items(items.iter()); + state.record_items(items.iter().cloned(), task_kind); } async fn replace_history(&self, items: Vec) { @@ -800,13 +813,9 @@ impl Session { } } - /// Record a user input item to conversation history and also persist a - /// corresponding UserMessage EventMsg to rollout. - async fn record_input_and_rollout_usermsg(&self, response_input: &ResponseInputItem) { + /// persist a corresponding UserMessage EventMsg to rollout. + async fn persist_user_msg_to_rollout(&self, response_input: &ResponseInputItem) { let response_item: ResponseItem = response_input.clone().into(); - // Add to conversation history and persist response item to rollout - self.record_conversation_items(std::slice::from_ref(&response_item)) - .await; // Derive user message events and persist only UserMessage to rollout let msgs = @@ -823,6 +832,38 @@ impl Session { } } + async fn initialize_task_history( + &self, + response_input: &ResponseInputItem, + task_kind: TaskKind, + initial_context: Vec, + ) { + match task_kind { + TaskKind::Regular | TaskKind::Compact => { + let response_item: ResponseItem = response_input.clone().into(); + self.record_conversation_items(std::slice::from_ref(&response_item), task_kind) + .await; + self.persist_user_msg_to_rollout(response_input).await; + } + TaskKind::Review => { + let mut state = self.state.lock().await; + state.initialize_review_history(response_input, initial_context); + } + } + } + + async fn prepare_prompt_input( + &self, + pending_input: Vec, + task_kind: TaskKind, + ) -> Vec { + if !pending_input.is_empty() && matches!(task_kind, TaskKind::Regular | TaskKind::Compact) { + self.persist_rollout_response_items(&pending_input).await; + } + let mut state = self.state.lock().await; + state.prepare_prompt_input(task_kind, pending_input) + } + async fn on_exec_command_begin( &self, turn_diff_tracker: SharedTurnDiffTracker, @@ -1222,13 +1263,16 @@ async fn submission_loop( // Optionally persist changes to model / effort if cwd.is_some() || approval_policy.is_some() || sandbox_policy.is_some() { - sess.record_conversation_items(&[ResponseItem::from(EnvironmentContext::new( - cwd, - approval_policy, - sandbox_policy, - // Shell is not configurable from turn to turn - None, - ))]) + sess.record_conversation_items( + &[ResponseItem::from(EnvironmentContext::new( + cwd, + approval_policy, + sandbox_policy, + // Shell is not configurable from turn to turn + None, + ))], + TaskKind::Regular, + ) .await; } } @@ -1322,8 +1366,11 @@ async fn submission_loop( let new_env_context = EnvironmentContext::from(&fresh_turn_context); if !new_env_context.equals_except_shell(&previous_env_context) { let env_response_item = ResponseItem::from(new_env_context); - sess.record_conversation_items(std::slice::from_ref(&env_response_item)) - .await; + sess.record_conversation_items( + std::slice::from_ref(&env_response_item), + TaskKind::Regular, + ) + .await; for msg in map_response_item_to_event_messages( &env_response_item, sess.show_raw_agent_reasoning(), @@ -1649,19 +1696,9 @@ pub(crate) async fn run_task( sess.send_event(event).await; let initial_input_for_turn: ResponseInputItem = ResponseInputItem::from(input); - // For review threads, keep an isolated in-memory history so the - // model sees a fresh conversation without the parent session's history. - // For normal turns, continue recording to the session history as before. - let is_review_mode = turn_context.is_review_mode; - let mut review_thread_history: Vec = Vec::new(); - if is_review_mode { - // Seed review threads with environment context so the model knows the working directory. - review_thread_history.extend(sess.build_initial_context(turn_context.as_ref())); - review_thread_history.push(initial_input_for_turn.into()); - } else { - sess.record_input_and_rollout_usermsg(&initial_input_for_turn) - .await; - } + let initial_context = sess.build_initial_context(turn_context.as_ref()); + sess.initialize_task_history(&initial_input_for_turn, task_kind, initial_context) + .await; let mut last_agent_message: Option = None; // Although from the perspective of codex.rs, TurnDiffTracker has the lifecycle of a Task which contains @@ -1680,27 +1717,9 @@ pub(crate) async fn run_task( .map(ResponseItem::from) .collect::>(); - // Construct the input that we will send to the model. - // - // - For review threads, use the isolated in-memory history so the - // model sees a fresh conversation (no parent history/user_instructions). - // - // - For normal turns, use the session's full history. When using the - // chat completions API (or ZDR clients), the model needs the full - // conversation history on each turn. The rollout file, however, should - // only record the new items that originated in this turn so that it - // represents an append-only log without duplicates. - let turn_input: Vec = if is_review_mode { - if !pending_input.is_empty() { - review_thread_history.extend(pending_input); - } - review_thread_history.clone() - } else { - sess.record_conversation_items(&pending_input).await; - sess.turn_input_with_history(pending_input).await - }; + let prompt_input = sess.prepare_prompt_input(pending_input, task_kind).await; - let turn_input_messages: Vec = turn_input + let turn_input_messages: Vec = prompt_input .iter() .filter_map(|item| match item { ResponseItem::Message { content, .. } => Some(content), @@ -1718,7 +1737,7 @@ pub(crate) async fn run_task( Arc::clone(&turn_context), Arc::clone(&turn_diff_tracker), sub_id.clone(), - turn_input, + prompt_input, task_kind, ) .await @@ -1833,13 +1852,11 @@ pub(crate) async fn run_task( // Only attempt to take the lock if there is something to record. if !items_to_record_in_conversation_history.is_empty() { - if is_review_mode { - review_thread_history - .extend(items_to_record_in_conversation_history.clone()); - } else { - sess.record_conversation_items(&items_to_record_in_conversation_history) - .await; - } + sess.record_conversation_items( + &items_to_record_in_conversation_history, + task_kind, + ) + .await; } if token_limit_reached { @@ -2051,61 +2068,6 @@ async fn try_run_turn( prompt: &Prompt, task_kind: TaskKind, ) -> CodexResult { - // call_ids that are part of this response. - let completed_call_ids = prompt - .input - .iter() - .filter_map(|ri| match ri { - ResponseItem::FunctionCallOutput { call_id, .. } => Some(call_id), - ResponseItem::LocalShellCall { - call_id: Some(call_id), - .. - } => Some(call_id), - ResponseItem::CustomToolCallOutput { call_id, .. } => Some(call_id), - _ => None, - }) - .collect::>(); - - // call_ids that were pending but are not part of this response. - // This usually happens because the user interrupted the model before we responded to one of its tool calls - // and then the user sent a follow-up message. - let missing_calls = { - prompt - .input - .iter() - .filter_map(|ri| match ri { - ResponseItem::FunctionCall { call_id, .. } => Some(call_id), - ResponseItem::LocalShellCall { - call_id: Some(call_id), - .. - } => Some(call_id), - ResponseItem::CustomToolCall { call_id, .. } => Some(call_id), - _ => None, - }) - .filter_map(|call_id| { - if completed_call_ids.contains(&call_id) { - None - } else { - Some(call_id.clone()) - } - }) - .map(|call_id| ResponseItem::CustomToolCallOutput { - call_id, - output: "aborted".to_string(), - }) - .collect::>() - }; - let prompt: Cow = if missing_calls.is_empty() { - Cow::Borrowed(prompt) - } else { - // Add the synthetic aborted missing calls to the beginning of the input to ensure all call ids have responses. - let input = [missing_calls, prompt.input.clone()].concat(); - Cow::Owned(Prompt { - input, - ..prompt.clone() - }) - }; - let rollout_item = RolloutItem::TurnContext(TurnContextItem { cwd: turn_context.cwd.clone(), approval_policy: turn_context.approval_policy, @@ -2118,7 +2080,7 @@ async fn try_run_turn( let mut stream = turn_context .client .clone() - .stream_with_task_kind(prompt.as_ref(), task_kind) + .stream_with_task_kind(prompt, task_kind) .await?; let tool_runtime = ToolCallRuntime::new( @@ -2440,12 +2402,16 @@ pub(crate) async fn exit_review_mode( } session - .record_conversation_items(&[ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ContentItem::InputText { text: user_message }], - }]) + .record_conversation_items( + &[ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { text: user_message }], + }], + TaskKind::Regular, + ) .await; + session.clear_review_thread().await; } use crate::executor::errors::ExecError; @@ -3025,7 +2991,7 @@ mod tests { for item in &initial_context { rollout_items.push(RolloutItem::ResponseItem(item.clone())); } - live_history.record_items(initial_context.iter()); + live_history.record_items(initial_context.iter().cloned(), TaskKind::Regular); let user1 = ResponseItem::Message { id: None, @@ -3034,8 +3000,8 @@ mod tests { text: "first user".to_string(), }], }; - live_history.record_items(std::iter::once(&user1)); - rollout_items.push(RolloutItem::ResponseItem(user1.clone())); + live_history.record_items(std::iter::once(user1.clone()), TaskKind::Regular); + rollout_items.push(RolloutItem::ResponseItem(user1)); let assistant1 = ResponseItem::Message { id: None, @@ -3044,8 +3010,8 @@ mod tests { text: "assistant reply one".to_string(), }], }; - live_history.record_items(std::iter::once(&assistant1)); - rollout_items.push(RolloutItem::ResponseItem(assistant1.clone())); + live_history.record_items(std::iter::once(assistant1.clone()), TaskKind::Regular); + rollout_items.push(RolloutItem::ResponseItem(assistant1)); let summary1 = "summary one"; let snapshot1 = live_history.contents(); @@ -3067,8 +3033,8 @@ mod tests { text: "second user".to_string(), }], }; - live_history.record_items(std::iter::once(&user2)); - rollout_items.push(RolloutItem::ResponseItem(user2.clone())); + live_history.record_items(std::iter::once(user2.clone()), TaskKind::Regular); + rollout_items.push(RolloutItem::ResponseItem(user2)); let assistant2 = ResponseItem::Message { id: None, @@ -3077,8 +3043,8 @@ mod tests { text: "assistant reply two".to_string(), }], }; - live_history.record_items(std::iter::once(&assistant2)); - rollout_items.push(RolloutItem::ResponseItem(assistant2.clone())); + live_history.record_items(std::iter::once(assistant2.clone()), TaskKind::Regular); + rollout_items.push(RolloutItem::ResponseItem(assistant2)); let summary2 = "summary two"; let snapshot2 = live_history.contents(); @@ -3100,8 +3066,8 @@ mod tests { text: "third user".to_string(), }], }; - live_history.record_items(std::iter::once(&user3)); - rollout_items.push(RolloutItem::ResponseItem(user3.clone())); + live_history.record_items(std::iter::once(user3.clone()), TaskKind::Regular); + rollout_items.push(RolloutItem::ResponseItem(user3)); let assistant3 = ResponseItem::Message { id: None, @@ -3110,8 +3076,8 @@ mod tests { text: "assistant reply three".to_string(), }], }; - live_history.record_items(std::iter::once(&assistant3)); - rollout_items.push(RolloutItem::ResponseItem(assistant3.clone())); + live_history.record_items(std::iter::once(assistant3.clone()), TaskKind::Regular); + rollout_items.push(RolloutItem::ResponseItem(assistant3)); (rollout_items, live_history.contents()) } diff --git a/codex-rs/core/src/codex/compact.rs b/codex-rs/core/src/codex/compact.rs index 93bbfa79c6..9a8921d74d 100644 --- a/codex-rs/core/src/codex/compact.rs +++ b/codex-rs/core/src/codex/compact.rs @@ -274,7 +274,8 @@ async fn drain_to_completed( }; match event { Ok(ResponseEvent::OutputItemDone(item)) => { - sess.record_into_history(std::slice::from_ref(&item)).await; + sess.record_into_history(std::slice::from_ref(&item), TaskKind::Compact) + .await; } Ok(ResponseEvent::RateLimits(snapshot)) => { sess.update_rate_limits(sub_id, snapshot).await; diff --git a/codex-rs/core/src/conversation_history.rs b/codex-rs/core/src/conversation_history.rs index 7c23e4dc29..3d45006997 100644 --- a/codex-rs/core/src/conversation_history.rs +++ b/codex-rs/core/src/conversation_history.rs @@ -1,15 +1,22 @@ +use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; +use crate::state::TaskKind; + /// Transcript of conversation history #[derive(Debug, Clone, Default)] pub(crate) struct ConversationHistory { /// The oldest items are at the beginning of the vector. items: Vec, + review_thread_history: Vec, } impl ConversationHistory { pub(crate) fn new() -> Self { - Self { items: Vec::new() } + Self { + items: Vec::new(), + review_thread_history: Vec::new(), + } } /// Returns a clone of the contents in the transcript. @@ -17,24 +24,113 @@ impl ConversationHistory { self.items.clone() } + pub(crate) fn review_thread_contents(&self) -> Vec { + self.review_thread_history.clone() + } + + pub(crate) fn clear_review_thread(&mut self) { + self.review_thread_history.clear(); + } + /// `items` is ordered from oldest to newest. - pub(crate) fn record_items(&mut self, items: I) + pub(crate) fn record_items(&mut self, items: I, task_kind: TaskKind) where - I: IntoIterator, - I::Item: std::ops::Deref, + I: IntoIterator, { for item in items { if !is_api_message(&item) { continue; } - self.items.push(item.clone()); + match task_kind { + TaskKind::Regular | TaskKind::Compact => { + self.items.push(item); + } + TaskKind::Review => { + self.review_thread_history.push(item); + } + } } } pub(crate) fn replace(&mut self, items: Vec) { self.items = items; } + + pub(crate) fn initialize_review_history( + &mut self, + response_input: &ResponseInputItem, + initial_context: Vec, + ) { + self.clear_review_thread(); + self.record_items(initial_context, TaskKind::Review); + self.record_items( + std::iter::once(ResponseItem::from(response_input.clone())), + TaskKind::Review, + ); + } + + pub(crate) fn add_pending_input( + &mut self, + pending_input: Vec, + task_kind: TaskKind, + ) { + self.record_items(pending_input, task_kind); + } + + pub(crate) fn handle_missing_tool_call_output(&mut self, task_kind: TaskKind) { + // call_ids that are part of this response. + let content = match task_kind { + TaskKind::Regular => self.contents(), + TaskKind::Review => self.review_thread_contents(), + TaskKind::Compact => self.contents(), + }; + let completed_call_ids = content + .iter() + .filter_map(|ri| match ri { + ResponseItem::FunctionCallOutput { call_id, .. } => Some(call_id), + ResponseItem::CustomToolCallOutput { call_id, .. } => Some(call_id), + _ => None, + }) + .collect::>(); + + // call_ids that were pending but are not part of this response. + // This usually happens because the user interrupted the model before we responded to one of its tool calls + // and then the user sent a follow-up message. + let missing_calls = { + content + .iter() + .filter_map(|ri| match ri { + ResponseItem::FunctionCall { call_id, .. } => Some(call_id), + ResponseItem::LocalShellCall { + call_id: Some(call_id), + .. + } => Some(call_id), + ResponseItem::CustomToolCall { call_id, .. } => Some(call_id), + _ => None, + }) + .filter_map(|call_id| { + if completed_call_ids.contains(&call_id) { + None + } else { + Some(call_id.clone()) + } + }) + .map(|call_id| ResponseItem::CustomToolCallOutput { + call_id, + output: "aborted".to_string(), + }) + .collect::>() + }; + self.record_items(missing_calls, task_kind); + } + + pub(crate) fn prompt(&self, task_kind: TaskKind) -> Vec { + match task_kind { + TaskKind::Regular | TaskKind::Compact => self.contents(), + TaskKind::Review => self.review_thread_contents(), + } + } } /// Anything that is not a system message or "reasoning" message is considered @@ -89,12 +185,12 @@ mod tests { text: "ignored".to_string(), }], }; - h.record_items([&system, &ResponseItem::Other]); + h.record_items([system, ResponseItem::Other], TaskKind::Regular); // User and assistant should be retained. let u = user_msg("hi"); let a = assistant_msg("hello"); - h.record_items([&u, &a]); + h.record_items([u, a], TaskKind::Regular); let items = h.contents(); assert_eq!( @@ -114,7 +210,7 @@ mod tests { text: "hello".to_string() }] } - ] + ], ); } } diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 8310d91c0c..ff3b5fe094 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -1,11 +1,13 @@ //! Session-wide mutable state. +use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; use crate::conversation_history::ConversationHistory; use crate::protocol::RateLimitSnapshot; use crate::protocol::TokenUsage; use crate::protocol::TokenUsageInfo; +use crate::state::TaskKind; /// Persistent, session-scoped state previously stored directly on `Session`. #[derive(Default)] @@ -25,12 +27,11 @@ impl SessionState { } // History helpers - pub(crate) fn record_items(&mut self, items: I) + pub(crate) fn record_items(&mut self, items: I, task_kind: TaskKind) where - I: IntoIterator, - I::Item: std::ops::Deref, + I: IntoIterator, { - self.history.record_items(items) + self.history.record_items(items, task_kind) } pub(crate) fn history_snapshot(&self) -> Vec { @@ -41,6 +42,31 @@ impl SessionState { self.history.replace(items); } + pub(crate) fn clear_review_thread(&mut self) { + self.history.clear_review_thread(); + } + + pub(crate) fn initialize_review_history( + &mut self, + response_input: &ResponseInputItem, + initial_context: Vec, + ) { + self.history + .initialize_review_history(response_input, initial_context); + } + + pub(crate) fn prepare_prompt_input( + &mut self, + task_kind: TaskKind, + pending_input: Vec, + ) -> Vec { + if !pending_input.is_empty() { + self.history.add_pending_input(pending_input, task_kind); + } + self.history.handle_missing_tool_call_output(task_kind); + self.history.prompt(task_kind) + } + // Token/rate limit helpers pub(crate) fn update_token_info_from_usage( &mut self,