Skip to content
Source
Sign inSign up

ThreadFind.ts

471 lines18.5 KB

1import {
2 OrchestrationV2TurnItemJson,
3 type OrchestrationV2ProjectedTurnItem,
4 OrchestrationV2SearchThreadInput,
5 type OrchestrationV2SearchThreadResult,
6 type OrchestrationV2TurnItem,
7 ThreadId,
8 TurnItemId,
9} from "@t3tools/contracts";
10import { searchableMessageSegments, searchablePlanSegments } from "@t3tools/shared/threadFindText";
11import { countThreadSearchOccurrences } from "@t3tools/shared/threadSearch";
12import { htmlRenderFromToolItem, mcpAppFromToolItem } from "@t3tools/shared/toolOutput";
13import * as Cache from "effect/Cache";
14import * as Data from "effect/Data";
15import * as Effect from "effect/Effect";
16import * as RcRef from "effect/RcRef";
17import * as Schema from "effect/Schema";
18import * as Stream from "effect/Stream";
19import * as NodeSqliteClient from "@t3tools/shared/nodeSqliteClient";
20import * as SqlClient from "effect/sql/SqlClient";
21import {
22 ProjectionStoreReadError,
23 ProjectionStoreThreadNotFoundError,
24 type ProjectionStoreV2Error,
25} from "./ProjectionStore.ts";
26​
27interface FindRow {
28 readonly position: number;
29 readonly visibility: OrchestrationV2ProjectedTurnItem["visibility"];
30 readonly sourceThreadId: ThreadId;
31 readonly sourceItemId: TurnItemId;
33interface FindPayload extends FindRow {
34 readonly payload: string;
35 readonly entryId: string;
36 readonly runId: OrchestrationV2TurnItem["runId"];
37 readonly type: OrchestrationV2TurnItem["type"];
38 readonly emptyText: boolean;
40interface FindDocument extends FindRow {
41 readonly entryId: string;
42 readonly runId: OrchestrationV2TurnItem["runId"];
43 readonly count: number;
45type CompletedScans = Map<
46 string,
47 {
48 sequence: number;
49 documents: readonly FindDocument[];
50 sourceThreadIds: readonly ThreadId[];
51 }
52>;
53class TextKey extends Data.Class<{
54 text: string;
55 role: "user" | "assistant" | "plan";
56 streaming: boolean;
57 hasContext: boolean;
58 cwd: string | undefined;
59 skills: NonNullable<OrchestrationV2SearchThreadInput["skills"]>;
60}> {}
61​
62function textKey(
63 item: OrchestrationV2TurnItem,
64 cwd: string | undefined,
65 skills: TextKey["skills"],
66) {
67 return new TextKey({
68 text:
69 item.type === "proposed_plan"
70 ? item.markdown
71 : item.type === "user_message" || item.type === "assistant_message"
72 ? item.text
73 : "",
74 role:
75 item.type === "proposed_plan" ? "plan" : item.type === "user_message" ? "user" : "assistant",
76 streaming: item.type === "assistant_message" && item.streaming,
77 hasContext: item.type === "user_message" && item.context !== undefined,
78 cwd,
79 skills,
80 });
82​
83function parseText(key: TextKey): readonly string[] {
84 return key.role === "plan"
85 ? searchablePlanSegments(key.text, key.cwd)
86 : (searchableMessageSegments(
87 {
88 role: key.role,
89 text: key.text,
90 streaming: key.streaming,
91 context: key.hasContext ? { version: 1, records: [] } : undefined,
92 },
93 key.cwd,
94 key.skills,
95 ) ?? []);
97​
98function documentFor(row: OrchestrationV2ProjectedTurnItem, count: number): FindDocument {
99 return {
100 position: row.position,
101 visibility: row.visibility,
102 sourceThreadId: row.sourceThreadId,
103 sourceItemId: row.sourceItemId,
104 entryId:
105 row.item.type === "user_message" || row.item.type === "assistant_message"
106 ? row.item.messageId
107 : row.item.id,
108 runId: row.item.runId,
109 count,
110 };
112​
113function countSegments(segments: readonly string[], query: string) {
114 return segments.reduce((sum, text) => sum + countThreadSearchOccurrences(text, query), 0);
116​
117function selectMatch(
118 documents: readonly FindDocument[],
119 input: Pick<OrchestrationV2SearchThreadInput, "index" | "start" | "offset">,
120) {
121 const totalMatches = documents.reduce((sum, doc) => sum + doc.count, 0);
122 let requestedIndex = input.index ?? 0;
123 if (input.index === undefined && input.start) {
124 const startIndex = documents.findIndex((doc) => doc.entryId === input.start?.entryId);
125 if (startIndex >= 0) {
126 requestedIndex =
127 documents.slice(0, startIndex).reduce((sum, doc) => sum + doc.count, 0) +
128 Math.min(input.start.occurrence, documents[startIndex]!.count);
129 if (requestedIndex >= totalMatches) requestedIndex = 0;
130 }
131 }
132 const activeIndex =
133 input.index === undefined && totalMatches > 0
134 ? (((requestedIndex + (input.offset ?? 0)) % totalMatches) + totalMatches) % totalMatches
135 : Math.min(requestedIndex, Math.max(0, totalMatches - 1));
136 let occurrence = activeIndex;
137 for (const document of documents) {
138 if (occurrence < document.count) return { totalMatches, activeIndex, document, occurrence };
139 occurrence -= document.count;
140 }
141 return { totalMatches, activeIndex, document: null, occurrence: 0 };
143​
144function resultForSelection(
145 selected: ReturnType<typeof selectMatch>,
146 documents: readonly FindDocument[],
147 snapshotSequence: number,
148): OrchestrationV2SearchThreadResult {
149 let startIndex = 0;
150 const entries = documents.flatMap(({ entryId, runId, count }) => {
151 const entry = { entryId, runId, count, startIndex };
152 startIndex += count;
153 return count > 0 ? [entry] : [];
154 });
155 const selectedEntry = entries.findIndex((entry) => entry.entryId === selected.document?.entryId);
156 const navigation =
157 entries.length <= 17
158 ? entries
159 : Array.from(
160 { length: 17 },
161 (_, i) => entries[(selectedEntry + i - 8 + entries.length) % entries.length]!,
162 );
163 return {
164 snapshotSequence,
165 totalMatches: selected.totalMatches,
166 activeIndex: selected.activeIndex,
167 navigation,
168 match:
169 selected.document === null
170 ? null
171 : {
172 entryId: selected.document.entryId,
173 runId: selected.document.runId,
174 occurrence: selected.occurrence,
175 },
176 };
178​
179export const makeThreadFind = Effect.fn("makeThreadFind")(function* (
180 readIndex: (threadId: ThreadId) => Effect.Effect<
181 readonly (Omit<FindRow, "position"> & {
182 readonly item: Pick<OrchestrationV2TurnItem, "type" | "runId">;
183 })[],
184 ProjectionStoreV2Error
185 >,
186) {
187 const sql = yield* SqlClient.SqlClient;
188 const encodeCacheKey = Schema.encodeEffect(
189 Schema.fromJsonString(
190 Schema.Struct({
191 threadId: ThreadId,
192 cwd: Schema.optional(Schema.String),
193 query: Schema.String,
194 skills: OrchestrationV2SearchThreadInput.fields.skills,
195 }),
196 ),
197 );
198 const encodeThreadIds = Schema.encodeEffect(Schema.fromJsonString(Schema.Array(ThreadId)));
199 const encodeSources = Schema.encodeEffect(
200 Schema.fromJsonString(Schema.Array(Schema.Struct({ threadId: ThreadId, id: TurnItemId }))),
201 );
202 const decodeItem = Schema.decodeUnknownEffect(Schema.fromJsonString(OrchestrationV2TurnItemJson));
203 const isSnapshotError = Schema.is(
204 Schema.Union([ProjectionStoreReadError, ProjectionStoreThreadNotFoundError]),
205 );
206 const load = Effect.fn("ThreadFind.load")(function* (
207 threadId: ThreadId,
208 rows: readonly FindRow[],
209 ) {
210 if (rows.length === 0) return [];
211 const sources = yield* encodeSources(
212 rows.map((row) => ({ threadId: row.sourceThreadId, id: row.sourceItemId })),
213 ).pipe(Effect.mapError((cause) => new ProjectionStoreReadError({ threadId, cause })));
214 const payloads = yield* sql<{
215 thread_id: string;
216 turn_item_id: string;
217 payload_json: string;
218 entry_id: string;
219 run_id: OrchestrationV2TurnItem["runId"];
220 item_type: OrchestrationV2TurnItem["type"];
221 empty_text: number | null;
222 }>`
223 SELECT item.thread_id, item.turn_item_id, item.payload_json,
224 COALESCE(json_extract(item.payload_json, '$.messageId'), item.turn_item_id) AS entry_id,
225 json_extract(item.payload_json, '$.runId') AS run_id,
226 json_extract(item.payload_json, '$.type') AS item_type,
227 json_extract(item.payload_json, '$.text') = '' AS empty_text
228 FROM orchestration_v2_projection_turn_items AS item
229 JOIN json_each(${sources}) AS wanted
230 ON item.thread_id = json_extract(wanted.value, '$.threadId')
231 AND item.turn_item_id = json_extract(wanted.value, '$.id')
232 `.pipe(Effect.mapError((cause) => new ProjectionStoreReadError({ threadId, cause })));
233 const byThread = new Map<string, Map<string, (typeof payloads)[number]>>();
234 for (const row of payloads) {
235 let items = byThread.get(row.thread_id);
236 if (!items) byThread.set(row.thread_id, (items = new Map()));
237 items.set(row.turn_item_id, row);
238 }
239 return yield* Effect.forEach(rows, (row) => {
240 const payload = byThread.get(row.sourceThreadId)?.get(row.sourceItemId);
241 return payload === undefined
242 ? Effect.fail(new ProjectionStoreReadError({ threadId }))
243 : Effect.succeed({
244 ...row,
245 payload: payload.payload_json,
246 entryId: payload.entry_id,
247 runId: payload.run_id,
248 type: payload.item_type,
249 emptyText: payload.empty_text === 1,
250 });
251 });
252 });
253 const makeTextCache = (capacity: number) =>
254 Cache.make({ capacity, lookup: (key: TextKey) => Effect.sync(() => parseText(key)) });
255 // Text keys are message content, so entries never go stale. Cache TTL only expires on access
256 // and would keep the last searched threads forever. Each search borrows this for its whole
257 // stream, and all of it is dropped once Find is idle.
258 const findCaches = yield* RcRef.make({
259 acquire: Effect.all({
260 // Scans visit every message; a cache smaller than the thread churns without hits.
261 text: makeTextCache(8192),
262 // Bounded separately because each entry can hold a large body; typing reuses the parse.
263 largeText: makeTextCache(32),
264 scans: Effect.sync((): CompletedScans => new Map()),
265 }),
266 idleTimeToLive: "1 minute",
267 });
268 const readRevision = Effect.fn("ThreadFind.readRevision")(function* (
269 threadIds: readonly ThreadId[],
270 ) {
271 const sources = yield* encodeThreadIds(threadIds);
272 const revision = yield* sql<{ sequence: number }>`
273 WITH RECURSIVE source_threads(thread_id) AS (
274 SELECT value FROM json_each(${sources})
275 UNION
276 SELECT json_extract(t.payload_json, '$.forkedFrom.threadId')
277 FROM orchestration_v2_projection_threads t
278 JOIN source_threads s ON s.thread_id = t.thread_id
279 WHERE json_extract(t.payload_json, '$.forkedFrom.type') = 'run'
280 )
281 SELECT COALESCE(MAX(sequence), 0) AS sequence FROM orchestration_events
282 -- The planner otherwise walks every V2 event backwards from the newest,
283 -- which takes seconds for a thread that has been idle a long time.
284 INDEXED BY idx_orch_events_stream_sequence
285 WHERE application_event_version = 2 AND aggregate_kind = 'thread'
286 AND stream_id IN (SELECT thread_id FROM source_threads)
287 `;
288 return revision[0]?.sequence ?? 0;
289 });
290 const snapshot = Effect.fn("ThreadFind.snapshot")(
291 function* (input: OrchestrationV2SearchThreadInput, scans: CompletedScans) {
292 const { threadId } = input;
293 const active = yield* sql<{ cwd: string | null }>`
294 SELECT COALESCE(json_extract(t.payload_json, '$.worktreePath'), p.workspace_root) AS cwd
295 FROM orchestration_v2_projection_threads t JOIN projection_projects p ON p.project_id = t.project_id
296 WHERE t.thread_id = ${threadId} AND t.deleted_at IS NULL AND p.deleted_at IS NULL
297 `;
298 if (!active[0]) return yield* new ProjectionStoreThreadNotFoundError({ threadId });
299 const cwd = active[0].cwd ?? undefined;
300 const cacheKey = yield* encodeCacheKey({
301 threadId,
302 cwd,
303 query: input.query,
304 skills: input.skills ?? [],
305 });
306 const cached = scans.get(cacheKey);
307 if (cached && cached.sequence === (yield* readRevision(cached.sourceThreadIds))) {
308 const selected = selectMatch(cached.documents, input);
309 scans.delete(cacheKey);
310 scans.set(cacheKey, cached);
311 return { result: resultForSelection(selected, cached.documents, cached.sequence) };
312 }
313 const index = yield* readIndex(threadId);
314 const sourceThreadIds = [...new Set([threadId, ...index.map((row) => row.sourceThreadId)])];
315 const sequence = yield* readRevision(sourceThreadIds);
316 const candidates = index
317 .map((row, position) => ({ ...row, position }))
318 .filter((row) =>
319 ["user_message", "assistant_message", "proposed_plan", "user_input_request"].includes(
320 row.item.type,
321 ),
322 );
323 const loadAll = Effect.fnUntraced(function* (rows: readonly FindRow[]) {
324 const loaded: FindPayload[] = [];
325 for (let start = 0; start < rows.length; start += 128) {
326 loaded.push(...(yield* load(threadId, rows.slice(start, start + 128))));
327 }
328 return loaded;
329 });
330 const payloads = yield* loadAll(candidates);
331 // Only runs with an empty reply can hide its fallback, so other tool payloads stay unread.
332 const emptyReplyRunIds = new Set(
333 payloads.flatMap((row) =>
334 row.type === "assistant_message" && row.emptyText && row.runId !== null
335 ? [row.runId]
336 : [],
337 ),
338 );
339 const toolPayloads = yield* loadAll(
340 index
341 .map((row, position) => ({ ...row, position }))
342 .filter(
343 (row) =>
344 row.item.type === "dynamic_tool" &&
345 row.item.runId !== null &&
346 emptyReplyRunIds.has(row.item.runId),
347 ),
348 );
349 return { payloads, toolPayloads, cwd, sequence, cacheKey, sourceThreadIds };
350 },
351 (effect) => sql.withTransaction(effect).pipe(NodeSqliteClient.readOnly),
352 (effect, input) =>
353 effect.pipe(
354 Effect.mapError((cause) =>
355 isSnapshotError(cause)
356 ? cause
357 : new ProjectionStoreReadError({ threadId: input.threadId, cause }),
358 ),
359 ),
360 );
361 const searchThreadStream = (input: OrchestrationV2SearchThreadInput) =>
362 Stream.unwrap(
363 Effect.gen(function* () {
364 const { scans, ...textCache } = yield* RcRef.get(findCaches);
365 // Release SQLite before decoding and parsing. Both frames describe one snapshot.
366 const data = yield* snapshot(input, scans);
367 if (data.result) return Stream.succeed(data.result);
368 const { payloads, toolPayloads, cwd, sequence, cacheKey, sourceThreadIds } = data;
369 const decode = (row: FindPayload) =>
370 decodeItem(row.payload).pipe(
371 Effect.map((item) => ({ ...row, item })),
372 Effect.mapError(
373 (cause) => new ProjectionStoreReadError({ threadId: input.threadId, cause }),
374 ),
375 );
376 // Answers rendered inside question cards must not appear as separate matches.
377 const foldedAnswerRunIds = new Map<string, OrchestrationV2TurnItem["runId"]>();
378 for (const payload of payloads) {
379 if (payload.type !== "user_input_request") continue;
380 const { item } = yield* decode(payload);
381 if (item.type === "user_input_request" && item.questionAnswer)
382 foldedAnswerRunIds.set(`async-answer:${item.questionAnswer.requestId}`, item.runId);
383 }
384 // The timeline hides an empty reply's fallback when its run shows an HTML render or MCP app.
385 const embeddedResultRunIds = new Set<OrchestrationV2TurnItem["runId"]>();
386 for (const payload of toolPayloads) {
387 const { item } = yield* decode(payload);
388 if (
389 item.type === "dynamic_tool" &&
390 ((item.status === "completed" && htmlRenderFromToolItem(item)) ||
391 mcpAppFromToolItem(item))
392 )
393 embeddedResultRunIds.add(item.runId);
394 }
395 const rows = payloads.filter(
396 (row) =>
397 row.type !== "user_input_request" &&
398 !(row.type === "user_message" && foldedAnswerRunIds.get(row.entryId) === row.runId),
399 );
400 const documents = new Map<number, FindDocument>();
401 const count = Effect.fnUntraced(function* (row: FindPayload) {
402 const cached = documents.get(row.position);
403 if (cached) return cached;
404 const decoded = yield* decode(row);
405 if (
406 decoded.item.type === "assistant_message" &&
407 decoded.item.text === "" &&
408 embeddedResultRunIds.has(decoded.item.runId)
409 ) {
410 const document = documentFor(decoded, 0);
411 documents.set(row.position, document);
412 return document;
413 }
414 const key = textKey(decoded.item, cwd, input.skills ?? []);
415 const segments = yield* Cache.get(
416 key.text.length <= 32_768 ? textCache.text : textCache.largeText,
417 key,
418 );
419 const document = documentFor(decoded, countSegments(segments, input.query));
420 documents.set(row.position, document);
421 yield* Effect.yieldNow;
422 return document;
423 });
424 const finish = Effect.gen(function* () {
425 for (const row of rows) yield* count(row);
426 const ordered = rows.map((row) => documents.get(row.position)!);
427 scans.delete(cacheKey);
428 scans.set(cacheKey, { sequence, documents: ordered, sourceThreadIds });
429 if (scans.size > 8) {
430 const oldest = scans.keys().next().value;
431 if (oldest !== undefined) scans.delete(oldest);
432 }
433 return resultForSelection(selectMatch(ordered, input), ordered, sequence);
434 });
435 if (input.index !== undefined || (input.offset ?? 0) !== 0)
436 return Stream.fromEffect(finish);
437 const anchor = Math.max(
438 0,
439 rows.findIndex((row) => row.entryId === input.start?.entryId),
440 );
441 for (let step = 0; step <= rows.length; step++) {
442 const index = (anchor + step) % rows.length;
443 const row = rows[index];
444 if (!row) break;
445 const document = yield* count(row);
446 const occurrence =
447 step === 0 && row.entryId === input.start?.entryId
448 ? Math.min(input.start.occurrence, document.count)
449 : 0;
450 if (occurrence >= document.count) continue;
451 const first: OrchestrationV2SearchThreadResult = {
452 complete: false,
453 snapshotSequence: sequence,
454 totalMatches: 1,
455 activeIndex: 0,
456 match: { entryId: document.entryId, runId: document.runId, occurrence },
457 };
458 return Stream.concat(Stream.succeed(first), Stream.fromEffect(finish));
459 }
460 return Stream.fromEffect(finish);
461 }),
462 );
463 return {
464 searchThreadStream,
465 searchThread: (input: OrchestrationV2SearchThreadInput) =>
466 searchThreadStream(input).pipe(
467 Stream.runCollect,
468 Effect.map((results) => results[results.length - 1]!),
469 ),
470 };
471});