From 86adf53235f3885dace914a53b00d67daa4e76eb Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Mon, 26 Jan 2026 19:33:48 -0800 Subject: [PATCH] fix: handle all web_search actions and in progress invocations (#9960) ### Summary - Parse all `web_search` tool actions (`search`, `find_in_page`, `open_page`). - Previously we only parsed + displayed `search`, which made the TUI appear to pause when the other actions were being used. - Show in progress `web_search` calls as `Searching the web` - Previously we only showed completed tool calls image ### Tests Added + updated tests, tested locally ### Follow ups Update VSCode extension to display these as well --- .../app-server-protocol/src/protocol/v2.rs | 4 + codex-rs/codex-api/src/sse/responses.rs | 2 +- codex-rs/core/src/codex.rs | 4 +- codex-rs/core/src/event_mapping.rs | 117 +++++++++++++++--- codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/web_search.rs | 24 ++++ codex-rs/core/tests/common/responses.rs | 5 +- codex-rs/core/tests/suite/client.rs | 4 +- codex-rs/core/tests/suite/items.rs | 23 ++-- .../src/event_processor_with_human_output.rs | 20 ++- .../src/event_processor_with_jsonl_output.rs | 34 ++++- codex-rs/exec/src/exec_events.rs | 3 + .../tests/event_processor_with_json_output.rs | 78 +++++++++++- codex-rs/protocol/src/items.rs | 5 +- codex-rs/protocol/src/models.rs | 45 +++++-- codex-rs/protocol/src/protocol.rs | 6 + codex-rs/tui/src/chatwidget.rs | 35 +++++- codex-rs/tui/src/history_cell.rs | 114 +++++++++++++++-- 18 files changed, 462 insertions(+), 62 deletions(-) create mode 100644 codex-rs/core/src/web_search.rs diff --git a/codex-rs/app-server-protocol/src/protocol/v2.rs b/codex-rs/app-server-protocol/src/protocol/v2.rs index 699ad5f52a..e4b3b90589 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2.rs @@ -2621,6 +2621,7 @@ mod tests { use codex_protocol::items::TurnItem; use codex_protocol::items::UserMessageItem; use codex_protocol::items::WebSearchItem; + use codex_protocol::models::WebSearchAction; use codex_protocol::protocol::NetworkAccess as CoreNetworkAccess; use codex_protocol::user_input::UserInput as CoreUserInput; use pretty_assertions::assert_eq; @@ -2728,6 +2729,9 @@ mod tests { let search_item = TurnItem::WebSearch(WebSearchItem { id: "search-1".to_string(), query: "docs".to_string(), + action: WebSearchAction::Search { + query: Some("docs".to_string()), + }, }); assert_eq!( diff --git a/codex-rs/codex-api/src/sse/responses.rs b/codex-rs/codex-api/src/sse/responses.rs index f23975f8dd..ad94cf5ea1 100644 --- a/codex-rs/codex-api/src/sse/responses.rs +++ b/codex-rs/codex-api/src/sse/responses.rs @@ -291,7 +291,7 @@ pub fn process_responses_event( if let Ok(item) = serde_json::from_value::(item_val) { return Ok(Some(ResponseEvent::OutputItemAdded(item))); } - debug!("failed to parse ResponseItem from output_item.done"); + debug!("failed to parse ResponseItem from output_item.added"); } } "response.reasoning_summary_part.added" => { diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 0f696ea84d..5e535f5130 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -3439,10 +3439,8 @@ async fn try_run_sampling_request( } ResponseEvent::OutputItemAdded(item) => { if let Some(turn_item) = handle_non_tool_response_item(&item).await { - let tracked_item = turn_item.clone(); sess.emit_turn_item_started(&turn_context, &turn_item).await; - - active_item = Some(tracked_item); + active_item = Some(turn_item); } } ResponseEvent::ServerReasoningIncluded(included) => { diff --git a/codex-rs/core/src/event_mapping.rs b/codex-rs/core/src/event_mapping.rs index 55a1037542..32382aad71 100644 --- a/codex-rs/core/src/event_mapping.rs +++ b/codex-rs/core/src/event_mapping.rs @@ -21,6 +21,7 @@ use crate::instructions::SkillInstructions; use crate::instructions::UserInstructions; use crate::session_prefix::is_session_prefix; use crate::user_shell_command::is_user_shell_command_text; +use crate::web_search::web_search_action_detail; fn parse_user_message(message: &[ContentItem]) -> Option { if UserInstructions::is_user_instructions(message) @@ -127,14 +128,17 @@ pub fn parse_turn_item(item: &ResponseItem) -> Option { raw_content, })) } - ResponseItem::WebSearchCall { - id, - action: WebSearchAction::Search { query }, - .. - } => Some(TurnItem::WebSearch(WebSearchItem { - id: id.clone().unwrap_or_default(), - query: query.clone().unwrap_or_default(), - })), + ResponseItem::WebSearchCall { id, action, .. } => { + let (action, query) = match action { + Some(action) => (action.clone(), web_search_action_detail(action)), + None => (WebSearchAction::Other, String::new()), + }; + Some(TurnItem::WebSearch(WebSearchItem { + id: id.clone().unwrap_or_default(), + query, + action, + })) + } _ => None, } } @@ -144,6 +148,7 @@ mod tests { use super::parse_turn_item; use codex_protocol::items::AgentMessageContent; use codex_protocol::items::TurnItem; + use codex_protocol::items::WebSearchItem; use codex_protocol::models::ContentItem; use codex_protocol::models::ReasoningItemContent; use codex_protocol::models::ReasoningItemReasoningSummary; @@ -419,18 +424,102 @@ mod tests { let item = ResponseItem::WebSearchCall { id: Some("ws_1".to_string()), status: Some("completed".to_string()), - action: WebSearchAction::Search { + action: Some(WebSearchAction::Search { query: Some("weather".to_string()), - }, + }), }; let turn_item = parse_turn_item(&item).expect("expected web search turn item"); match turn_item { - TurnItem::WebSearch(search) => { - assert_eq!(search.id, "ws_1"); - assert_eq!(search.query, "weather"); - } + TurnItem::WebSearch(search) => assert_eq!( + search, + WebSearchItem { + id: "ws_1".to_string(), + query: "weather".to_string(), + action: WebSearchAction::Search { + query: Some("weather".to_string()), + }, + } + ), + other => panic!("expected TurnItem::WebSearch, got {other:?}"), + } + } + + #[test] + fn parses_web_search_open_page_call() { + let item = ResponseItem::WebSearchCall { + id: Some("ws_open".to_string()), + status: Some("completed".to_string()), + action: Some(WebSearchAction::OpenPage { + url: Some("https://example.com".to_string()), + }), + }; + + let turn_item = parse_turn_item(&item).expect("expected web search turn item"); + + match turn_item { + TurnItem::WebSearch(search) => assert_eq!( + search, + WebSearchItem { + id: "ws_open".to_string(), + query: "https://example.com".to_string(), + action: WebSearchAction::OpenPage { + url: Some("https://example.com".to_string()), + }, + } + ), + other => panic!("expected TurnItem::WebSearch, got {other:?}"), + } + } + + #[test] + fn parses_web_search_find_in_page_call() { + let item = ResponseItem::WebSearchCall { + id: Some("ws_find".to_string()), + status: Some("completed".to_string()), + action: Some(WebSearchAction::FindInPage { + url: Some("https://example.com".to_string()), + pattern: Some("needle".to_string()), + }), + }; + + let turn_item = parse_turn_item(&item).expect("expected web search turn item"); + + match turn_item { + TurnItem::WebSearch(search) => assert_eq!( + search, + WebSearchItem { + id: "ws_find".to_string(), + query: "'needle' in https://example.com".to_string(), + action: WebSearchAction::FindInPage { + url: Some("https://example.com".to_string()), + pattern: Some("needle".to_string()), + }, + } + ), + other => panic!("expected TurnItem::WebSearch, got {other:?}"), + } + } + + #[test] + fn parses_partial_web_search_call_without_action_as_other() { + let item = ResponseItem::WebSearchCall { + id: Some("ws_partial".to_string()), + status: Some("in_progress".to_string()), + action: None, + }; + + let turn_item = parse_turn_item(&item).expect("expected web search turn item"); + match turn_item { + TurnItem::WebSearch(search) => assert_eq!( + search, + WebSearchItem { + id: "ws_partial".to_string(), + query: String::new(), + action: WebSearchAction::Other, + } + ), other => panic!("expected TurnItem::WebSearch, got {other:?}"), } } diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index f662f41b21..010ba3874a 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -69,6 +69,7 @@ mod event_mapping; pub mod review_format; pub mod review_prompts; mod thread_manager; +pub mod web_search; pub use codex_protocol::protocol::InitialHistory; pub use thread_manager::NewThread; pub use thread_manager::ThreadManager; diff --git a/codex-rs/core/src/web_search.rs b/codex-rs/core/src/web_search.rs new file mode 100644 index 0000000000..458758b70e --- /dev/null +++ b/codex-rs/core/src/web_search.rs @@ -0,0 +1,24 @@ +use codex_protocol::models::WebSearchAction; + +pub fn web_search_action_detail(action: &WebSearchAction) -> String { + match action { + WebSearchAction::Search { query } => query.clone().unwrap_or_default(), + WebSearchAction::OpenPage { url } => url.clone().unwrap_or_default(), + WebSearchAction::FindInPage { url, pattern } => match (pattern, url) { + (Some(pattern), Some(url)) => format!("'{pattern}' in {url}"), + (Some(pattern), None) => format!("'{pattern}'"), + (None, Some(url)) => url.clone(), + (None, None) => String::new(), + }, + WebSearchAction::Other => String::new(), + } +} + +pub fn web_search_detail(action: Option<&WebSearchAction>, query: &str) -> String { + let detail = action.map(web_search_action_detail).unwrap_or_default(); + if detail.is_empty() { + query.to_string() + } else { + detail + } +} diff --git a/codex-rs/core/tests/common/responses.rs b/codex-rs/core/tests/common/responses.rs index 3ca45ca735..e36d9592f1 100644 --- a/codex-rs/core/tests/common/responses.rs +++ b/codex-rs/core/tests/common/responses.rs @@ -494,14 +494,13 @@ pub fn ev_reasoning_text_delta(delta: &str) -> Value { }) } -pub fn ev_web_search_call_added(id: &str, status: &str, query: &str) -> Value { +pub fn ev_web_search_call_added_partial(id: &str, status: &str) -> Value { serde_json::json!({ "type": "response.output_item.added", "item": { "type": "web_search_call", "id": id, - "status": status, - "action": {"type": "search", "query": query} + "status": status } }) } diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 1ed398fe9a..bdc2f311c0 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1211,9 +1211,9 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() { prompt.input.push(ResponseItem::WebSearchCall { id: Some("web-search-id".into()), status: Some("completed".into()), - action: WebSearchAction::Search { + action: Some(WebSearchAction::Search { query: Some("weather".into()), - }, + }), }); prompt.input.push(ResponseItem::FunctionCall { id: Some("function-id".into()), diff --git a/codex-rs/core/tests/suite/items.rs b/codex-rs/core/tests/suite/items.rs index 60d0dbc75f..08ac598388 100644 --- a/codex-rs/core/tests/suite/items.rs +++ b/codex-rs/core/tests/suite/items.rs @@ -6,6 +6,7 @@ use codex_core::protocol::ItemCompletedEvent; use codex_core::protocol::ItemStartedEvent; use codex_core::protocol::Op; use codex_protocol::items::TurnItem; +use codex_protocol::models::WebSearchAction; use codex_protocol::user_input::ByteRange; use codex_protocol::user_input::TextElement; use codex_protocol::user_input::UserInput; @@ -18,7 +19,7 @@ use core_test_support::responses::ev_reasoning_item_added; use core_test_support::responses::ev_reasoning_summary_text_delta; use core_test_support::responses::ev_reasoning_text_delta; use core_test_support::responses::ev_response_created; -use core_test_support::responses::ev_web_search_call_added; +use core_test_support::responses::ev_web_search_call_added_partial; use core_test_support::responses::ev_web_search_call_done; use core_test_support::responses::mount_sse_once; use core_test_support::responses::sse; @@ -208,8 +209,7 @@ async fn web_search_item_is_emitted() -> anyhow::Result<()> { let TestCodex { codex, .. } = test_codex().build(&server).await?; - let web_search_added = - ev_web_search_call_added("web-search-1", "in_progress", "weather seattle"); + let web_search_added = ev_web_search_call_added_partial("web-search-1", "in_progress"); let web_search_done = ev_web_search_call_done("web-search-1", "completed", "weather seattle"); let first_response = sse(vec![ @@ -230,11 +230,8 @@ async fn web_search_item_is_emitted() -> anyhow::Result<()> { }) .await?; - let started = wait_for_event_match(&codex, |ev| match ev { - EventMsg::ItemStarted(ItemStartedEvent { - item: TurnItem::WebSearch(item), - .. - }) => Some(item.clone()), + let begin = wait_for_event_match(&codex, |ev| match ev { + EventMsg::WebSearchBegin(event) => Some(event.clone()), _ => None, }) .await; @@ -247,8 +244,14 @@ async fn web_search_item_is_emitted() -> anyhow::Result<()> { }) .await; - assert_eq!(started.id, completed.id); - assert_eq!(completed.query, "weather seattle"); + assert_eq!(begin.call_id, "web-search-1"); + assert_eq!(completed.id, begin.call_id); + assert_eq!( + completed.action, + WebSearchAction::Search { + query: Some("weather seattle".to_string()), + } + ); Ok(()) } diff --git a/codex-rs/exec/src/event_processor_with_human_output.rs b/codex-rs/exec/src/event_processor_with_human_output.rs index 6275ceaf8d..fb7aa54ab4 100644 --- a/codex-rs/exec/src/event_processor_with_human_output.rs +++ b/codex-rs/exec/src/event_processor_with_human_output.rs @@ -32,6 +32,7 @@ use codex_core::protocol::TurnCompleteEvent; use codex_core::protocol::TurnDiffEvent; use codex_core::protocol::WarningEvent; use codex_core::protocol::WebSearchEndEvent; +use codex_core::web_search::web_search_detail; use codex_protocol::num_format::format_with_separators; use owo_colors::OwoColorize; use owo_colors::Style; @@ -370,8 +371,20 @@ impl EventProcessor for EventProcessorWithHumanOutput { } } } - EventMsg::WebSearchEnd(WebSearchEndEvent { call_id: _, query }) => { - ts_msg!(self, "🌐 Searched: {query}"); + EventMsg::WebSearchBegin(_) => { + ts_msg!(self, "🌐 Searching the web..."); + } + EventMsg::WebSearchEnd(WebSearchEndEvent { + call_id: _, + query, + action, + }) => { + let detail = web_search_detail(Some(&action), &query); + if detail.is_empty() { + ts_msg!(self, "🌐 Searched the web"); + } else { + ts_msg!(self, "🌐 Searched: {detail}"); + } } EventMsg::PatchApplyBegin(PatchApplyBeginEvent { call_id, @@ -737,8 +750,7 @@ impl EventProcessor for EventProcessorWithHumanOutput { ); } EventMsg::ShutdownComplete => return CodexStatus::Shutdown, - EventMsg::WebSearchBegin(_) - | EventMsg::ExecApprovalRequest(_) + EventMsg::ExecApprovalRequest(_) | EventMsg::ApplyPatchApprovalRequest(_) | EventMsg::TerminalInteraction(_) | EventMsg::ExecCommandOutputDelta(_) diff --git a/codex-rs/exec/src/event_processor_with_jsonl_output.rs b/codex-rs/exec/src/event_processor_with_jsonl_output.rs index 779ec9292d..0dfce33cb6 100644 --- a/codex-rs/exec/src/event_processor_with_jsonl_output.rs +++ b/codex-rs/exec/src/event_processor_with_jsonl_output.rs @@ -49,6 +49,7 @@ use codex_core::protocol::CollabCloseBeginEvent; use codex_core::protocol::CollabCloseEndEvent; use codex_core::protocol::CollabWaitingBeginEvent; use codex_core::protocol::CollabWaitingEndEvent; +use codex_protocol::models::WebSearchAction; use codex_protocol::plan_tool::StepStatus; use codex_protocol::plan_tool::UpdatePlanArgs; use serde_json::Value as JsonValue; @@ -66,6 +67,7 @@ pub struct EventProcessorWithJsonOutput { last_total_token_usage: Option, running_mcp_tool_calls: HashMap, running_collab_tool_calls: HashMap, + running_web_search_calls: HashMap, last_critical_error: Option, } @@ -107,6 +109,7 @@ impl EventProcessorWithJsonOutput { last_total_token_usage: None, running_mcp_tool_calls: HashMap::new(), running_collab_tool_calls: HashMap::new(), + running_web_search_calls: HashMap::new(), last_critical_error: None, } } @@ -138,7 +141,7 @@ impl EventProcessorWithJsonOutput { protocol::EventMsg::CollabCloseEnd(ev) => self.handle_collab_close_end(ev), protocol::EventMsg::PatchApplyBegin(ev) => self.handle_patch_apply_begin(ev), protocol::EventMsg::PatchApplyEnd(ev) => self.handle_patch_apply_end(ev), - protocol::EventMsg::WebSearchBegin(_) => Vec::new(), + protocol::EventMsg::WebSearchBegin(ev) => self.handle_web_search_begin(ev), protocol::EventMsg::WebSearchEnd(ev) => self.handle_web_search_end(ev), protocol::EventMsg::TokenCount(ev) => { if let Some(info) = &ev.info { @@ -195,11 +198,36 @@ impl EventProcessorWithJsonOutput { })] } - fn handle_web_search_end(&self, ev: &protocol::WebSearchEndEvent) -> Vec { + fn handle_web_search_begin(&mut self, ev: &protocol::WebSearchBeginEvent) -> Vec { + if self.running_web_search_calls.contains_key(&ev.call_id) { + return Vec::new(); + } + let item_id = self.get_next_item_id(); + self.running_web_search_calls + .insert(ev.call_id.clone(), item_id.clone()); let item = ThreadItem { - id: self.get_next_item_id(), + id: item_id, details: ThreadItemDetails::WebSearch(WebSearchItem { + id: ev.call_id.clone(), + query: String::new(), + action: WebSearchAction::Other, + }), + }; + + vec![ThreadEvent::ItemStarted(ItemStartedEvent { item })] + } + + fn handle_web_search_end(&mut self, ev: &protocol::WebSearchEndEvent) -> Vec { + let item_id = self + .running_web_search_calls + .remove(&ev.call_id) + .unwrap_or_else(|| self.get_next_item_id()); + let item = ThreadItem { + id: item_id, + details: ThreadItemDetails::WebSearch(WebSearchItem { + id: ev.call_id.clone(), query: ev.query.clone(), + action: ev.action.clone(), }), }; diff --git a/codex-rs/exec/src/exec_events.rs b/codex-rs/exec/src/exec_events.rs index e113a5c4b4..47c67a6dca 100644 --- a/codex-rs/exec/src/exec_events.rs +++ b/codex-rs/exec/src/exec_events.rs @@ -1,3 +1,4 @@ +use codex_protocol::models::WebSearchAction; use mcp_types::ContentBlock as McpContentBlock; use serde::Deserialize; use serde::Serialize; @@ -280,7 +281,9 @@ pub struct McpToolCallItem { /// A web search request. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, TS)] pub struct WebSearchItem { + pub id: String, pub query: String, + pub action: WebSearchAction, } /// An error notification. diff --git a/codex-rs/exec/tests/event_processor_with_json_output.rs b/codex-rs/exec/tests/event_processor_with_json_output.rs index cb2a283b30..105824c60d 100644 --- a/codex-rs/exec/tests/event_processor_with_json_output.rs +++ b/codex-rs/exec/tests/event_processor_with_json_output.rs @@ -20,6 +20,7 @@ use codex_core::protocol::PatchApplyEndEvent; use codex_core::protocol::SandboxPolicy; use codex_core::protocol::SessionConfiguredEvent; use codex_core::protocol::WarningEvent; +use codex_core::protocol::WebSearchBeginEvent; use codex_core::protocol::WebSearchEndEvent; use codex_exec::event_processor_with_jsonl_output::EventProcessorWithJsonOutput; use codex_exec::exec_events::AgentMessageItem; @@ -54,6 +55,7 @@ use codex_exec::exec_events::TurnStartedEvent; use codex_exec::exec_events::Usage; use codex_exec::exec_events::WebSearchItem; use codex_protocol::ThreadId; +use codex_protocol::models::WebSearchAction; use codex_protocol::plan_tool::PlanItemArg; use codex_protocol::plan_tool::StepStatus; use codex_protocol::plan_tool::UpdatePlanArgs; @@ -124,11 +126,15 @@ fn task_started_produces_turn_started_event() { fn web_search_end_emits_item_completed() { let mut ep = EventProcessorWithJsonOutput::new(None); let query = "rust async await".to_string(); + let action = WebSearchAction::Search { + query: Some(query.clone()), + }; let out = ep.collect_thread_events(&event( "w1", EventMsg::WebSearchEnd(WebSearchEndEvent { call_id: "call-123".to_string(), query: query.clone(), + action: action.clone(), }), )); @@ -137,12 +143,82 @@ fn web_search_end_emits_item_completed() { vec![ThreadEvent::ItemCompleted(ItemCompletedEvent { item: ThreadItem { id: "item_0".to_string(), - details: ThreadItemDetails::WebSearch(WebSearchItem { query }), + details: ThreadItemDetails::WebSearch(WebSearchItem { + id: "call-123".to_string(), + query, + action, + }), }, })] ); } +#[test] +fn web_search_begin_emits_item_started() { + let mut ep = EventProcessorWithJsonOutput::new(None); + let out = ep.collect_thread_events(&event( + "w0", + EventMsg::WebSearchBegin(WebSearchBeginEvent { + call_id: "call-0".to_string(), + }), + )); + + assert_eq!(out.len(), 1); + let ThreadEvent::ItemStarted(ItemStartedEvent { item }) = &out[0] else { + panic!("expected ItemStarted"); + }; + assert!(item.id.starts_with("item_")); + assert_eq!( + item.details, + ThreadItemDetails::WebSearch(WebSearchItem { + id: "call-0".to_string(), + query: String::new(), + action: WebSearchAction::Other, + }) + ); +} + +#[test] +fn web_search_begin_then_end_reuses_item_id() { + let mut ep = EventProcessorWithJsonOutput::new(None); + let begin = ep.collect_thread_events(&event( + "w0", + EventMsg::WebSearchBegin(WebSearchBeginEvent { + call_id: "call-1".to_string(), + }), + )); + let ThreadEvent::ItemStarted(ItemStartedEvent { item: started_item }) = &begin[0] else { + panic!("expected ItemStarted"); + }; + let action = WebSearchAction::Search { + query: Some("rust async await".to_string()), + }; + let end = ep.collect_thread_events(&event( + "w1", + EventMsg::WebSearchEnd(WebSearchEndEvent { + call_id: "call-1".to_string(), + query: "rust async await".to_string(), + action: action.clone(), + }), + )); + let ThreadEvent::ItemCompleted(ItemCompletedEvent { + item: completed_item, + }) = &end[0] + else { + panic!("expected ItemCompleted"); + }; + + assert_eq!(completed_item.id, started_item.id); + assert_eq!( + completed_item.details, + ThreadItemDetails::WebSearch(WebSearchItem { + id: "call-1".to_string(), + query: "rust async await".to_string(), + action, + }) + ); +} + #[test] fn plan_update_emits_todo_list_started_updated_and_completed() { let mut ep = EventProcessorWithJsonOutput::new(None); diff --git a/codex-rs/protocol/src/items.rs b/codex-rs/protocol/src/items.rs index 9276a759cc..12fa6a0f51 100644 --- a/codex-rs/protocol/src/items.rs +++ b/codex-rs/protocol/src/items.rs @@ -1,3 +1,4 @@ +use crate::models::WebSearchAction; use crate::protocol::AgentMessageEvent; use crate::protocol::AgentReasoningEvent; use crate::protocol::AgentReasoningRawContentEvent; @@ -49,10 +50,11 @@ pub struct ReasoningItem { pub raw_content: Vec, } -#[derive(Debug, Clone, Deserialize, Serialize, TS, JsonSchema)] +#[derive(Debug, Clone, Deserialize, Serialize, TS, JsonSchema, PartialEq)] pub struct WebSearchItem { pub id: String, pub query: String, + pub action: WebSearchAction, } impl UserMessageItem { @@ -181,6 +183,7 @@ impl WebSearchItem { EventMsg::WebSearchEnd(WebSearchEndEvent { call_id: self.id.clone(), query: self.query.clone(), + action: self.action.clone(), }) } } diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index 57241bb8b6..b0b8156916 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -157,7 +157,9 @@ pub enum ResponseItem { #[serde(default, skip_serializing_if = "Option::is_none")] #[ts(optional)] status: Option, - action: WebSearchAction, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] + action: Option, }, // Generated by the harness but considered exactly as a model response. GhostSnapshot { @@ -1034,10 +1036,12 @@ mod tests { "query": "weather seattle" } }"#, - WebSearchAction::Search { + None, + Some(WebSearchAction::Search { query: Some("weather seattle".into()), - }, + }), Some("completed".into()), + true, ), ( r#"{ @@ -1048,10 +1052,12 @@ mod tests { "url": "https://example.com" } }"#, - WebSearchAction::OpenPage { + None, + Some(WebSearchAction::OpenPage { url: Some("https://example.com".into()), - }, + }), Some("open".into()), + true, ), ( r#"{ @@ -1063,26 +1069,43 @@ mod tests { "pattern": "installation" } }"#, - WebSearchAction::FindInPage { + None, + Some(WebSearchAction::FindInPage { url: Some("https://example.com/docs".into()), pattern: Some("installation".into()), - }, + }), Some("in_progress".into()), + true, + ), + ( + r#"{ + "type": "web_search_call", + "status": "in_progress", + "id": "ws_partial" + }"#, + Some("ws_partial".into()), + None, + Some("in_progress".into()), + false, ), ]; - for (json_literal, expected_action, expected_status) in cases { + for (json_literal, expected_id, expected_action, expected_status, expect_roundtrip) in cases + { let parsed: ResponseItem = serde_json::from_str(json_literal)?; let expected = ResponseItem::WebSearchCall { - id: None, + id: expected_id.clone(), status: expected_status.clone(), action: expected_action.clone(), }; assert_eq!(parsed, expected); let serialized = serde_json::to_value(&parsed)?; - let original_value: serde_json::Value = serde_json::from_str(json_literal)?; - assert_eq!(serialized, original_value); + let mut expected_serialized: serde_json::Value = serde_json::from_str(json_literal)?; + if !expect_roundtrip && let Some(obj) = expected_serialized.as_object_mut() { + obj.remove("id"); + } + assert_eq!(serialized, expected_serialized); } Ok(()) diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index e9cb8147d7..c60b13d72f 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -24,6 +24,7 @@ use crate::message_history::HistoryEntry; use crate::models::BaseInstructions; use crate::models::ContentItem; use crate::models::ResponseItem; +use crate::models::WebSearchAction; use crate::num_format::format_with_separators; use crate::openai_models::ReasoningEffort as ReasoningEffortConfig; use crate::parse_command::ParsedCommand; @@ -1041,6 +1042,7 @@ impl HasLegacyEvent for ReasoningRawContentDeltaEvent { impl HasLegacyEvent for EventMsg { fn as_legacy_events(&self, show_raw_agent_reasoning: bool) -> Vec { match self { + EventMsg::ItemStarted(event) => event.as_legacy_events(show_raw_agent_reasoning), EventMsg::ItemCompleted(event) => event.as_legacy_events(show_raw_agent_reasoning), EventMsg::AgentMessageContentDelta(event) => { event.as_legacy_events(show_raw_agent_reasoning) @@ -1402,6 +1404,7 @@ pub struct WebSearchBeginEvent { pub struct WebSearchEndEvent { pub call_id: String, pub query: String, + pub action: WebSearchAction, } // Conversation kept for backward compatibility. @@ -2375,6 +2378,9 @@ mod tests { item: TurnItem::WebSearch(WebSearchItem { id: "search-1".into(), query: "find docs".into(), + action: WebSearchAction::Search { + query: Some("find docs".into()), + }, }), }; diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 6f45ec25ce..9dc1bc4bfe 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -162,6 +162,7 @@ use crate::history_cell::AgentMessageCell; use crate::history_cell::HistoryCell; use crate::history_cell::McpToolCallCell; use crate::history_cell::PlainHistoryCell; +use crate::history_cell::WebSearchCell; use crate::key_hint; use crate::key_hint::KeyBinding; use crate::markdown::append_markdown; @@ -1490,13 +1491,43 @@ impl ChatWidget { self.defer_or_handle(|q| q.push_mcp_end(ev), |s| s.handle_mcp_end_now(ev2)); } - fn on_web_search_begin(&mut self, _ev: WebSearchBeginEvent) { + fn on_web_search_begin(&mut self, ev: WebSearchBeginEvent) { self.flush_answer_stream_with_separator(); + self.flush_active_cell(); + self.active_cell = Some(Box::new(history_cell::new_active_web_search_call( + ev.call_id, + String::new(), + self.config.animations, + ))); + self.bump_active_cell_revision(); + self.request_redraw(); } fn on_web_search_end(&mut self, ev: WebSearchEndEvent) { self.flush_answer_stream_with_separator(); - self.add_to_history(history_cell::new_web_search_call(ev.query)); + let WebSearchEndEvent { + call_id, + query, + action, + } = ev; + let mut handled = false; + if let Some(cell) = self + .active_cell + .as_mut() + .and_then(|cell| cell.as_any_mut().downcast_mut::()) + && cell.call_id() == call_id + { + cell.update(action.clone(), query.clone()); + cell.complete(); + self.bump_active_cell_revision(); + self.flush_active_cell(); + handled = true; + } + + if !handled { + self.add_to_history(history_cell::new_web_search_call(call_id, query, action)); + } + self.had_work_activity = true; } fn on_collab_event(&mut self, cell: PlainHistoryCell) { diff --git a/codex-rs/tui/src/history_cell.rs b/codex-rs/tui/src/history_cell.rs index 94da257d9d..02a148df4c 100644 --- a/codex-rs/tui/src/history_cell.rs +++ b/codex-rs/tui/src/history_cell.rs @@ -43,6 +43,8 @@ use codex_core::protocol::FileChange; use codex_core::protocol::McpAuthStatus; use codex_core::protocol::McpInvocation; use codex_core::protocol::SessionConfiguredEvent; +use codex_core::web_search::web_search_detail; +use codex_protocol::models::WebSearchAction; use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; use codex_protocol::plan_tool::PlanItemArg; use codex_protocol::plan_tool::StepStatus; @@ -1342,9 +1344,89 @@ pub(crate) fn new_active_mcp_tool_call( McpToolCallCell::new(call_id, invocation, animations_enabled) } -pub(crate) fn new_web_search_call(query: String) -> PrefixedWrappedHistoryCell { - let text: Text<'static> = Line::from(vec!["Searched".bold(), " ".into(), query.into()]).into(); - PrefixedWrappedHistoryCell::new(text, "• ".dim(), " ") +fn web_search_header(completed: bool) -> &'static str { + if completed { + "Searched" + } else { + "Searching the web" + } +} + +#[derive(Debug)] +pub(crate) struct WebSearchCell { + call_id: String, + query: String, + action: Option, + start_time: Instant, + completed: bool, + animations_enabled: bool, +} + +impl WebSearchCell { + pub(crate) fn new( + call_id: String, + query: String, + action: Option, + animations_enabled: bool, + ) -> Self { + Self { + call_id, + query, + action, + start_time: Instant::now(), + completed: false, + animations_enabled, + } + } + + pub(crate) fn call_id(&self) -> &str { + &self.call_id + } + + pub(crate) fn update(&mut self, action: WebSearchAction, query: String) { + self.action = Some(action); + self.query = query; + } + + pub(crate) fn complete(&mut self) { + self.completed = true; + } +} + +impl HistoryCell for WebSearchCell { + fn display_lines(&self, width: u16) -> Vec> { + let bullet = if self.completed { + "•".dim() + } else { + spinner(Some(self.start_time), self.animations_enabled) + }; + let header = web_search_header(self.completed); + let detail = web_search_detail(self.action.as_ref(), &self.query); + let text: Text<'static> = if detail.is_empty() { + Line::from(vec![header.bold()]).into() + } else { + Line::from(vec![header.bold(), " ".into(), detail.into()]).into() + }; + PrefixedWrappedHistoryCell::new(text, vec![bullet, " ".into()], " ").display_lines(width) + } +} + +pub(crate) fn new_active_web_search_call( + call_id: String, + query: String, + animations_enabled: bool, +) -> WebSearchCell { + WebSearchCell::new(call_id, query, None, animations_enabled) +} + +pub(crate) fn new_web_search_call( + call_id: String, + query: String, + action: WebSearchAction, +) -> WebSearchCell { + let mut cell = WebSearchCell::new(call_id, query, Some(action), false); + cell.complete(); + cell } /// If the first content is an image, return a new cell with the image. @@ -1837,6 +1919,7 @@ mod tests { use codex_core::config::types::McpServerConfig; use codex_core::config::types::McpServerTransportConfig; use codex_core::protocol::McpAuthStatus; + use codex_protocol::models::WebSearchAction; use codex_protocol::parse_command::ParsedCommand; use dirs::home_dir; use pretty_assertions::assert_eq; @@ -2060,8 +2143,12 @@ mod tests { #[test] fn web_search_history_cell_snapshot() { + let query = + "example search query with several generic words to exercise wrapping".to_string(); let cell = new_web_search_call( - "example search query with several generic words to exercise wrapping".to_string(), + "call-1".to_string(), + query.clone(), + WebSearchAction::Search { query: Some(query) }, ); let rendered = render_lines(&cell.display_lines(64)).join("\n"); @@ -2070,8 +2157,12 @@ mod tests { #[test] fn web_search_history_cell_wraps_with_indented_continuation() { + let query = + "example search query with several generic words to exercise wrapping".to_string(); let cell = new_web_search_call( - "example search query with several generic words to exercise wrapping".to_string(), + "call-1".to_string(), + query.clone(), + WebSearchAction::Search { query: Some(query) }, ); let rendered = render_lines(&cell.display_lines(64)); @@ -2086,7 +2177,12 @@ mod tests { #[test] fn web_search_history_cell_short_query_does_not_wrap() { - let cell = new_web_search_call("short query".to_string()); + let query = "short query".to_string(); + let cell = new_web_search_call( + "call-1".to_string(), + query.clone(), + WebSearchAction::Search { query: Some(query) }, + ); let rendered = render_lines(&cell.display_lines(64)); assert_eq!(rendered, vec!["• Searched short query".to_string()]); @@ -2094,8 +2190,12 @@ mod tests { #[test] fn web_search_history_cell_transcript_snapshot() { + let query = + "example search query with several generic words to exercise wrapping".to_string(); let cell = new_web_search_call( - "example search query with several generic words to exercise wrapping".to_string(), + "call-1".to_string(), + query.clone(), + WebSearchAction::Search { query: Some(query) }, ); let rendered = render_lines(&cell.transcript_lines(64)).join("\n");