feat(app-server): add within-thread occurrence search

This commit is contained in:
Brent Traut
2026-07-08 10:43:14 -07:00
parent f1affbac5e
commit 786430670b
24 changed files with 1551 additions and 5 deletions
+2
View File
@@ -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",
@@ -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": {
@@ -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": {
@@ -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, };
@@ -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, };
@@ -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";
@@ -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,
@@ -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(&params).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::<ThreadListParams>(json!({
@@ -1223,6 +1223,65 @@ pub struct ThreadSearchResponse {
pub backwards_cursor: Option<String>,
}
#[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<String>,
/// Optional page size; the server caps a query generation at 250 occurrences.
#[ts(optional = nullable)]
pub limit: Option<u32>,
}
#[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<ThreadSearchOccurrence>,
/// Opaque cursor for the next page of this point-in-time query snapshot.
pub next_cursor: Option<String>,
/// 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/")]
+33
View File
@@ -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
@@ -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
}
@@ -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;
@@ -707,6 +707,15 @@ impl ThreadRequestProcessor {
.map(|response| Some(response.into()))
}
pub(crate) async fn thread_search_occurrences(
&self,
params: ThreadSearchOccurrencesParams,
) -> Result<Option<ClientResponsePayload>, 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<ThreadSearchOccurrencesResponse, JSONRPCErrorError> {
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::<Result<Vec<_>, 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,
@@ -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<i64> {
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,
@@ -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::<ThreadSearchOccurrencesResponse>(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::<ThreadSearchOccurrencesResponse>(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::<ThreadTurnsListResponse>(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::<ThreadSearchOccurrencesResponse>(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;
+32 -4
View File
@@ -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<SystemTime>,
}
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<SystemTime> {
self.source_modified_at
}
/// Reads the next JSONL record from the rollout.
pub async fn next_line(&mut self) -> io::Result<Option<String>> {
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<dyn Read + Send>).lines(),
)
let lines = io::BufReader::new(Box::new(decoder) as Box<dyn Read + Send>).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(),
})
}
}
+2
View File
@@ -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 }
+4
View File
@@ -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;
+14
View File
@@ -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<Mutex<HashMap<ThreadId, LiveRecorderEntry>>>,
occurrence_search_cache: Arc<Mutex<search_thread_occurrences::ThreadOccurrenceSearchCache>>,
state_db: Option<StateDbHandle>,
}
@@ -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,
@@ -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,
@@ -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<u64, ThreadOccurrenceSearchSnapshot>,
projections: HashMap<codex_protocol::ThreadId, CachedSearchProjection>,
}
#[derive(Clone)]
struct ThreadOccurrenceSearchSnapshot {
thread_id: codex_protocol::ThreadId,
search_term: String,
items: Vec<StoredThreadOccurrence>,
is_capped: bool,
expires_at: Instant,
completed: bool,
}
#[derive(Clone, PartialEq, Eq)]
struct RolloutFingerprint {
path: PathBuf,
len: u64,
modified_at: Option<SystemTime>,
}
#[derive(Clone)]
struct SearchProjection {
fingerprint: RolloutFingerprint,
items: Vec<SearchableThreadItem>,
}
struct CachedSearchProjection {
projection: Arc<SearchProjection>,
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<ThreadOccurrenceSearchPage> {
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, &params, cursor, page_size).await;
}
let snapshot = build_snapshot(store, &params).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<ThreadOccurrenceSearchPage> {
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<ThreadOccurrenceSearchSnapshot> {
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<Arc<SearchProjection>> {
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::<String, i64>::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::<RolloutLine>(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<Cow<'_, str>> {
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::<Vec<_>>().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<std::ops::Range<usize>> {
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::<usize>();
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<StoredThreadOccurrence> {
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<String> {
serde_json::to_string(&cursor).map_err(|err| ThreadStoreError::Internal {
message: format!("failed to serialize search cursor: {err}"),
})
}
fn parse_cursor(cursor: &str) -> ThreadStoreResult<ThreadOccurrenceSearchCursor> {
serde_json::from_str(cursor).map_err(|_| ThreadStoreError::InvalidRequest {
message: format!("invalid cursor: {cursor}"),
})
}
impl ThreadOccurrenceSearchCache {
fn insert_snapshot(
&mut self,
snapshot: ThreadOccurrenceSearchSnapshot,
) -> ThreadStoreResult<u64> {
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);
}
}
@@ -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<_>>(),
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 <kbd>needle</kbd>"),
"before <kbd>needle</kbd>"
);
assert_eq!(super::markdown_to_search_text("foo \nbar"), "foobar");
}
fn search_params(
thread_id: ThreadId,
cursor: Option<String>,
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,
}))
}
+14
View File
@@ -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 {
+44
View File
@@ -284,6 +284,50 @@ pub struct ThreadSearchPage {
pub next_cursor: Option<String>,
}
/// 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<String>,
/// 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<StoredThreadOccurrence>,
pub next_cursor: Option<String>,
pub is_capped: bool,
}
/// Requested amount of item detail for stored turns.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub enum StoredTurnItemsView {