From 8f8ca17da0c12608b0f9a14f982124b96f17006e Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Thu, 23 Oct 2025 18:58:08 -0700 Subject: [PATCH] input-validation --- codex-rs/core/src/codex.rs | 146 ++++++++++++++-------- codex-rs/core/src/codex/compact.rs | 24 ++-- codex-rs/core/src/conversation_history.rs | 70 ++++++++++- codex-rs/core/src/error.rs | 3 + codex-rs/core/src/response_processing.rs | 9 +- codex-rs/core/src/state/session.rs | 6 +- codex-rs/core/src/tasks/compact.rs | 3 +- codex-rs/core/src/tasks/mod.rs | 23 +++- codex-rs/core/src/tasks/regular.rs | 3 +- codex-rs/core/src/tasks/review.rs | 5 +- codex-rs/core/tests/suite/mod.rs | 1 + 11 files changed, 208 insertions(+), 85 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 77ac35ae2a..1a3be60534 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -568,7 +568,9 @@ impl Session { // Dispatch the SessionConfiguredEvent first and then report any errors. // If resuming, include converted initial messages in the payload so UIs can render them immediately. let initial_messages = initial_history.get_event_msgs(); - sess.record_initial_history(initial_history).await; + sess.record_initial_history(initial_history) + .await + .map_err(anyhow::Error::new)?; let events = std::iter::once(Event { id: INITIAL_SUBMIT_ID.to_owned(), @@ -601,13 +603,16 @@ impl Session { format!("auto-compact-{id}") } - async fn record_initial_history(&self, conversation_history: InitialHistory) { + async fn record_initial_history( + &self, + conversation_history: InitialHistory, + ) -> CodexResult<()> { let turn_context = self.new_turn(SessionSettingsUpdate::default()).await; match conversation_history { 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).await?; } InitialHistory::Resumed(_) | InitialHistory::Forked(_) => { let rollout_items = conversation_history.get_rollout_items(); @@ -615,9 +620,9 @@ impl Session { // Always add response items to conversation history let reconstructed_history = - self.reconstruct_history_from_rollout(&turn_context, &rollout_items); + 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).await?; } // If persisting, persist all rollout items as-is (recorder filters) @@ -626,6 +631,7 @@ impl Session { } } } + Ok(()) } pub(crate) async fn update_settings(&self, updates: SessionSettingsUpdate) { @@ -856,21 +862,25 @@ impl Session { /// Records input items: always append to conversation history and /// persist these response items to rollout. - pub(crate) async fn record_conversation_items(&self, items: &[ResponseItem]) { - self.record_into_history(items).await; + pub(crate) async fn record_conversation_items( + &self, + items: &[ResponseItem], + ) -> CodexResult<()> { + self.record_into_history(items).await?; self.persist_rollout_response_items(items).await; + Ok(()) } fn reconstruct_history_from_rollout( &self, turn_context: &TurnContext, rollout_items: &[RolloutItem], - ) -> Vec { + ) -> CodexResult> { let mut history = ConversationHistory::new(); 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))?; } RolloutItem::Compacted(compacted) => { let snapshot = history.get_history(); @@ -885,13 +895,14 @@ impl Session { _ => {} } } - history.get_history() + Ok(history.get_history()) } /// 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]) -> CodexResult<()> { let mut state = self.state.lock().await; - state.record_items(items.iter()); + state.record_items(items.iter())?; + Ok(()) } async fn replace_history(&self, items: Vec) { @@ -1000,11 +1011,11 @@ impl Session { &self, turn_context: &TurnContext, response_input: &ResponseInputItem, - ) { + ) -> CodexResult<()> { 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; + .await?; // Derive user message events and persist only UserMessage to rollout let turn_item = parse_turn_item(&response_item); @@ -1013,6 +1024,7 @@ impl Session { self.emit_turn_item_started_completed(turn_context, item, false) .await; } + Ok(()) } /// Helper that emits a BackgroundEvent with the given message. This keeps @@ -1193,9 +1205,17 @@ async fn submission_loop(sess: Arc, config: Arc, rx_sub: Receiv if let Err(items) = sess.inject_input(items).await { if let Some(env_item) = sess .build_environment_update_item(previous_context.as_ref(), ¤t_context) + && let Err(err) = sess + .record_conversation_items(std::slice::from_ref(&env_item)) + .await { - sess.record_conversation_items(std::slice::from_ref(&env_item)) - .await; + sess.send_event( + current_context.as_ref(), + EventMsg::Error(ErrorEvent { + message: err.to_string(), + }), + ) + .await; } sess.spawn_task(Arc::clone(¤t_context), items, RegularTask) @@ -1508,9 +1528,9 @@ pub(crate) async fn run_task( input: Vec, task_kind: TaskKind, cancellation_token: CancellationToken, -) -> Option { +) -> CodexResult> { if input.is_empty() { - return None; + return Ok(None); } let event = EventMsg::TaskStarted(TaskStartedEvent { model_context_window: turn_context.client.get_model_context_window(), @@ -1527,11 +1547,11 @@ pub(crate) async fn run_task( if is_review_mode { // Seed review threads with environment context so the model knows the working directory. review_thread_history - .record_items(sess.build_initial_context(turn_context.as_ref()).iter()); - review_thread_history.record_items(std::iter::once(&initial_input_for_turn.into())); + .record_items(sess.build_initial_context(turn_context.as_ref()).iter())?; + review_thread_history.record_items(std::iter::once(&initial_input_for_turn.into()))?; } else { sess.record_input_and_rollout_usermsg(turn_context.as_ref(), &initial_input_for_turn) - .await; + .await?; } let mut last_agent_message: Option = None; @@ -1563,11 +1583,11 @@ pub(crate) async fn run_task( // represents an append-only log without duplicates. let turn_input: Vec = if is_review_mode { if !pending_input.is_empty() { - review_thread_history.record_items(&pending_input); + review_thread_history.record_items(&pending_input)?; } review_thread_history.get_history() } else { - sess.record_conversation_items(&pending_input).await; + sess.record_conversation_items(&pending_input).await?; sess.history_snapshot().await }; @@ -1615,7 +1635,7 @@ pub(crate) async fn run_task( &mut review_thread_history, &sess, ) - .await; + .await?; if token_limit_reached { if auto_compact_recently_attempted { @@ -1632,7 +1652,8 @@ pub(crate) async fn run_task( break; } auto_compact_recently_attempted = true; - compact::run_inline_auto_compact_task(sess.clone(), turn_context.clone()).await; + compact::run_inline_auto_compact_task(sess.clone(), turn_context.clone()) + .await?; continue; } @@ -1657,13 +1678,13 @@ pub(crate) async fn run_task( Err(CodexErr::TurnAborted { dangling_artifacts: processed_items, }) => { - let _ = process_items( + process_items( processed_items, is_review_mode, &mut review_thread_history, &sess, ) - .await; + .await?; // Aborted turn is reported via a different event. break; } @@ -1692,10 +1713,10 @@ pub(crate) async fn run_task( Arc::clone(&turn_context), last_agent_message.as_deref().map(parse_review_output_event), ) - .await; + .await?; } - last_agent_message + Ok(last_agent_message) } /// Parse the review output; when not valid JSON, build a structured @@ -2134,7 +2155,7 @@ pub(crate) async fn exit_review_mode( session: Arc, turn_context: Arc, review_output: Option, -) { +) -> CodexResult<()> { let event = EventMsg::ExitedReviewMode(ExitedReviewModeEvent { review_output: review_output.clone(), }); @@ -2177,7 +2198,8 @@ pub(crate) async fn exit_review_mode( role: "user".to_string(), content: vec![ContentItem::InputText { text: user_message }], }]) - .await; + .await?; + Ok(()) } fn mcp_init_error_display( @@ -2225,6 +2247,7 @@ mod tests { use crate::config::ConfigToml; use crate::config_types::McpServerConfig; use crate::config_types::McpServerTransportConfig; + use crate::error::Result as CodexResult; use crate::exec::ExecToolCallOutput; use crate::mcp::auth::McpAuthStatusEntry; use crate::tools::format_exec_output_str; @@ -2265,9 +2288,12 @@ mod tests { #[test] fn reconstruct_history_matches_live_compactions() { let (session, turn_context) = make_session_and_context(); - let (rollout_items, expected) = sample_rollout(&session, &turn_context); + let (rollout_items, expected) = + sample_rollout(&session, &turn_context).expect("sample rollout"); - let reconstructed = session.reconstruct_history_from_rollout(&turn_context, &rollout_items); + let reconstructed = session + .reconstruct_history_from_rollout(&turn_context, &rollout_items) + .expect("reconstruct history"); assert_eq!(expected, reconstructed); } @@ -2275,15 +2301,19 @@ mod tests { #[test] fn record_initial_history_reconstructs_resumed_transcript() { let (session, turn_context) = make_session_and_context(); - let (rollout_items, expected) = sample_rollout(&session, &turn_context); + let (rollout_items, expected) = + sample_rollout(&session, &turn_context).expect("sample rollout"); - tokio_test::block_on(session.record_initial_history(InitialHistory::Resumed( - ResumedHistory { - conversation_id: ConversationId::default(), - history: rollout_items, - rollout_path: PathBuf::from("/tmp/resume.jsonl"), - }, - ))); + tokio_test::block_on(async { + session + .record_initial_history(InitialHistory::Resumed(ResumedHistory { + conversation_id: ConversationId::default(), + history: rollout_items, + rollout_path: PathBuf::from("/tmp/resume.jsonl"), + })) + .await + .expect("record resumed history"); + }); let actual = tokio_test::block_on(async { session.state.lock().await.history_snapshot() }); assert_eq!(expected, actual); @@ -2292,9 +2322,15 @@ mod tests { #[test] fn record_initial_history_reconstructs_forked_transcript() { let (session, turn_context) = make_session_and_context(); - let (rollout_items, expected) = sample_rollout(&session, &turn_context); + let (rollout_items, expected) = + sample_rollout(&session, &turn_context).expect("sample rollout"); - tokio_test::block_on(session.record_initial_history(InitialHistory::Forked(rollout_items))); + tokio_test::block_on(async { + session + .record_initial_history(InitialHistory::Forked(rollout_items)) + .await + .expect("record forked history"); + }); let actual = tokio_test::block_on(async { session.state.lock().await.history_snapshot() }); assert_eq!(expected, actual); @@ -2654,10 +2690,10 @@ mod tests { _ctx: Arc, _input: Vec, cancellation_token: CancellationToken, - ) -> Option { + ) -> CodexResult> { if self.listen_to_cancellation_token { cancellation_token.cancelled().await; - return None; + return Ok(None); } loop { sleep(Duration::from_secs(60)).await; @@ -2666,7 +2702,7 @@ mod tests { async fn abort(&self, session: Arc, ctx: Arc) { if let TaskKind::Review = self.kind { - exit_review_mode(session.clone_session(), ctx, None).await; + let _ = exit_review_mode(session.clone_session(), ctx, None).await; } } } @@ -2822,7 +2858,7 @@ mod tests { fn sample_rollout( session: &Session, turn_context: &TurnContext, - ) -> (Vec, Vec) { + ) -> CodexResult<(Vec, Vec)> { let mut rollout_items = Vec::new(); let mut live_history = ConversationHistory::new(); @@ -2830,7 +2866,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())?; let user1 = ResponseItem::Message { id: None, @@ -2839,7 +2875,7 @@ mod tests { text: "first user".to_string(), }], }; - live_history.record_items(std::iter::once(&user1)); + live_history.record_items(std::iter::once(&user1))?; rollout_items.push(RolloutItem::ResponseItem(user1.clone())); let assistant1 = ResponseItem::Message { @@ -2849,7 +2885,7 @@ mod tests { text: "assistant reply one".to_string(), }], }; - live_history.record_items(std::iter::once(&assistant1)); + live_history.record_items(std::iter::once(&assistant1))?; rollout_items.push(RolloutItem::ResponseItem(assistant1.clone())); let summary1 = "summary one"; @@ -2872,7 +2908,7 @@ mod tests { text: "second user".to_string(), }], }; - live_history.record_items(std::iter::once(&user2)); + live_history.record_items(std::iter::once(&user2))?; rollout_items.push(RolloutItem::ResponseItem(user2.clone())); let assistant2 = ResponseItem::Message { @@ -2882,7 +2918,7 @@ mod tests { text: "assistant reply two".to_string(), }], }; - live_history.record_items(std::iter::once(&assistant2)); + live_history.record_items(std::iter::once(&assistant2))?; rollout_items.push(RolloutItem::ResponseItem(assistant2.clone())); let summary2 = "summary two"; @@ -2905,7 +2941,7 @@ mod tests { text: "third user".to_string(), }], }; - live_history.record_items(std::iter::once(&user3)); + live_history.record_items(std::iter::once(&user3))?; rollout_items.push(RolloutItem::ResponseItem(user3.clone())); let assistant3 = ResponseItem::Message { @@ -2915,10 +2951,10 @@ mod tests { text: "assistant reply three".to_string(), }], }; - live_history.record_items(std::iter::once(&assistant3)); + live_history.record_items(std::iter::once(&assistant3))?; rollout_items.push(RolloutItem::ResponseItem(assistant3.clone())); - (rollout_items, live_history.get_history()) + Ok((rollout_items, live_history.get_history())) } #[tokio::test] diff --git a/codex-rs/core/src/codex/compact.rs b/codex-rs/core/src/codex/compact.rs index 26f9f9eb9e..00787df6aa 100644 --- a/codex-rs/core/src/codex/compact.rs +++ b/codex-rs/core/src/codex/compact.rs @@ -39,35 +39,35 @@ struct HistoryBridgeTemplate<'a> { pub(crate) async fn run_inline_auto_compact_task( sess: Arc, turn_context: Arc, -) { +) -> CodexResult<()> { let input = vec![UserInput::Text { text: SUMMARIZATION_PROMPT.to_string(), }]; - run_compact_task_inner(sess, turn_context, input).await; + run_compact_task_inner(sess, turn_context, input).await } pub(crate) async fn run_compact_task( sess: Arc, turn_context: Arc, input: Vec, -) -> Option { +) -> CodexResult> { let start_event = EventMsg::TaskStarted(TaskStartedEvent { model_context_window: turn_context.client.get_model_context_window(), }); sess.send_event(&turn_context, start_event).await; - run_compact_task_inner(sess.clone(), turn_context, input).await; - None + run_compact_task_inner(sess.clone(), turn_context, input).await?; + Ok(None) } async fn run_compact_task_inner( sess: Arc, turn_context: Arc, input: Vec, -) { +) -> CodexResult<()> { let initial_input_for_turn: ResponseInputItem = ResponseInputItem::from(input); let mut history = sess.clone_history().await; - history.record_items(&[initial_input_for_turn.into()]); + history.record_items(&[initial_input_for_turn.into()])?; let mut truncated_count = 0usize; @@ -106,7 +106,7 @@ async fn run_compact_task_inner( break; } Err(CodexErr::Interrupted) => { - return; + return Ok(()); } Err(e @ CodexErr::ContextWindowExceeded) => { if turn_input.len() > 1 { @@ -124,7 +124,7 @@ async fn run_compact_task_inner( message: e.to_string(), }); sess.send_event(&turn_context, event).await; - return; + return Ok(()); } Err(e) => { if retries < max_retries { @@ -142,7 +142,7 @@ async fn run_compact_task_inner( message: e.to_string(), }); sess.send_event(&turn_context, event).await; - return; + return Ok(()); } } } @@ -164,6 +164,7 @@ async fn run_compact_task_inner( message: "Compact task completed".to_string(), }); sess.send_event(&turn_context, event).await; + Ok(()) } pub fn content_items_to_text(content: &[ContentItem]) -> Option { @@ -252,7 +253,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)) + .await?; } Ok(ResponseEvent::RateLimits(snapshot)) => { sess.update_rate_limits(turn_context, snapshot).await; diff --git a/codex-rs/core/src/conversation_history.rs b/codex-rs/core/src/conversation_history.rs index b7f153d29f..225092715c 100644 --- a/codex-rs/core/src/conversation_history.rs +++ b/codex-rs/core/src/conversation_history.rs @@ -2,8 +2,11 @@ use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; +use codex_utils_tokenizer::Tokenizer; use tracing::error; +use crate::error::CodexErr; + /// Transcript of conversation history #[derive(Debug, Clone, Default)] pub(crate) struct ConversationHistory { @@ -34,7 +37,7 @@ impl ConversationHistory { } /// `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) -> Result<(), CodexErr> where I: IntoIterator, I::Item: std::ops::Deref, @@ -43,9 +46,10 @@ impl ConversationHistory { if !is_api_message(&item) { continue; } - + self.validate_input(&item)?; self.items.push(item.clone()); } + Ok(()) } pub(crate) fn get_history(&mut self) -> Vec { @@ -81,6 +85,62 @@ impl ConversationHistory { self.items.clone() } + fn validate_input(&self, item: &ResponseItem) -> Result<(), CodexErr> { + match item { + ResponseItem::Message { content, .. } => { + self.validate_input_content_item(content)?; + Ok(()) + } + ResponseItem::FunctionCall { .. } + | ResponseItem::FunctionCallOutput { .. } + | ResponseItem::CustomToolCall { .. } + | ResponseItem::CustomToolCallOutput { .. } + | ResponseItem::LocalShellCall { .. } + | ResponseItem::Reasoning { .. } + | ResponseItem::WebSearchCall { .. } => Ok(()), + ResponseItem::Other => Err(CodexErr::InvalidInput(format!("invalid input: {item:?}"))), + } + } + + fn validate_input_content_item( + &self, + content: &[codex_protocol::models::ContentItem], + ) -> Result<(), CodexErr> { + let Some(info) = &self.token_info else { + return Ok(()); + }; + let Some(context_window) = info.model_context_window else { + return Ok(()); + }; + + let tokenizer = Tokenizer::try_default() + .map_err(|e| CodexErr::InvalidInput(format!("tokenizer error: {e}")))?; + + let mut input_tokens: i64 = 0; + for item in content { + match item { + codex_protocol::models::ContentItem::InputText { text } => { + input_tokens += tokenizer.count(text); + } + codex_protocol::models::ContentItem::InputImage { .. } => { + // no validation currently + } + codex_protocol::models::ContentItem::OutputText { .. } => { + // no validation currently + } + } + } + + let last_turn_total = info.last_token_usage.total_tokens; + let combined_tokens = input_tokens + last_turn_total; + let threshold = context_window * 95 / 100; + if combined_tokens > threshold { + return Err(CodexErr::InvalidInput("input too large".to_string())); + } + + Ok(()) + } + fn ensure_call_outputs_present(&mut self) { // Collect synthetic outputs to insert immediately after their calls. // Store the insertion position (index of call) alongside the item so @@ -381,7 +441,7 @@ mod tests { fn create_history_with_items(items: Vec) -> ConversationHistory { let mut h = ConversationHistory::new(); - h.record_items(items.iter()); + h.record_items(items.iter()).unwrap(); h } @@ -406,12 +466,12 @@ mod tests { text: "ignored".to_string(), }], }; - h.record_items([&system, &ResponseItem::Other]); + h.record_items([&system, &ResponseItem::Other]).unwrap(); // 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]).unwrap(); let items = h.contents(); assert_eq!( diff --git a/codex-rs/core/src/error.rs b/codex-rs/core/src/error.rs index e733b3c67f..ebee28d02a 100644 --- a/codex-rs/core/src/error.rs +++ b/codex-rs/core/src/error.rs @@ -158,6 +158,9 @@ pub enum CodexErr { #[error("{0}")] EnvVar(EnvVarError), + + #[error("invalid input: {0}")] + InvalidInput(String), } impl From for CodexErr { diff --git a/codex-rs/core/src/response_processing.rs b/codex-rs/core/src/response_processing.rs index b9139ce6c0..5dc38e8883 100644 --- a/codex-rs/core/src/response_processing.rs +++ b/codex-rs/core/src/response_processing.rs @@ -1,5 +1,6 @@ use crate::codex::Session; use crate::conversation_history::ConversationHistory; +use crate::error::Result as CodexResult; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; @@ -13,7 +14,7 @@ pub(crate) async fn process_items( is_review_mode: bool, review_thread_history: &mut ConversationHistory, sess: &Session, -) -> (Vec, Vec) { +) -> CodexResult<(Vec, Vec)> { let mut items_to_record_in_conversation_history = Vec::::new(); let mut responses = Vec::::new(); for processed_response_item in processed_items { @@ -102,11 +103,11 @@ pub(crate) async fn process_items( // 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.record_items(items_to_record_in_conversation_history.iter()); + review_thread_history.record_items(items_to_record_in_conversation_history.iter())?; } else { sess.record_conversation_items(&items_to_record_in_conversation_history) - .await; + .await?; } } - (responses, items_to_record_in_conversation_history) + Ok((responses, items_to_record_in_conversation_history)) } diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 7c4603d9ec..8d47e6fe7f 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -4,6 +4,7 @@ use codex_protocol::models::ResponseItem; use crate::codex::SessionConfiguration; use crate::conversation_history::ConversationHistory; +use crate::error::CodexErr; use crate::protocol::RateLimitSnapshot; use crate::protocol::TokenUsage; use crate::protocol::TokenUsageInfo; @@ -26,12 +27,13 @@ impl SessionState { } // History helpers - pub(crate) fn record_items(&mut self, items: I) + pub(crate) fn record_items(&mut self, items: I) -> Result<(), CodexErr> where I: IntoIterator, I::Item: std::ops::Deref, { - self.history.record_items(items) + self.history.record_items(items)?; + Ok(()) } pub(crate) fn history_snapshot(&mut self) -> Vec { diff --git a/codex-rs/core/src/tasks/compact.rs b/codex-rs/core/src/tasks/compact.rs index 64b2a9d2b4..f690b5d3ae 100644 --- a/codex-rs/core/src/tasks/compact.rs +++ b/codex-rs/core/src/tasks/compact.rs @@ -5,6 +5,7 @@ use tokio_util::sync::CancellationToken; use crate::codex::TurnContext; use crate::codex::compact; +use crate::error::Result as CodexResult; use crate::state::TaskKind; use codex_protocol::user_input::UserInput; @@ -26,7 +27,7 @@ impl SessionTask for CompactTask { ctx: Arc, input: Vec, _cancellation_token: CancellationToken, - ) -> Option { + ) -> CodexResult> { compact::run_compact_task(session.clone_session(), ctx, input).await } } diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 79527814e3..7248698cbd 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -15,6 +15,8 @@ use tracing::warn; use crate::codex::Session; use crate::codex::TurnContext; +use crate::error::Result as CodexResult; +use crate::protocol::ErrorEvent; use crate::protocol::EventMsg; use crate::protocol::TaskCompleteEvent; use crate::protocol::TurnAbortReason; @@ -56,7 +58,7 @@ pub(crate) trait SessionTask: Send + Sync + 'static { ctx: Arc, input: Vec, cancellation_token: CancellationToken, - ) -> Option; + ) -> CodexResult>; async fn abort(&self, session: Arc, ctx: Arc) { let _ = (session, ctx); @@ -86,7 +88,7 @@ impl Session { let task_cancellation_token = cancellation_token.child_token(); tokio::spawn(async move { let ctx_for_finish = Arc::clone(&ctx); - let last_agent_message = task_for_run + let run_result = task_for_run .run( Arc::clone(&session_ctx), ctx, @@ -98,8 +100,21 @@ impl Session { if !task_cancellation_token.is_cancelled() { // Emit completion uniformly from spawn site so all tasks share the same lifecycle. let sess = session_ctx.clone_session(); - sess.on_task_finished(ctx_for_finish, last_agent_message) - .await; + match run_result { + Ok(last_agent_message) => { + sess.on_task_finished(ctx_for_finish, last_agent_message) + .await; + } + Err(err) => { + let message = err.to_string(); + sess.send_event( + ctx_for_finish.as_ref(), + EventMsg::Error(ErrorEvent { message }), + ) + .await; + sess.on_task_finished(ctx_for_finish, None).await; + } + } } done_clone.notify_waiters(); }) diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index 58ecedd446..14a6b21cc0 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -5,6 +5,7 @@ use tokio_util::sync::CancellationToken; use crate::codex::TurnContext; use crate::codex::run_task; +use crate::error::Result as CodexResult; use crate::state::TaskKind; use codex_protocol::user_input::UserInput; @@ -26,7 +27,7 @@ impl SessionTask for RegularTask { ctx: Arc, input: Vec, cancellation_token: CancellationToken, - ) -> Option { + ) -> CodexResult> { let sess = session.clone_session(); run_task(sess, ctx, input, TaskKind::Regular, cancellation_token).await } diff --git a/codex-rs/core/src/tasks/review.rs b/codex-rs/core/src/tasks/review.rs index fbf553b349..9d45cdef3d 100644 --- a/codex-rs/core/src/tasks/review.rs +++ b/codex-rs/core/src/tasks/review.rs @@ -6,6 +6,7 @@ use tokio_util::sync::CancellationToken; use crate::codex::TurnContext; use crate::codex::exit_review_mode; use crate::codex::run_task; +use crate::error::Result as CodexResult; use crate::state::TaskKind; use codex_protocol::user_input::UserInput; @@ -27,12 +28,12 @@ impl SessionTask for ReviewTask { ctx: Arc, input: Vec, cancellation_token: CancellationToken, - ) -> Option { + ) -> CodexResult> { let sess = session.clone_session(); run_task(sess, ctx, input, TaskKind::Review, cancellation_token).await } async fn abort(&self, session: Arc, ctx: Arc) { - exit_review_mode(session.clone_session(), ctx, None).await; + let _ = exit_review_mode(session.clone_session(), ctx, None).await; } } diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index bfaef15a92..173eaf1c1e 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -13,6 +13,7 @@ mod compact_resume_fork; mod exec; mod fork_conversation; mod grep_files; +mod input_validation; mod items; mod json_result; mod list_dir;