use std::borrow::Cow; use std::collections::HashMap; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::UserInput; use codex_protocol::ThreadId; use codex_protocol::protocol::strip_user_message_prefix; use futures::TryStreamExt; use pulldown_cmark::Event; use pulldown_cmark::Parser; use pulldown_cmark::TagEnd; use serde::Deserialize; use serde::Serialize; use sqlx::Row; use super::super::LocalThreadStore; use super::read::CursorScope; use super::read::serialize_cursor; use super::read::validate_thread_for_paginated_reads; use super::sqlite_integer; use super::thread_history_error; use super::turn_lookup::find_visible_turn; use crate::SearchTextRange; use crate::SearchThreadOccurrencesParams; use crate::StoredThreadOccurrence; use crate::ThreadOccurrenceSearchPage; use crate::ThreadStoreError; use crate::ThreadStoreResult; const SNIPPET_CONTEXT_BEFORE_CHARS: usize = 48; const SNIPPET_CONTEXT_AFTER_CHARS: usize = 96; #[derive(Deserialize, Serialize)] #[serde(rename_all = "camelCase")] struct SearchCursor { thread_id: ThreadId, search_term: String, next_rollout_ordinal: i64, next_occurrence_index: usize, } struct CandidateRow { turn_id: String, item_id: String, rollout_ordinal: i64, item_json: String, turn_rollout_ordinal: i64, } pub(in crate::local) async fn search_thread_occurrences( store: &LocalThreadStore, params: SearchThreadOccurrencesParams, ) -> ThreadStoreResult { if params.search_term.trim().is_empty() { return Err(ThreadStoreError::InvalidRequest { message: "thread/searchOccurrences requires search_term".to_string(), }); } if params.page_size == 0 { return Err(ThreadStoreError::InvalidRequest { message: "thread/searchOccurrences requires page_size greater than zero".to_string(), }); } validate_thread_for_paginated_reads( store, params.thread_id, /*include_archived*/ true, "thread/searchOccurrences", ) .await?; let cursor = parse_cursor( params.cursor.as_deref(), params.thread_id, ¶ms.search_term, )?; let matcher = LiteralMatcher::new(params.search_term.as_str()); let lineage = store.resolve_rollout_lineage(params.thread_id).await?; let cursor_segment = cursor .as_ref() .map(|cursor| { lineage .segment_index_for_ordinal( u64::try_from(cursor.next_rollout_ordinal) .map_err(|_| invalid_cursor("negative rollout ordinal"))?, ) .ok_or_else(|| invalid_cursor("position outside thread lineage")) }) .transpose()?; let pool = store.thread_history_db().await?; let mut items = Vec::with_capacity(params.page_size); let mut effective_turn_ordinals = HashMap::new(); for (segment_index, segment) in lineage .segments() .iter() .enumerate() .skip(cursor_segment.unwrap_or(0)) { let segment_start_ordinal = sqlite_integer(segment.start_ordinal(), "rollout ordinal")?; let next_rollout_ordinal = if Some(segment_index) == cursor_segment { cursor .as_ref() .map_or(0, |cursor| cursor.next_rollout_ordinal) } else { segment_start_ordinal }; let end_rollout_ordinal = segment .end_ordinal() .map(|ordinal| sqlite_integer(ordinal, "rollout ordinal")) .transpose()? .unwrap_or(i64::MAX); let mut rows = sqlx::query( r#" SELECT turn_id, item_id, rollout_ordinal, item_json, turn_rollout_ordinal FROM ( SELECT items.turn_id, items.item_id, items.rollout_ordinal, items.item_json, turns.rollout_ordinal AS turn_rollout_ordinal FROM thread_items AS items JOIN thread_turns AS turns ON turns.thread_id = items.thread_id AND turns.turn_id = items.turn_id WHERE items.thread_id = ? AND items.item_type = 'userMessage' AND items.rollout_ordinal >= ? AND items.rollout_ordinal < ? AND turns.rollout_ordinal >= ? AND turns.rollout_ordinal < ? UNION ALL SELECT items.turn_id, items.item_id, items.rollout_ordinal, items.item_json, turns.rollout_ordinal AS turn_rollout_ordinal FROM thread_turns AS turns JOIN thread_items AS items ON items.thread_id = turns.thread_id AND items.turn_id = turns.turn_id AND items.item_id = turns.final_agent_item_id WHERE turns.thread_id = ? AND turns.final_agent_item_id IS NOT NULL AND items.rollout_ordinal >= ? AND items.rollout_ordinal < ? AND turns.rollout_ordinal >= ? AND turns.rollout_ordinal < ? ) ORDER BY rollout_ordinal ASC "#, ) .bind(segment.rollout_id().to_string()) .bind(next_rollout_ordinal) .bind(end_rollout_ordinal) .bind(segment_start_ordinal) .bind(end_rollout_ordinal) .bind(segment.rollout_id().to_string()) .bind(next_rollout_ordinal) .bind(end_rollout_ordinal) .bind(segment_start_ordinal) .bind(end_rollout_ordinal) .fetch(pool); let mut matching_rows = Vec::new(); let mut matching_occurrences = 0_usize; while let Some(row) = rows.try_next().await.map_err(thread_history_error)? { let row = candidate_row(row)?; let item = serde_json::from_str::(row.item_json.as_str()).map_err(|err| { ThreadStoreError::Internal { message: format!("failed to deserialize stored thread item: {err}"), } })?; let Some(text) = searchable_text(&item) else { continue; }; let first_occurrence_index = cursor .as_ref() .filter(|cursor| cursor.next_rollout_ordinal == row.rollout_ordinal) .map_or(0, |cursor| cursor.next_occurrence_index); let remaining = params .page_size .saturating_add(1) .saturating_sub(items.len()) .saturating_sub(matching_occurrences); let matches = matcher.find_ranges( text.as_ref(), first_occurrence_index.saturating_add(remaining), ); if matches.len() <= first_occurrence_index { continue; } matching_occurrences += matches.len() - first_occurrence_index; matching_rows.push((row, text.into_owned(), matches, first_occurrence_index)); if matching_occurrences == params .page_size .saturating_add(1) .saturating_sub(items.len()) { break; } } drop(rows); for (row, text, matches, first_occurrence_index) in matching_rows { let turn_rollout_ordinal = if segment_index + 1 == lineage.segments().len() { row.turn_rollout_ordinal } else if let Some(ordinal) = effective_turn_ordinals.get(&row.turn_id) { *ordinal } else { let ordinal = find_visible_turn(pool, &lineage, row.turn_id.as_str()) .await? .rollout_ordinal; effective_turn_ordinals.insert(row.turn_id.clone(), ordinal); ordinal }; let turn_cursor = serialize_cursor( params.thread_id, CursorScope::Turns, turn_rollout_ordinal, /*include_anchor*/ true, )?; for (occurrence_index, matched) in matches.into_iter().enumerate().skip(first_occurrence_index) { if items.len() == params.page_size { return Ok(ThreadOccurrenceSearchPage { items, next_cursor: Some(serialize_cursor_for_search(SearchCursor { thread_id: params.thread_id, search_term: params.search_term, next_rollout_ordinal: row.rollout_ordinal, next_occurrence_index: occurrence_index, })?), }); } items.push(occurrence_in_item( row.turn_id.as_str(), row.item_id.as_str(), text.as_str(), matched, turn_cursor.as_str(), )); } } } Ok(ThreadOccurrenceSearchPage { items, next_cursor: None, }) } fn candidate_row(row: sqlx::sqlite::SqliteRow) -> ThreadStoreResult { let rollout_ordinal = row.try_get::("rollout_ordinal")?; let turn_rollout_ordinal = row.try_get::("turn_rollout_ordinal")?; if rollout_ordinal < 0 || turn_rollout_ordinal < 0 { return Err(ThreadStoreError::Internal { message: "invalid stored thread history ordinal".to_string(), }); } Ok(CandidateRow { turn_id: row.try_get("turn_id")?, item_id: row.try_get("item_id")?, rollout_ordinal, item_json: row.try_get("item_json")?, turn_rollout_ordinal, }) } fn parse_cursor( cursor: Option<&str>, thread_id: ThreadId, search_term: &str, ) -> ThreadStoreResult> { let Some(cursor) = cursor else { return Ok(None); }; let cursor_value: SearchCursor = serde_json::from_str(cursor).map_err(|_| invalid_cursor(cursor))?; if cursor_value.thread_id != thread_id || cursor_value.search_term != search_term || cursor_value.next_rollout_ordinal < 0 { return Err(invalid_cursor(cursor)); } Ok(Some(cursor_value)) } fn serialize_cursor_for_search(cursor: SearchCursor) -> ThreadStoreResult { serde_json::to_string(&cursor).map_err(thread_history_error) } fn invalid_cursor(cursor: &str) -> ThreadStoreError { ThreadStoreError::InvalidRequest { message: format!("invalid cursor: {cursor}"), } } fn searchable_text(item: &ThreadItem) -> Option> { match item { ThreadItem::UserMessage { content, .. } => { let mut text_parts = content .iter() .filter_map(|input| match input { UserInput::Text { text, .. } => Some(strip_user_message_prefix(text)), UserInput::Image { .. } | UserInput::LocalImage { .. } | UserInput::Audio { .. } | UserInput::LocalAudio { .. } | UserInput::Skill { .. } | UserInput::Mention { .. } => None, }) .filter(|text| !text.is_empty()) .peekable(); let first = text_parts.next()?; match text_parts.next() { None => Some(Cow::Borrowed(first)), Some(second) => { let mut parts = vec![first, second]; parts.extend(text_parts); Some(Cow::Owned(parts.concat())) } } } ThreadItem::AgentMessage { text, .. } => { let text = markdown_to_search_text(text); (!text.is_empty()).then_some(Cow::Owned(text)) } ThreadItem::HookPrompt { .. } | ThreadItem::Plan { .. } | ThreadItem::Reasoning { .. } | ThreadItem::CommandExecution { .. } | ThreadItem::FileChange { .. } | ThreadItem::McpToolCall { .. } | ThreadItem::DynamicToolCall { .. } | ThreadItem::CollabAgentToolCall { .. } | ThreadItem::SubAgentActivity { .. } | ThreadItem::WebSearch(_) | ThreadItem::ImageView { .. } | ThreadItem::Sleep(_) | ThreadItem::ImageGeneration(_) | ThreadItem::EnteredReviewMode { .. } | ThreadItem::ExitedReviewMode { .. } | ThreadItem::ContextCompaction { .. } => None, } } fn markdown_to_search_text(markdown: &str) -> String { let mut text = String::new(); for event in Parser::new(markdown.trim()) { match event { Event::Text(value) | Event::Code(value) | Event::Html(value) | Event::InlineHtml(value) => text.push_str(&value), Event::SoftBreak | Event::HardBreak | Event::Rule => text.push(' '), Event::End( TagEnd::Paragraph | TagEnd::Heading(_) | TagEnd::BlockQuote | TagEnd::CodeBlock | TagEnd::List(_) | TagEnd::Item | TagEnd::Table | TagEnd::TableHead | TagEnd::TableRow | TagEnd::TableCell, ) => text.push(' '), Event::Start(_) | Event::End( TagEnd::Emphasis | TagEnd::Strong | TagEnd::Strikethrough | TagEnd::Link | TagEnd::HtmlBlock | TagEnd::FootnoteDefinition | TagEnd::Image | TagEnd::MetadataBlock(_), ) | Event::FootnoteReference(_) | Event::TaskListMarker(_) => {} } } text.split_whitespace().collect::>().join(" ") } struct LiteralMatcher { lowercase_needle: String, } impl LiteralMatcher { fn new(needle: &str) -> Self { Self { lowercase_needle: needle.to_lowercase(), } } fn find_ranges(&self, text: &str, limit: usize) -> Vec> { let lowercase_text = text.to_lowercase(); let mut spans = Vec::with_capacity(text.chars().count()); let mut lowercase_start = 0; for (original_start, character) in text.char_indices() { let lowercase_end = lowercase_start + character.to_lowercase().map(char::len_utf8).sum::(); spans.push(( lowercase_start..lowercase_end, original_start..original_start + character.len_utf8(), )); lowercase_start = lowercase_end; } lowercase_text .match_indices(self.lowercase_needle.as_str()) .take(limit) .filter_map(|(start, matched)| { let end = start.saturating_add(matched.len()); let original_start = spans .iter() .find(|(lowercase, _)| lowercase.contains(&start))? .1 .start; let original_end = spans .iter() .find(|(lowercase, _)| lowercase.contains(&end.saturating_sub(1)))? .1 .end; Some(original_start..original_end) }) .collect() } } fn occurrence_in_item( turn_id: &str, item_id: &str, text: &str, matched: std::ops::Range, turn_cursor: &str, ) -> StoredThreadOccurrence { let snippet_start = char_start_before(text, matched.start, SNIPPET_CONTEXT_BEFORE_CHARS); let snippet_end = char_end_after(text, matched.end, SNIPPET_CONTEXT_AFTER_CHARS); let leading_ellipsis = snippet_start > 0; let trailing_ellipsis = snippet_end < text.len(); let mut snippet = String::new(); if leading_ellipsis { snippet.push_str("... "); } snippet.push_str(&text[snippet_start..snippet_end]); if trailing_ellipsis { snippet.push_str(" ..."); } let snippet_match_start = if leading_ellipsis { 4 } else { 0 } + utf16_len(&text[snippet_start..matched.start]); let match_len = utf16_len(&text[matched]); StoredThreadOccurrence { turn_id: turn_id.to_string(), item_id: item_id.to_string(), snippet, snippet_match_range: SearchTextRange { start: snippet_match_start, end: snippet_match_start.saturating_add(match_len), }, turn_cursor: turn_cursor.to_string(), } } fn utf16_len(text: &str) -> u32 { u32::try_from(text.encode_utf16().count()).unwrap_or(u32::MAX) } fn char_start_before(text: &str, byte_index: usize, chars_before: usize) -> usize { text[..byte_index] .char_indices() .rev() .nth(chars_before) .map(|(index, _)| index) .unwrap_or(0) } fn char_end_after(text: &str, byte_index: usize, chars_after: usize) -> usize { text[byte_index..] .char_indices() .nth(chars_after) .map(|(offset, _)| byte_index.saturating_add(offset)) .unwrap_or(text.len()) }