diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 8062d60f4b..d719a520b7 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -5130,6 +5130,7 @@ dependencies = [ name = "codex-thread-store" version = "0.0.0" dependencies = [ + "caseless", "chrono", "codex-app-server-protocol", "codex-extension-items", diff --git a/codex-rs/app-server/tests/suite/v2/thread_read.rs b/codex-rs/app-server/tests/suite/v2/thread_read.rs index 24222b867d..3a624d5d52 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_read.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_read.rs @@ -623,6 +623,7 @@ async fn thread_turns_list_supports_requested_items_view() -> Result<()> { #[tokio::test] async fn thread_search_occurrences_reads_paginated_projection() -> Result<()> { + let partial_answer = "😀 ΟΣ ß partial ꟒"; let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; MockResponsesConfig::new(&server.uri()).write(codex_home.path())?; @@ -727,7 +728,7 @@ async fn thread_search_occurrences_reads_paginated_projection() -> Result<()> { moderation: None, id: "partial-1".to_string(), content: vec![AgentMessageContent::Text { - text: "stable partial answer".to_string(), + text: partial_answer.to_string(), }], phase: Some(MessagePhase::PartialAnswer), memory_citation: None, @@ -828,23 +829,34 @@ async fn thread_search_occurrences_reads_paginated_projection() -> Result<()> { assert_eq!(data[2].snippet_match_range.end, 15); assert_eq!(next_cursor, None); - let request_id = mcp - .send_thread_search_occurrences_request(ThreadSearchOccurrencesParams { - thread_id: thread_id.to_string(), - search_term: "partial".to_string(), - cursor: None, - limit: Some(3), - }) - .await?; - let ThreadSearchOccurrencesResponse { data, next_cursor } = - timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; - assert_eq!( - data.iter() - .map(|occurrence| (occurrence.item_id.as_str(), occurrence.snippet.as_str())) - .collect::>(), - vec![("partial-1", "stable partial answer")] - ); - assert_eq!(next_cursor, None); + for (term, expected_range) in [ + ("partial", 8..15), + ("Σ", 4..5), + ("σ", 4..5), + ("ς", 4..5), + ("ΟΣ", 3..5), + ("ss", 6..7), + ("ꟓ", 16..17), + ] { + let request_id = mcp + .send_thread_search_occurrences_request(ThreadSearchOccurrencesParams { + thread_id: thread_id.to_string(), + search_term: term.to_string(), + cursor: None, + limit: Some(3), + }) + .await?; + let ThreadSearchOccurrencesResponse { data, next_cursor } = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(request_id)).await??; + let [hit] = data.as_slice() else { + panic!("expected one match for {term}, got {data:?}"); + }; + assert_eq!(hit.item_id, "partial-1"); + assert_eq!(hit.snippet, partial_answer); + let range = &hit.snippet_match_range; + assert_eq!(range.start..range.end, expected_range, "search term {term}"); + assert_eq!(next_cursor, None); + } let fork_request_id = mcp .send_thread_fork_request(ThreadForkParams { diff --git a/codex-rs/thread-store/Cargo.toml b/codex-rs/thread-store/Cargo.toml index e0ba6c7df0..e885672752 100644 --- a/codex-rs/thread-store/Cargo.toml +++ b/codex-rs/thread-store/Cargo.toml @@ -13,6 +13,7 @@ doctest = false workspace = true [dependencies] +caseless = { workspace = true } chrono = { workspace = true, features = ["serde"] } codex-app-server-protocol = { workspace = true } codex-extension-items = { workspace = true } diff --git a/codex-rs/thread-store/src/local/thread_history/search.rs b/codex-rs/thread-store/src/local/thread_history/search.rs index e3896447a9..8a74437ea5 100644 --- a/codex-rs/thread-store/src/local/thread_history/search.rs +++ b/codex-rs/thread-store/src/local/thread_history/search.rs @@ -1,6 +1,7 @@ use std::borrow::Cow; use std::collections::HashMap; +use caseless::Caseless; use codex_app_server_protocol::ThreadItemPayload; use codex_app_server_protocol::UserInput; use codex_protocol::ThreadId; @@ -398,49 +399,50 @@ fn markdown_to_search_text(markdown: &str) -> String { } struct LiteralMatcher { - lowercase_needle: String, + folded_needle: String, +} + +fn fold_char(character: char) -> impl Iterator { + // Rust's lowercase mappings can cover newer Unicode characters than the folding tables. + character.to_lowercase().default_case_fold() } impl LiteralMatcher { fn new(needle: &str) -> Self { Self { - lowercase_needle: needle.to_lowercase(), + folded_needle: needle.chars().flat_map(fold_char).collect(), } } 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; + let mut folded_text = String::with_capacity(text.len()); + let mut boundaries = Vec::with_capacity(text.chars().count() + 1); 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; + boundaries.push((folded_text.len(), original_start)); + folded_text.extend(fold_char(character)); } + boundaries.push((folded_text.len(), text.len())); - // Use two-pointer method to find matches in linear time. - let mut start_span = 0; - let mut end_span = 0; - lowercase_text - .match_indices(self.lowercase_needle.as_str()) - .take(limit) - .filter_map(|(start, matched)| { - let end = start.saturating_add(matched.len()); - while spans.get(start_span)?.0.end <= start { - start_span += 1; - } - while spans.get(end_span)?.0.end <= end.saturating_sub(1) { - end_span += 1; - } - let original_start = spans.get(start_span)?.1.start; - let original_end = spans.get(end_span)?.1.end; - Some(original_start..original_end) - }) - .collect() + // Advance through source character boundaries even when a match ends partway through a fold. + let mut cursor = 0; + let mut matches = Vec::new(); + while matches.len() < limit { + let offset = boundaries[cursor].0; + let Some(relative_start) = folded_text[offset..].find(&self.folded_needle) else { + break; + }; + let start = offset + relative_start; + let end = start + self.folded_needle.len(); + while boundaries[cursor + 1].0 <= start { + cursor += 1; + } + let original_start = boundaries[cursor].1; + while boundaries[cursor].0 < end { + cursor += 1; + } + matches.push(original_start..boundaries[cursor].1); + } + matches } }