diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 9e250998ff..46b34cc693 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -4065,6 +4065,7 @@ name = "codex-thread-store" version = "0.0.0" dependencies = [ "chrono", + "codex-app-server-protocol", "codex-git-utils", "codex-install-context", "codex-otel", @@ -4073,6 +4074,7 @@ dependencies = [ "codex-state", "codex-utils-path", "pretty_assertions", + "pulldown-cmark", "serde", "serde_json", "tempfile", diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index b7526b080d..7a8f7ca30b 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -19325,6 +19325,70 @@ "title": "ThreadRollbackResponse", "type": "object" }, + "ThreadSearchOccurrence": { + "properties": { + "itemId": { + "type": "string" + }, + "matchRange": { + "allOf": [ + { + "$ref": "#/definitions/v2/ThreadSearchTextRange" + } + ], + "description": "Match range within the complete item text, in UTF-16 code units." + }, + "occurrenceId": { + "description": "Stable identity for this occurrence within the persisted query snapshot.", + "type": "string" + }, + "occurrenceIndex": { + "description": "Zero-based occurrence index within the user or assistant message item.", + "format": "uint32", + "minimum": 0.0, + "type": "integer" + }, + "snippet": { + "type": "string" + }, + "snippetMatchRange": { + "allOf": [ + { + "$ref": "#/definitions/v2/ThreadSearchTextRange" + } + ], + "description": "Match range within `snippet`, in UTF-16 code units." + }, + "threadId": { + "type": "string" + }, + "turnCursor": { + "description": "Opaque inclusive cursor accepted by `thread/turns/list` for this turn.", + "type": "string" + }, + "turnId": { + "type": "string" + }, + "turnStartedAt": { + "description": "Unix timestamp in seconds used to place the matching turn chronologically.", + "format": "int64", + "type": "integer" + } + }, + "required": [ + "itemId", + "matchRange", + "occurrenceId", + "occurrenceIndex", + "snippet", + "snippetMatchRange", + "threadId", + "turnCursor", + "turnId", + "turnStartedAt" + ], + "type": "object" + }, "ThreadSearchResult": { "properties": { "snippet": { @@ -19340,6 +19404,27 @@ ], "type": "object" }, + "ThreadSearchTextRange": { + "properties": { + "end": { + "description": "Exclusive UTF-16 code-unit offset.", + "format": "uint32", + "minimum": 0.0, + "type": "integer" + }, + "start": { + "description": "Inclusive UTF-16 code-unit offset.", + "format": "uint32", + "minimum": 0.0, + "type": "integer" + } + }, + "required": [ + "end", + "start" + ], + "type": "object" + }, "ThreadSetNameParams": { "$schema": "http://json-schema.org/draft-07/schema#", "properties": { diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index e1afc7883a..af1d2cfa2a 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -17104,6 +17104,70 @@ "title": "ThreadRollbackResponse", "type": "object" }, + "ThreadSearchOccurrence": { + "properties": { + "itemId": { + "type": "string" + }, + "matchRange": { + "allOf": [ + { + "$ref": "#/definitions/ThreadSearchTextRange" + } + ], + "description": "Match range within the complete item text, in UTF-16 code units." + }, + "occurrenceId": { + "description": "Stable identity for this occurrence within the persisted query snapshot.", + "type": "string" + }, + "occurrenceIndex": { + "description": "Zero-based occurrence index within the user or assistant message item.", + "format": "uint32", + "minimum": 0.0, + "type": "integer" + }, + "snippet": { + "type": "string" + }, + "snippetMatchRange": { + "allOf": [ + { + "$ref": "#/definitions/ThreadSearchTextRange" + } + ], + "description": "Match range within `snippet`, in UTF-16 code units." + }, + "threadId": { + "type": "string" + }, + "turnCursor": { + "description": "Opaque inclusive cursor accepted by `thread/turns/list` for this turn.", + "type": "string" + }, + "turnId": { + "type": "string" + }, + "turnStartedAt": { + "description": "Unix timestamp in seconds used to place the matching turn chronologically.", + "format": "int64", + "type": "integer" + } + }, + "required": [ + "itemId", + "matchRange", + "occurrenceId", + "occurrenceIndex", + "snippet", + "snippetMatchRange", + "threadId", + "turnCursor", + "turnId", + "turnStartedAt" + ], + "type": "object" + }, "ThreadSearchResult": { "properties": { "snippet": { @@ -17119,6 +17183,27 @@ ], "type": "object" }, + "ThreadSearchTextRange": { + "properties": { + "end": { + "description": "Exclusive UTF-16 code-unit offset.", + "format": "uint32", + "minimum": 0.0, + "type": "integer" + }, + "start": { + "description": "Inclusive UTF-16 code-unit offset.", + "format": "uint32", + "minimum": 0.0, + "type": "integer" + } + }, + "required": [ + "end", + "start" + ], + "type": "object" + }, "ThreadSetNameParams": { "$schema": "http://json-schema.org/draft-07/schema#", "properties": { diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadSearchOccurrence.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadSearchOccurrence.ts new file mode 100644 index 0000000000..eba8c49dbf --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadSearchOccurrence.ts @@ -0,0 +1,30 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. +import type { ThreadSearchTextRange } from "./ThreadSearchTextRange"; + +export type ThreadSearchOccurrence = { threadId: string, turnId: string, itemId: string, +/** + * Stable identity for this occurrence within the persisted query snapshot. + */ +occurrenceId: string, +/** + * Zero-based occurrence index within the user or assistant message item. + */ +occurrenceIndex: number, snippet: string, +/** + * Match range within the complete item text, in UTF-16 code units. + */ +matchRange: ThreadSearchTextRange, +/** + * Match range within `snippet`, in UTF-16 code units. + */ +snippetMatchRange: ThreadSearchTextRange, +/** + * Unix timestamp in seconds used to place the matching turn chronologically. + */ +turnStartedAt: number, +/** + * Opaque inclusive cursor accepted by `thread/turns/list` for this turn. + */ +turnCursor: string, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadSearchTextRange.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadSearchTextRange.ts new file mode 100644 index 0000000000..89b83c7e3e --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadSearchTextRange.ts @@ -0,0 +1,13 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +export type ThreadSearchTextRange = { +/** + * Inclusive UTF-16 code-unit offset. + */ +start: number, +/** + * Exclusive UTF-16 code-unit offset. + */ +end: number, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts index 070d21eab8..1d4389cbdc 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts @@ -446,7 +446,9 @@ export type { ThreadResumeParams } from "./ThreadResumeParams"; export type { ThreadResumeResponse } from "./ThreadResumeResponse"; export type { ThreadRollbackParams } from "./ThreadRollbackParams"; export type { ThreadRollbackResponse } from "./ThreadRollbackResponse"; +export type { ThreadSearchOccurrence } from "./ThreadSearchOccurrence"; export type { ThreadSearchResult } from "./ThreadSearchResult"; +export type { ThreadSearchTextRange } from "./ThreadSearchTextRange"; export type { ThreadSetNameParams } from "./ThreadSetNameParams"; export type { ThreadSetNameResponse } from "./ThreadSetNameResponse"; export type { ThreadSettings } from "./ThreadSettings"; diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index 27a8b0205d..177ad74338 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -630,6 +630,13 @@ client_request_definitions! { serialization: None, response: v2::ThreadSearchResponse, }, + #[experimental("thread/searchOccurrences")] + ThreadSearchOccurrences => "thread/searchOccurrences" { + params: v2::ThreadSearchOccurrencesParams, + // Explicitly concurrent: this reads an immutable persisted-history snapshot. + serialization: None, + response: v2::ThreadSearchOccurrencesResponse, + }, ThreadLoadedList => "thread/loaded/list" { params: v2::ThreadLoadedListParams, serialization: None, diff --git a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs index 377ecba2c2..7e1e94ebfc 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs @@ -299,6 +299,63 @@ fn thread_items_list_round_trips() { ); } +#[test] +fn thread_search_occurrences_round_trips_anchors() { + let params = ThreadSearchOccurrencesParams { + thread_id: "thr_123".to_string(), + search_term: "needle".to_string(), + cursor: Some("search-cursor".to_string()), + limit: Some(50), + }; + + assert_eq!( + serde_json::to_value(¶ms).expect("serialize params"), + json!({ + "threadId": "thr_123", + "searchTerm": "needle", + "cursor": "search-cursor", + "limit": 50, + }) + ); + + let response = ThreadSearchOccurrencesResponse { + data: vec![ThreadSearchOccurrence { + thread_id: "thr_123".to_string(), + turn_id: "turn_456".to_string(), + item_id: "item-2".to_string(), + occurrence_id: "turn_456:item-2:7".to_string(), + occurrence_index: 0, + snippet: "before needle after".to_string(), + match_range: ThreadSearchTextRange { start: 7, end: 13 }, + snippet_match_range: ThreadSearchTextRange { start: 7, end: 13 }, + turn_started_at: 1_700_000_000, + turn_cursor: r#"{"turnId":"turn_456","includeAnchor":true}"#.to_string(), + }], + next_cursor: Some("next-search-cursor".to_string()), + is_capped: false, + }; + + assert_eq!( + serde_json::to_value(&response).expect("serialize response"), + json!({ + "data": [{ + "threadId": "thr_123", + "turnId": "turn_456", + "itemId": "item-2", + "occurrenceId": "turn_456:item-2:7", + "occurrenceIndex": 0, + "snippet": "before needle after", + "matchRange": {"start": 7, "end": 13}, + "snippetMatchRange": {"start": 7, "end": 13}, + "turnStartedAt": 1_700_000_000, + "turnCursor": r#"{"turnId":"turn_456","includeAnchor":true}"#, + }], + "nextCursor": "next-search-cursor", + "isCapped": false, + }) + ); +} + #[test] fn thread_list_params_accepts_single_cwd() { let params = serde_json::from_value::(json!({ diff --git a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs index 367a4a64c8..72eeacd186 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs @@ -1223,6 +1223,65 @@ pub struct ThreadSearchResponse { pub backwards_cursor: Option, } +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ThreadSearchOccurrencesParams { + pub thread_id: String, + /// Case-insensitive literal substring to find in user-visible user and assistant messages. + pub search_term: String, + /// Opaque pagination cursor returned by a previous call for the same thread and search term. + #[ts(optional = nullable)] + pub cursor: Option, + /// Optional page size; the server caps a query generation at 250 occurrences. + #[ts(optional = nullable)] + pub limit: Option, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ThreadSearchTextRange { + /// Inclusive UTF-16 code-unit offset. + pub start: u32, + /// Exclusive UTF-16 code-unit offset. + pub end: u32, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ThreadSearchOccurrence { + pub thread_id: String, + pub turn_id: String, + pub item_id: String, + /// Stable identity for this occurrence within the persisted query snapshot. + pub occurrence_id: String, + /// Zero-based occurrence index within the user or assistant message item. + pub occurrence_index: u32, + pub snippet: String, + /// Match range within the complete item text, in UTF-16 code units. + pub match_range: ThreadSearchTextRange, + /// Match range within `snippet`, in UTF-16 code units. + pub snippet_match_range: ThreadSearchTextRange, + /// Unix timestamp in seconds used to place the matching turn chronologically. + #[ts(type = "number")] + pub turn_started_at: i64, + /// Opaque inclusive cursor accepted by `thread/turns/list` for this turn. + pub turn_cursor: String, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ThreadSearchOccurrencesResponse { + pub data: Vec, + /// Opaque cursor for the next page of this point-in-time query snapshot. + pub next_cursor: Option, + /// Whether this query generation reached the server's 250-occurrence cap. + pub is_capped: bool, +} + #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, JsonSchema, TS)] #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 6ee6baf2dc..fbb12a6b51 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -145,6 +145,7 @@ Example with notification opt-out: - `thread/loaded/list` — list the thread ids currently loaded in memory. - `thread/read` — read a stored thread by id without resuming it; optionally include turns via `includeTurns`. The returned `thread` includes `status` (`ThreadStatus`), defaulting to `notLoaded` when the thread is not currently loaded. - `thread/turns/list` — experimental; page through a stored thread’s turn history without resuming it; supports cursor-based pagination with `sortDirection`, `itemsView`, `nextCursor`, and `backwardsCursor`. +- `thread/searchOccurrences` — experimental; find case-insensitive literal occurrences in one stored thread’s visible user and assistant messages. Results are chronological, point-in-time paginated, capped at 250 occurrences, and include an inclusive cursor accepted by `thread/turns/list` for bounded hydration around a selected match. - `thread/items/list` — experimental; page through persisted thread items without resuming the thread. Pass `turnId` to restrict results to one turn, or omit it to page items across the thread. The active thread store must support item pagination. - `thread/metadata/update` — patch stored thread metadata in sqlite; currently supports updating persisted `gitInfo` fields and returns the refreshed `thread`. - `thread/settings/update` — experimental; queue a partial update to a loaded thread’s next-turn settings without starting a turn or adding transcript items. Omitted fields leave settings unchanged; `serviceTier: null` clears the tier; deprecated `multiAgentMode` is ignored, while Ultra reasoning effort enables proactive multi-agent behavior; `sandboxPolicy` and `permissions` cannot be combined. Returns `{}` when the update is accepted and emits `thread/settings/updated` with the full effective settings only if they actually change. `turn/start` settings overrides emit the same notification when they change the stored settings. @@ -515,6 +516,38 @@ Every returned `Turn` includes `itemsView`, which tells clients whether the `ite } } ``` +Search one thread without hydrating its turns by calling `thread/searchOccurrences`. Matching uses +the current Cmd-F semantics over visible user text and the final assistant message in each turn: +lowercase literal substrings over the client's Markdown-to-search-text projection. Each result +contains UTF-16 ranges for the full item and snippet plus a `turnCursor` that +can be passed directly to `thread/turns/list` with either sort direction; the anchor turn is included +in that page. Unconsumed search pagination cursors remain valid for ten minutes within the running +app-server process; a completed page remains retryable for 30 seconds. + +```json +{ "method": "thread/searchOccurrences", "id": 25, "params": { + "threadId": "thr_123", + "searchTerm": "needle", + "limit": 50 +} } +{ "id": 25, "result": { + "data": [{ + "threadId": "thr_123", + "turnId": "turn_456", + "itemId": "item-3", + "occurrenceId": "turn_456:item-3:120", + "occurrenceIndex": 0, + "snippet": "... surrounding needle text ...", + "matchRange": { "start": 120, "end": 126 }, + "snippetMatchRange": { "start": 16, "end": 22 }, + "turnStartedAt": 1736078400, + "turnCursor": "opaque-inclusive-turn-cursor" + }], + "nextCursor": "next-search-page-or-null", + "isCapped": false +} } +``` + `thread/items/list` pages full persisted items across a thread, optionally filtered to one turn: ```json diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index ca122920bd..e7f5acb061 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -1140,6 +1140,11 @@ impl MessageProcessor { ClientRequest::ThreadTurnsList { params, .. } => { self.thread_processor.thread_turns_list(params).await } + ClientRequest::ThreadSearchOccurrences { params, .. } => { + self.thread_processor + .thread_search_occurrences(params) + .await + } ClientRequest::ThreadItemsList { params, .. } => { self.thread_processor.thread_items_list(params).await } diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 42ec2063b4..20db88d9ed 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -251,9 +251,13 @@ use codex_app_server_protocol::ThreadResumeInitialTurnsPageParams; use codex_app_server_protocol::ThreadResumeParams; use codex_app_server_protocol::ThreadResumeResponse; use codex_app_server_protocol::ThreadRollbackParams; +use codex_app_server_protocol::ThreadSearchOccurrence; +use codex_app_server_protocol::ThreadSearchOccurrencesParams; +use codex_app_server_protocol::ThreadSearchOccurrencesResponse; use codex_app_server_protocol::ThreadSearchParams; use codex_app_server_protocol::ThreadSearchResponse; use codex_app_server_protocol::ThreadSearchResult; +use codex_app_server_protocol::ThreadSearchTextRange; use codex_app_server_protocol::ThreadSetNameParams; use codex_app_server_protocol::ThreadSetNameResponse; use codex_app_server_protocol::ThreadSettings; @@ -457,6 +461,7 @@ use codex_thread_store::ListThreadsParams as StoreListThreadsParams; use codex_thread_store::LocalThreadStore; use codex_thread_store::ReadThreadByRolloutPathParams as StoreReadThreadByRolloutPathParams; use codex_thread_store::ReadThreadParams as StoreReadThreadParams; +use codex_thread_store::SearchThreadOccurrencesParams as StoreSearchThreadOccurrencesParams; use codex_thread_store::SearchThreadsParams as StoreSearchThreadsParams; use codex_thread_store::SortDirection as StoreSortDirection; use codex_thread_store::StoredThread; diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index af8525873f..c367deeb6d 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -707,6 +707,15 @@ impl ThreadRequestProcessor { .map(|response| Some(response.into())) } + pub(crate) async fn thread_search_occurrences( + &self, + params: ThreadSearchOccurrencesParams, + ) -> Result, JSONRPCErrorError> { + self.thread_search_occurrences_response_inner(params) + .await + .map(|response| Some(response.into())) + } + pub(crate) async fn thread_items_list( &self, params: ThreadItemsListParams, @@ -2440,6 +2449,73 @@ impl ThreadRequestProcessor { ) } + async fn thread_search_occurrences_response_inner( + &self, + params: ThreadSearchOccurrencesParams, + ) -> Result { + let ThreadSearchOccurrencesParams { + thread_id, + search_term, + cursor, + limit, + } = params; + let thread_id = ThreadId::from_string(&thread_id) + .map_err(|err| invalid_request(format!("invalid thread id: {err}")))?; + if search_term.is_empty() { + return Err(invalid_request( + "thread/searchOccurrences requires a non-empty searchTerm", + )); + } + let page_size = limit + .map(|value| value as usize) + .unwrap_or(THREAD_SEARCH_OCCURRENCES_DEFAULT_LIMIT) + .clamp(1, THREAD_SEARCH_OCCURRENCES_MAX_RESULTS); + let page = self + .thread_store + .search_thread_occurrences(StoreSearchThreadOccurrencesParams { + thread_id, + search_term, + cursor, + page_size, + max_results: THREAD_SEARCH_OCCURRENCES_MAX_RESULTS, + }) + .await + .map_err(thread_search_occurrences_error)?; + let data = page + .items + .into_iter() + .map(|item| { + Ok(ThreadSearchOccurrence { + thread_id: item.thread_id.to_string(), + turn_cursor: serialize_thread_turns_cursor( + item.turn_id.as_str(), + /*include_anchor*/ true, + )?, + turn_id: item.turn_id, + item_id: item.item_id, + occurrence_id: item.occurrence_id, + occurrence_index: item.occurrence_index, + snippet: item.snippet, + match_range: ThreadSearchTextRange { + start: item.match_range.start, + end: item.match_range.end, + }, + snippet_match_range: ThreadSearchTextRange { + start: item.snippet_match_range.start, + end: item.snippet_match_range.end, + }, + turn_started_at: item.turn_started_at, + }) + }) + .collect::, JSONRPCErrorError>>()?; + + Ok(ThreadSearchOccurrencesResponse { + data, + next_cursor: page.next_cursor, + is_capped: page.is_capped, + }) + } + async fn thread_items_list_response_inner( &self, params: ThreadItemsListParams, @@ -3846,6 +3922,8 @@ fn xcode_26_4_mcp_elicitations_auto_deny( const THREAD_TURNS_DEFAULT_LIMIT: usize = 25; const THREAD_TURNS_MAX_LIMIT: usize = 100; +const THREAD_SEARCH_OCCURRENCES_DEFAULT_LIMIT: usize = 50; +const THREAD_SEARCH_OCCURRENCES_MAX_RESULTS: usize = 250; const THREAD_ITEMS_DEFAULT_LIMIT: usize = 25; const THREAD_ITEMS_MAX_LIMIT: usize = 100; @@ -4163,6 +4241,19 @@ fn thread_turns_list_history_load_error( } } +fn thread_search_occurrences_error(err: ThreadStoreError) -> JSONRPCErrorError { + match err { + ThreadStoreError::InvalidRequest { message } => invalid_request(message), + ThreadStoreError::Unsupported { operation } => { + unsupported_thread_store_operation(operation) + } + ThreadStoreError::ThreadNotFound { thread_id } => { + invalid_request(format!("no rollout found for thread id {thread_id}")) + } + err => internal_error(format!("failed to search thread occurrences: {err}")), + } +} + fn thread_read_history_load_error( thread_id: ThreadId, err: ThreadStoreError, diff --git a/codex-rs/app-server/tests/common/test_app_server.rs b/codex-rs/app-server/tests/common/test_app_server.rs index e080ab0261..9135e24dbb 100644 --- a/codex-rs/app-server/tests/common/test_app_server.rs +++ b/codex-rs/app-server/tests/common/test_app_server.rs @@ -100,6 +100,7 @@ use codex_app_server_protocol::ThreadRealtimeStartParams; use codex_app_server_protocol::ThreadRealtimeStopParams; use codex_app_server_protocol::ThreadResumeParams; use codex_app_server_protocol::ThreadRollbackParams; +use codex_app_server_protocol::ThreadSearchOccurrencesParams; use codex_app_server_protocol::ThreadSearchParams; use codex_app_server_protocol::ThreadSetNameParams; use codex_app_server_protocol::ThreadSettingsUpdateParams; @@ -747,6 +748,15 @@ impl TestAppServer { self.send_request("thread/turns/list", params).await } + /// Send a `thread/searchOccurrences` JSON-RPC request. + pub async fn send_thread_search_occurrences_request( + &mut self, + params: ThreadSearchOccurrencesParams, + ) -> anyhow::Result { + let params = Some(serde_json::to_value(params)?); + self.send_request("thread/searchOccurrences", params).await + } + /// Send a `thread/items/list` JSON-RPC request. pub async fn send_thread_items_list_request( &mut self, 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 4162e7ebd6..cfd76602a9 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_read.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_read.rs @@ -1,5 +1,6 @@ use anyhow::Result; use app_test_support::TestAppServer; +use app_test_support::create_fake_rollout; use app_test_support::create_fake_rollout_with_text_elements; use app_test_support::create_mock_responses_server_repeating_assistant; use app_test_support::rollout_path; @@ -29,6 +30,8 @@ use codex_app_server_protocol::ThreadReadResponse; use codex_app_server_protocol::ThreadResumeInitialTurnsPageParams; use codex_app_server_protocol::ThreadResumeParams; use codex_app_server_protocol::ThreadResumeResponse; +use codex_app_server_protocol::ThreadSearchOccurrencesParams; +use codex_app_server_protocol::ThreadSearchOccurrencesResponse; use codex_app_server_protocol::ThreadSetNameParams; use codex_app_server_protocol::ThreadSetNameResponse; use codex_app_server_protocol::ThreadStartParams; @@ -139,6 +142,106 @@ async fn thread_read_returns_summary_without_turns() -> Result<()> { Ok(()) } +#[tokio::test] +async fn thread_search_occurrences_returns_hydratable_turn_cursor() -> Result<()> { + let codex_home = TempDir::new()?; + let thread_id = create_fake_rollout( + codex_home.path(), + "2025-01-05T12-00-00", + "2025-01-05T12:00:00Z", + "Needle then NEEDLE", + Some("mock_provider"), + None, + )?; + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .build() + .await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let search_id = mcp + .send_thread_search_occurrences_request(ThreadSearchOccurrencesParams { + thread_id: thread_id.clone(), + search_term: "needle".to_string(), + cursor: None, + limit: Some(1), + }) + .await?; + let search_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(search_id)), + ) + .await??; + let ThreadSearchOccurrencesResponse { + data, + next_cursor, + is_capped, + } = to_response::(search_resp)?; + assert_eq!(data.len(), 1); + assert_eq!(data[0].thread_id, thread_id); + assert_eq!(data[0].occurrence_index, 0); + assert_eq!(data[0].match_range.start, 0); + assert_eq!(data[0].match_range.end, 6); + let next_cursor = next_cursor.expect("second search page cursor"); + assert!(!is_capped); + + let next_search_id = mcp + .send_thread_search_occurrences_request(ThreadSearchOccurrencesParams { + thread_id: thread_id.clone(), + search_term: "needle".to_string(), + cursor: Some(next_cursor), + limit: Some(1), + }) + .await?; + let next_search_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(next_search_id)), + ) + .await??; + let next_page = to_response::(next_search_resp)?; + assert_eq!(next_page.data.len(), 1); + assert_eq!(next_page.data[0].occurrence_index, 1); + assert_eq!(next_page.next_cursor, None); + + let turns_id = mcp + .send_thread_turns_list_request(ThreadTurnsListParams { + thread_id: thread_id.clone(), + cursor: Some(data[0].turn_cursor.clone()), + limit: Some(1), + sort_direction: Some(SortDirection::Asc), + items_view: Some(TurnItemsView::NotLoaded), + }) + .await?; + let turns_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turns_id)), + ) + .await??; + let turns = to_response::(turns_resp)?; + assert_eq!(turns.data.len(), 1); + assert_eq!(turns.data[0].id, data[0].turn_id); + + let literal_whitespace_id = mcp + .send_thread_search_occurrences_request(ThreadSearchOccurrencesParams { + thread_id, + search_term: " needle ".to_string(), + cursor: None, + limit: Some(10), + }) + .await?; + let literal_whitespace_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(literal_whitespace_id)), + ) + .await??; + assert_eq!( + to_response::(literal_whitespace_resp)?.data, + Vec::new() + ); + + Ok(()) +} + #[tokio::test] async fn thread_read_can_include_turns() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; diff --git a/codex-rs/rollout/src/compression.rs b/codex-rs/rollout/src/compression.rs index e64cc318a6..fc8a7a9948 100644 --- a/codex-rs/rollout/src/compression.rs +++ b/codex-rs/rollout/src/compression.rs @@ -9,6 +9,7 @@ use std::path::PathBuf; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use std::time::Duration; +use std::time::SystemTime; #[cfg(unix)] use std::os::unix::fs::OpenOptionsExt; @@ -193,6 +194,9 @@ impl RolloutFile { /// Line-oriented rollout reader returned by [`open_rollout_line_reader`]. pub struct RolloutLineReader { inner: RolloutLineReaderInner, + source_path: PathBuf, + source_len: u64, + source_modified_at: Option, } enum RolloutLineReaderInner { @@ -201,6 +205,21 @@ enum RolloutLineReaderInner { } impl RolloutLineReader { + /// Returns the physical rollout representation opened by this reader. + pub fn source_path(&self) -> &Path { + self.source_path.as_path() + } + + /// Returns the source length captured when this reader opened the rollout. + pub fn source_len(&self) -> u64 { + self.source_len + } + + /// Returns the source modification time captured when this reader opened the rollout. + pub fn source_modified_at(&self) -> Option { + self.source_modified_at + } + /// Reads the next JSONL record from the rollout. pub async fn next_line(&mut self) -> io::Result> { match &mut self.inner { @@ -987,22 +1006,31 @@ mod reader { .await .unwrap_or_else(|| path.to_path_buf()); if path::is_compressed_rollout_path(path.as_path()) { + let source_path = path.clone(); let reader = tokio::task::spawn_blocking(move || { let input = File::open(path.as_path())?; + let metadata = input.metadata()?; let decoder = zstd::stream::read::Decoder::new(input)?; - Ok::<_, io::Error>( - io::BufReader::new(Box::new(decoder) as Box).lines(), - ) + let lines = io::BufReader::new(Box::new(decoder) as Box).lines(); + Ok::<_, io::Error>((lines, metadata.len(), metadata.modified().ok())) }) .await .map_err(io::Error::other)??; + let (reader, source_len, source_modified_at) = reader; return Ok(RolloutLineReader { inner: RolloutLineReaderInner::Blocking(Some(reader)), + source_path, + source_len, + source_modified_at, }); } - let file = tokio::fs::File::open(path).await?; + let file = tokio::fs::File::open(path.as_path()).await?; + let metadata = file.metadata().await?; Ok(RolloutLineReader { inner: RolloutLineReaderInner::Plain(tokio::io::BufReader::new(file).lines()), + source_path: path, + source_len: metadata.len(), + source_modified_at: metadata.modified().ok(), }) } } diff --git a/codex-rs/thread-store/Cargo.toml b/codex-rs/thread-store/Cargo.toml index 35b3919c60..752d4a0a41 100644 --- a/codex-rs/thread-store/Cargo.toml +++ b/codex-rs/thread-store/Cargo.toml @@ -14,6 +14,7 @@ workspace = true [dependencies] chrono = { workspace = true, features = ["serde"] } +codex-app-server-protocol = { workspace = true } codex-git-utils = { workspace = true } codex-install-context = { workspace = true } codex-otel = { workspace = true } @@ -21,6 +22,7 @@ codex-protocol = { workspace = true } codex-rollout = { workspace = true } codex-state = { workspace = true } codex-utils-path = { workspace = true } +pulldown-cmark = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } thiserror = { workspace = true } diff --git a/codex-rs/thread-store/src/lib.rs b/codex-rs/thread-store/src/lib.rs index 34a0ab2ba1..65ffd01cb8 100644 --- a/codex-rs/thread-store/src/lib.rs +++ b/codex-rs/thread-store/src/lib.rs @@ -37,17 +37,21 @@ pub use types::LoadThreadHistoryParams; pub use types::ReadThreadByRolloutPathParams; pub use types::ReadThreadParams; pub use types::ResumeThreadParams; +pub use types::SearchTextRange; +pub use types::SearchThreadOccurrencesParams; pub use types::SearchThreadsParams; pub use types::SortDirection; pub use types::StoredThread; pub use types::StoredThreadHistory; pub use types::StoredThreadItem; +pub use types::StoredThreadOccurrence; pub use types::StoredThreadSearchResult; pub use types::StoredTurn; pub use types::StoredTurnError; pub use types::StoredTurnItemsView; pub use types::StoredTurnStatus; pub use types::ThreadMetadataPatch; +pub use types::ThreadOccurrenceSearchPage; pub use types::ThreadPage; pub use types::ThreadPersistenceMetadata; pub use types::ThreadRelationFilter; diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 5939ea606e..1c4996d94c 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -5,6 +5,7 @@ mod helpers; mod list_threads; mod live_writer; mod read_thread; +mod search_thread_occurrences; mod search_threads; mod unarchive_thread; mod update_thread_metadata; @@ -31,9 +32,11 @@ use crate::LoadThreadHistoryParams; use crate::ReadThreadByRolloutPathParams; use crate::ReadThreadParams; use crate::ResumeThreadParams; +use crate::SearchThreadOccurrencesParams; use crate::SearchThreadsParams; use crate::StoredThread; use crate::StoredThreadHistory; +use crate::ThreadOccurrenceSearchPage; use crate::ThreadPage; use crate::ThreadSearchPage; use crate::ThreadStore; @@ -59,6 +62,7 @@ use crate::UpdateThreadMetadataParams; pub struct LocalThreadStore { pub(super) config: LocalThreadStoreConfig, live_recorders: Arc>>, + occurrence_search_cache: Arc>, state_db: Option, } @@ -106,6 +110,7 @@ impl LocalThreadStore { Self { config, live_recorders: Arc::new(Mutex::new(HashMap::new())), + occurrence_search_cache: Arc::new(Mutex::new(Default::default())), state_db, } } @@ -301,6 +306,15 @@ impl ThreadStore for LocalThreadStore { Box::pin(async move { search_threads::search_threads(self, params).await }) } + fn search_thread_occurrences( + &self, + params: SearchThreadOccurrencesParams, + ) -> ThreadStoreFuture<'_, ThreadOccurrenceSearchPage> { + Box::pin( + async move { search_thread_occurrences::search_thread_occurrences(self, params).await }, + ) + } + fn update_thread_metadata( &self, params: UpdateThreadMetadataParams, diff --git a/codex-rs/thread-store/src/local/read_thread.rs b/codex-rs/thread-store/src/local/read_thread.rs index 27ce3076e9..f255b8bd89 100644 --- a/codex-rs/thread-store/src/local/read_thread.rs +++ b/codex-rs/thread-store/src/local/read_thread.rs @@ -206,7 +206,7 @@ async fn attach_history_if_requested( Ok(()) } -async fn resolve_rollout_path( +pub(super) async fn resolve_rollout_path( store: &LocalThreadStore, thread_id: codex_protocol::ThreadId, include_archived: bool, diff --git a/codex-rs/thread-store/src/local/search_thread_occurrences.rs b/codex-rs/thread-store/src/local/search_thread_occurrences.rs new file mode 100644 index 0000000000..83610fd219 --- /dev/null +++ b/codex-rs/thread-store/src/local/search_thread_occurrences.rs @@ -0,0 +1,634 @@ +use std::borrow::Cow; +use std::collections::HashMap; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; +use std::time::SystemTime; + +use chrono::DateTime; +use codex_app_server_protocol::ThreadHistoryBuilder; +use codex_app_server_protocol::ThreadItem; +use codex_app_server_protocol::UserInput; +use codex_protocol::protocol::RolloutLine; +use codex_protocol::protocol::USER_MESSAGE_BEGIN; +use pulldown_cmark::Event; +use pulldown_cmark::Parser; +use pulldown_cmark::TagEnd; +use serde::Deserialize; +use serde::Serialize; + +use super::LocalThreadStore; +use crate::SearchTextRange; +use crate::SearchThreadOccurrencesParams; +use crate::StoredThreadOccurrence; +use crate::ThreadOccurrenceSearchPage; +use crate::ThreadStoreError; +use crate::ThreadStoreResult; + +const SNAPSHOT_TTL: Duration = Duration::from_secs(10 * 60); +const COMPLETED_SNAPSHOT_TTL: Duration = Duration::from_secs(30); +const MAX_ACTIVE_SNAPSHOTS: usize = 256; +const MAX_CACHED_PROJECTIONS: usize = 8; +const SNIPPET_CONTEXT_BEFORE_CHARS: usize = 48; +const SNIPPET_CONTEXT_AFTER_CHARS: usize = 96; + +#[cfg(test)] +#[path = "search_thread_occurrences_tests.rs"] +mod tests; + +#[derive(Default)] +pub(super) struct ThreadOccurrenceSearchCache { + next_snapshot_id: u64, + snapshots: HashMap, + projections: HashMap, +} + +#[derive(Clone)] +struct ThreadOccurrenceSearchSnapshot { + thread_id: codex_protocol::ThreadId, + search_term: String, + items: Vec, + is_capped: bool, + expires_at: Instant, + completed: bool, +} + +#[derive(Clone, PartialEq, Eq)] +struct RolloutFingerprint { + path: PathBuf, + len: u64, + modified_at: Option, +} + +#[derive(Clone)] +struct SearchProjection { + fingerprint: RolloutFingerprint, + items: Vec, +} + +struct CachedSearchProjection { + projection: Arc, + expires_at: Instant, +} + +#[derive(Clone)] +struct SearchableThreadItem { + turn_id: String, + item_id: String, + text: String, + turn_started_at: i64, +} + +#[derive(Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +struct ThreadOccurrenceSearchCursor { + snapshot_id: u64, + offset: usize, +} + +pub(super) async fn search_thread_occurrences( + store: &LocalThreadStore, + params: SearchThreadOccurrencesParams, +) -> ThreadStoreResult { + if params.search_term.is_empty() { + return Err(ThreadStoreError::InvalidRequest { + message: "search_thread_occurrences requires search_term".to_string(), + }); + } + if params.max_results == 0 { + return Err(ThreadStoreError::InvalidRequest { + message: "search_thread_occurrences requires max_results greater than zero".to_string(), + }); + } + let page_size = params.page_size.clamp(1, params.max_results); + + if let Some(cursor) = params.cursor.as_deref() { + return page_from_cached_snapshot(store, ¶ms, cursor, page_size).await; + } + + let snapshot = build_snapshot(store, ¶ms).await?; + let is_capped = snapshot.is_capped; + let page_end = page_size.min(snapshot.items.len()); + let items = snapshot.items[..page_end].to_vec(); + let next_cursor = if page_end < snapshot.items.len() { + let snapshot_id = store + .occurrence_search_cache + .lock() + .await + .insert_snapshot(snapshot)?; + Some(serialize_cursor(ThreadOccurrenceSearchCursor { + snapshot_id, + offset: page_end, + })?) + } else { + None + }; + + Ok(ThreadOccurrenceSearchPage { + items, + next_cursor, + is_capped, + }) +} + +async fn page_from_cached_snapshot( + store: &LocalThreadStore, + params: &SearchThreadOccurrencesParams, + cursor: &str, + page_size: usize, +) -> ThreadStoreResult { + let cursor = parse_cursor(cursor)?; + let mut cache = store.occurrence_search_cache.lock().await; + cache.remove_expired(); + let snapshot = cache.snapshots.get(&cursor.snapshot_id).ok_or_else(|| { + ThreadStoreError::InvalidRequest { + message: "invalid cursor: search snapshot is no longer available".to_string(), + } + })?; + if snapshot.thread_id != params.thread_id || snapshot.search_term != params.search_term { + return Err(ThreadStoreError::InvalidRequest { + message: "invalid cursor: thread or search term does not match snapshot".to_string(), + }); + } + if cursor.offset >= snapshot.items.len() { + return Err(ThreadStoreError::InvalidRequest { + message: "invalid cursor: search offset is out of range".to_string(), + }); + } + + let page_end = cursor + .offset + .saturating_add(page_size) + .min(snapshot.items.len()); + let items = snapshot.items[cursor.offset..page_end].to_vec(); + let has_next_page = page_end < snapshot.items.len(); + let next_cursor = has_next_page + .then(|| { + serialize_cursor(ThreadOccurrenceSearchCursor { + snapshot_id: cursor.snapshot_id, + offset: page_end, + }) + }) + .transpose()?; + + let page = ThreadOccurrenceSearchPage { + items, + next_cursor, + is_capped: snapshot.is_capped, + }; + if !has_next_page && let Some(snapshot) = cache.snapshots.get_mut(&cursor.snapshot_id) { + snapshot.completed = true; + snapshot.expires_at = Instant::now() + COMPLETED_SNAPSHOT_TTL; + } + Ok(page) +} + +async fn build_snapshot( + store: &LocalThreadStore, + params: &SearchThreadOccurrencesParams, +) -> ThreadStoreResult { + let projection = load_search_projection(store, params.thread_id).await?; + let matcher = LiteralMatcher::new(params.search_term.as_str()); + let mut items = Vec::with_capacity(params.max_results.saturating_add(1)); + for item in &projection.items { + let remaining = params + .max_results + .saturating_add(1) + .saturating_sub(items.len()); + items.extend(occurrences_in_item( + params.thread_id, + item.turn_id.clone(), + item.item_id.clone(), + item.text.as_str(), + item.turn_started_at, + &matcher, + remaining, + )); + if items.len() > params.max_results { + break; + } + } + let is_capped = items.len() > params.max_results; + items.truncate(params.max_results); + + Ok(ThreadOccurrenceSearchSnapshot { + thread_id: params.thread_id, + search_term: params.search_term.clone(), + items, + is_capped, + expires_at: Instant::now() + SNAPSHOT_TTL, + completed: false, + }) +} + +async fn load_search_projection( + store: &LocalThreadStore, + thread_id: codex_protocol::ThreadId, +) -> ThreadStoreResult> { + let path = + super::read_thread::resolve_rollout_path(store, thread_id, /*include_archived*/ true) + .await? + .ok_or(ThreadStoreError::ThreadNotFound { thread_id })?; + let mut lines = match codex_rollout::open_rollout_line_reader(path.as_path()).await { + Ok(lines) => lines, + Err(err) if err.kind() == std::io::ErrorKind::NotFound => { + return Err(ThreadStoreError::ThreadNotFound { thread_id }); + } + Err(err) => { + return Err(ThreadStoreError::Internal { + message: format!("failed to open rollout for search: {err}"), + }); + } + }; + let physical_path = lines.source_path().to_path_buf(); + let fingerprint = RolloutFingerprint { + path: physical_path.clone(), + len: lines.source_len(), + modified_at: lines.source_modified_at(), + }; + { + let mut cache = store.occurrence_search_cache.lock().await; + cache.remove_expired(); + if let Some(cached) = cache.projections.get_mut(&thread_id) + && cached.projection.fingerprint == fingerprint + { + cached.expires_at = Instant::now() + SNAPSHOT_TTL; + return Ok(Arc::clone(&cached.projection)); + } + } + + let plain_path = codex_rollout::plain_rollout_path(physical_path.as_path()); + let snapshot_len = if physical_path == plain_path { + Some(fingerprint.len) + } else { + // Compressed rollouts are immutable, so the decompressed byte count does not need a bound. + None + }; + let mut builder = ThreadHistoryBuilder::new(); + let mut turn_timestamps = HashMap::::new(); + let mut bytes_read = 0_u64; + let mut parse_errors = 0_usize; + + while let Some(line) = lines + .next_line() + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to read rollout during search: {err}"), + })? + { + if snapshot_len.is_some_and(|limit| bytes_read.saturating_add(line.len() as u64) > limit) { + break; + } + bytes_read = bytes_read.saturating_add(line.len() as u64 + 1); + let rollout_line = match serde_json::from_str::(line.trim()) { + Ok(rollout_line) => rollout_line, + Err(_) => { + parse_errors = parse_errors.saturating_add(1); + continue; + } + }; + let changes = builder.handle_rollout_item_with_changes(&rollout_line.item); + let line_timestamp = if changes + .changed_turns + .iter() + .any(|turn| turn.started_at.is_none()) + { + DateTime::parse_from_rfc3339(&rollout_line.timestamp) + .map(|timestamp| timestamp.timestamp()) + .unwrap_or_default() + } else { + 0 + }; + + for removed_turn_id in changes.removed_turn_ids { + turn_timestamps.remove(&removed_turn_id); + } + for turn in changes.changed_turns { + if let Some(started_at) = turn.started_at { + turn_timestamps.insert(turn.turn_id, started_at); + } else { + turn_timestamps + .entry(turn.turn_id) + .or_insert(line_timestamp); + } + } + } + if parse_errors > 0 { + tracing::warn!(%thread_id, parse_errors, "skipped malformed rollout records during thread search"); + } + + let mut items = Vec::new(); + for turn in builder.finish() { + let turn_started_at = turn + .started_at + .or_else(|| turn_timestamps.get(&turn.id).copied()) + .unwrap_or_default(); + let last_assistant_id = turn.items.iter().rev().find_map(|item| match item { + ThreadItem::AgentMessage { id, .. } => Some(id.clone()), + _ => None, + }); + for item in turn.items { + if matches!(&item, ThreadItem::AgentMessage { id, .. } if Some(id) != last_assistant_id.as_ref()) + { + continue; + } + let Some(text) = searchable_text(&item) else { + continue; + }; + items.push(SearchableThreadItem { + turn_id: turn.id.clone(), + item_id: item.id().to_string(), + text: text.into_owned(), + turn_started_at, + }); + } + } + let projection = Arc::new(SearchProjection { fingerprint, items }); + store + .occurrence_search_cache + .lock() + .await + .insert_projection( + thread_id, + CachedSearchProjection { + projection: Arc::clone(&projection), + expires_at: Instant::now() + SNAPSHOT_TTL, + }, + ); + Ok(projection) +} + +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::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.join("\n"))) + } + } + } + ThreadItem::AgentMessage { text, .. } => { + let text = markdown_to_search_text(text); + (!text.is_empty()).then_some(Cow::Owned(text)) + } + _ => None, + } +} + +fn strip_user_message_prefix(text: &str) -> &str { + match text.find(USER_MESSAGE_BEGIN) { + Some(index) => text[index + USER_MESSAGE_BEGIN.len()..].trim(), + None => text.trim(), + } +} + +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::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::HardBreak + | 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 occurrences_in_item( + thread_id: codex_protocol::ThreadId, + turn_id: String, + item_id: String, + text: &str, + turn_started_at: i64, + matcher: &LiteralMatcher, + limit: usize, +) -> Vec { + let mut previous_match_end = 0; + let mut utf16_offset = 0_u32; + matcher + .find_ranges(text, limit) + .into_iter() + .enumerate() + .map(|(occurrence_index, matched)| { + utf16_offset = + utf16_offset.saturating_add(utf16_len(&text[previous_match_end..matched.start])); + let match_start = utf16_offset; + let match_end = match_start.saturating_add(utf16_len(&text[matched.clone()])); + previous_match_end = matched.end; + utf16_offset = match_end; + 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]); + let occurrence_index = u32::try_from(occurrence_index).unwrap_or(u32::MAX); + + StoredThreadOccurrence { + thread_id, + occurrence_id: format!("{turn_id}:{item_id}:{match_start}"), + turn_id: turn_id.clone(), + item_id: item_id.clone(), + occurrence_index, + snippet, + match_range: SearchTextRange { + start: match_start, + end: match_end, + }, + snippet_match_range: SearchTextRange { + start: snippet_match_start, + end: snippet_match_start.saturating_add(match_len), + }, + turn_started_at, + } + }) + .collect() +} + +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()) +} + +fn serialize_cursor(cursor: ThreadOccurrenceSearchCursor) -> ThreadStoreResult { + serde_json::to_string(&cursor).map_err(|err| ThreadStoreError::Internal { + message: format!("failed to serialize search cursor: {err}"), + }) +} + +fn parse_cursor(cursor: &str) -> ThreadStoreResult { + serde_json::from_str(cursor).map_err(|_| ThreadStoreError::InvalidRequest { + message: format!("invalid cursor: {cursor}"), + }) +} + +impl ThreadOccurrenceSearchCache { + fn insert_snapshot( + &mut self, + snapshot: ThreadOccurrenceSearchSnapshot, + ) -> ThreadStoreResult { + self.remove_expired(); + if self.snapshots.len() >= MAX_ACTIVE_SNAPSHOTS + && let Some(completed_snapshot_id) = self + .snapshots + .iter() + .filter(|(_, snapshot)| snapshot.completed) + .min_by_key(|(_, snapshot)| snapshot.expires_at) + .map(|(snapshot_id, _)| *snapshot_id) + { + self.snapshots.remove(&completed_snapshot_id); + } + if self.snapshots.len() >= MAX_ACTIVE_SNAPSHOTS { + return Err(ThreadStoreError::Internal { + message: "too many active thread search snapshots".to_string(), + }); + } + self.next_snapshot_id = self.next_snapshot_id.wrapping_add(1); + let snapshot_id = self.next_snapshot_id; + self.snapshots.insert(snapshot_id, snapshot); + Ok(snapshot_id) + } + + fn insert_projection( + &mut self, + thread_id: codex_protocol::ThreadId, + projection: CachedSearchProjection, + ) { + self.remove_expired(); + if self.projections.len() >= MAX_CACHED_PROJECTIONS + && !self.projections.contains_key(&thread_id) + && let Some(oldest_thread_id) = self + .projections + .iter() + .min_by_key(|(_, projection)| projection.expires_at) + .map(|(thread_id, _)| *thread_id) + { + self.projections.remove(&oldest_thread_id); + } + self.projections.insert(thread_id, projection); + } + + fn remove_expired(&mut self) { + let now = Instant::now(); + self.snapshots + .retain(|_, snapshot| snapshot.expires_at > now); + self.projections + .retain(|_, projection| projection.expires_at > now); + } +} diff --git a/codex-rs/thread-store/src/local/search_thread_occurrences_tests.rs b/codex-rs/thread-store/src/local/search_thread_occurrences_tests.rs new file mode 100644 index 0000000000..b24e8d5c94 --- /dev/null +++ b/codex-rs/thread-store/src/local/search_thread_occurrences_tests.rs @@ -0,0 +1,219 @@ +use std::sync::Arc; + +use codex_protocol::ThreadId; +use codex_protocol::models::BaseInstructions; +use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::protocol::AgentMessageEvent; +use codex_protocol::protocol::AgentReasoningEvent; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadHistoryMode; +use codex_protocol::protocol::ThreadMemoryMode; +use codex_protocol::protocol::TurnCompleteEvent; +use codex_protocol::protocol::TurnStartedEvent; +use codex_protocol::protocol::USER_MESSAGE_BEGIN; +use codex_protocol::protocol::UserMessageEvent; +use pretty_assertions::assert_eq; +use tempfile::TempDir; + +use super::LocalThreadStore; +use crate::CreateThreadParams; +use crate::LiveThread; +use crate::SearchThreadOccurrencesParams; +use crate::ThreadPersistenceMetadata; +use crate::ThreadStore; +use crate::local::test_support::test_config; + +#[tokio::test] +async fn searches_visible_occurrences_and_keeps_paginated_snapshot_stable() { + let home = TempDir::new().expect("temp dir"); + let store = Arc::new(LocalThreadStore::new(test_config(home.path()), None)); + let thread_id = ThreadId::default(); + let live_thread = LiveThread::create(store.clone(), create_thread_params(thread_id)) + .await + .expect("create thread"); + + live_thread + .append_items(&[ + turn_started("turn-1", 1_700_000_000), + user_message(&format!( + "needle in hidden environment context\n{USER_MESSAGE_BEGIN}\nNeedle twice: NEEDLE" + )), + RolloutItem::EventMsg(EventMsg::AgentReasoning(AgentReasoningEvent { + text: "needle in hidden reasoning".to_string(), + })), + RolloutItem::ResponseItem(codex_protocol::models::ResponseItem::FunctionCallOutput { + id: None, + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload::from_text("needle in tool output".to_string()), + internal_chat_message_metadata_passthrough: None, + }), + agent_message("needle in an earlier assistant message"), + agent_message("**Needle** from assistant; [needle](https://example.com)"), + turn_complete("turn-1"), + ]) + .await + .expect("append searchable history"); + live_thread.flush().await.expect("flush history"); + + let first = store + .search_thread_occurrences(search_params(thread_id, None, 2)) + .await + .expect("first search page"); + assert_eq!(first.items.len(), 2); + assert_eq!(first.items[0].turn_id, "turn-1"); + assert_eq!(first.items[0].item_id, "item-1"); + assert_eq!(first.items[0].occurrence_index, 0); + assert_eq!(first.items[1].occurrence_index, 1); + assert_eq!(first.items[0].turn_started_at, 1_700_000_000); + assert_eq!(first.items[0].snippet_match_range.start, 0); + assert_eq!(first.items[0].snippet_match_range.end, 6); + let cursor = first.next_cursor.expect("second page cursor"); + + live_thread + .append_items(&[ + turn_started("turn-2", 1_700_000_100), + agent_message("new needle after the snapshot"), + turn_complete("turn-2"), + ]) + .await + .expect("append after snapshot"); + live_thread.flush().await.expect("flush appended history"); + + let second = store + .search_thread_occurrences(search_params(thread_id, Some(cursor), 2)) + .await + .expect("second search page"); + assert_eq!(second.items.len(), 2); + assert_eq!(second.items[0].turn_id, "turn-1"); + assert_eq!(second.items[0].item_id, "item-4"); + assert_eq!(second.items[1].item_id, "item-4"); + assert_eq!(second.next_cursor, None); + + let refreshed = store + .search_thread_occurrences(search_params(thread_id, None, 10)) + .await + .expect("refreshed search"); + assert_eq!(refreshed.items.len(), 5); + assert_eq!(refreshed.items.last().expect("new match").turn_id, "turn-2"); + + let capped = store + .search_thread_occurrences(SearchThreadOccurrencesParams { + max_results: 3, + page_size: 3, + ..search_params(thread_id, None, 3) + }) + .await + .expect("capped search"); + assert_eq!(capped.items.len(), 3); + assert!(capped.is_capped); +} + +#[test] +fn reports_utf16_ranges_for_javascript_clients() { + let matcher = super::LiteralMatcher::new("needle"); + let matches = super::occurrences_in_item( + ThreadId::default(), + "turn-1".to_string(), + "item-1".to_string(), + "🙂NEEDLE needle", + 1, + &matcher, + usize::MAX, + ); + + assert_eq!( + matches + .iter() + .map(|item| item.match_range.clone()) + .collect::>(), + vec![ + crate::SearchTextRange { start: 2, end: 8 }, + crate::SearchTextRange { start: 9, end: 15 }, + ] + ); + + let greek_matches = super::LiteralMatcher::new("ος").find_ranges("ΟΣ", usize::MAX); + assert_eq!(greek_matches, vec![0..4]); + let expanded_matches = super::LiteralMatcher::new("i").find_ranges("İ", usize::MAX); + assert_eq!(expanded_matches, vec![0..2]); + assert_eq!( + super::markdown_to_search_text("before needle"), + "before needle" + ); + assert_eq!(super::markdown_to_search_text("foo \nbar"), "foobar"); +} + +fn search_params( + thread_id: ThreadId, + cursor: Option, + page_size: usize, +) -> SearchThreadOccurrencesParams { + SearchThreadOccurrencesParams { + thread_id, + search_term: "needle".to_string(), + cursor, + page_size, + max_results: 250, + } +} + +fn create_thread_params(thread_id: ThreadId) -> CreateThreadParams { + CreateThreadParams { + session_id: thread_id.into(), + thread_id, + extra_config: None, + forked_from_id: None, + parent_thread_id: None, + source: SessionSource::Exec, + thread_source: None, + originator: "test_originator".to_string(), + base_instructions: BaseInstructions::default(), + dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), + multi_agent_version: None, + history_mode: ThreadHistoryMode::Legacy, + initial_window_id: uuid::Uuid::now_v7().to_string(), + metadata: ThreadPersistenceMetadata { + cwd: Some(std::env::current_dir().expect("cwd")), + model_provider: "test-provider".to_string(), + memory_mode: ThreadMemoryMode::Enabled, + }, + } +} + +fn turn_started(turn_id: &str, started_at: i64) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::TurnStarted(TurnStartedEvent { + turn_id: turn_id.to_string(), + trace_id: None, + started_at: Some(started_at), + model_context_window: None, + collaboration_mode_kind: Default::default(), + })) +} + +fn turn_complete(turn_id: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: turn_id.to_string(), + last_agent_message: None, + completed_at: None, + duration_ms: None, + time_to_first_token_ms: None, + })) +} + +fn user_message(message: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { + message: message.to_string(), + ..Default::default() + })) +} + +fn agent_message(message: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::AgentMessage(AgentMessageEvent { + message: message.to_string(), + phase: None, + memory_citation: None, + })) +} diff --git a/codex-rs/thread-store/src/store.rs b/codex-rs/thread-store/src/store.rs index 67ee603f49..be78a645a5 100644 --- a/codex-rs/thread-store/src/store.rs +++ b/codex-rs/thread-store/src/store.rs @@ -16,9 +16,11 @@ use crate::LoadThreadHistoryParams; use crate::ReadThreadByRolloutPathParams; use crate::ReadThreadParams; use crate::ResumeThreadParams; +use crate::SearchThreadOccurrencesParams; use crate::SearchThreadsParams; use crate::StoredThread; use crate::StoredThreadHistory; +use crate::ThreadOccurrenceSearchPage; use crate::ThreadPage; use crate::ThreadSearchPage; use crate::ThreadStoreError; @@ -102,6 +104,18 @@ pub trait ThreadStore: Any + Send + Sync { }) } + /// Searches visible user and assistant message occurrences within one thread. + fn search_thread_occurrences( + &self, + _params: SearchThreadOccurrencesParams, + ) -> ThreadStoreFuture<'_, ThreadOccurrenceSearchPage> { + Box::pin(async { + Err(ThreadStoreError::Unsupported { + operation: "search_thread_occurrences", + }) + }) + } + /// Lists turns within a stored thread. fn list_turns(&self, _params: ListTurnsParams) -> ThreadStoreFuture<'_, TurnPage> { Box::pin(async { diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index 4e688eca6f..1bc1c1d99e 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -284,6 +284,50 @@ pub struct ThreadSearchPage { pub next_cursor: Option, } +/// Parameters for searching occurrences within one persisted thread. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct SearchThreadOccurrencesParams { + /// Thread id to search. + pub thread_id: ThreadId, + /// Case-insensitive literal substring to find in visible message text. + pub search_term: String, + /// Opaque cursor returned by a previous search call for this query. + pub cursor: Option, + /// Maximum number of occurrences to return in this page. + pub page_size: usize, + /// Maximum number of occurrences retained for this query generation. + pub max_results: usize, +} + +/// UTF-16 code-unit range suitable for direct use by JavaScript clients. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct SearchTextRange { + pub start: u32, + pub end: u32, +} + +/// Lightweight stable anchor for one persisted message occurrence. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct StoredThreadOccurrence { + pub thread_id: ThreadId, + pub turn_id: String, + pub item_id: String, + pub occurrence_id: String, + pub occurrence_index: u32, + pub snippet: String, + pub match_range: SearchTextRange, + pub snippet_match_range: SearchTextRange, + pub turn_started_at: i64, +} + +/// Point-in-time page of occurrence search results. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct ThreadOccurrenceSearchPage { + pub items: Vec, + pub next_cursor: Option, + pub is_capped: bool, +} + /// Requested amount of item detail for stored turns. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] pub enum StoredTurnItemsView {