1import { OrchestrationV2ThreadProjectionJson } from "@t3tools/contracts";
2import {
3 omitLocalVisibleTurnItems,
4 restoreLocalVisibleTurnItems,
5} from "@t3tools/shared/orchestrationV2BoundedSnapshot";
6import * as Effect from "effect/Effect";
7import * as Option from "effect/Option";
8import * as Schema from "effect/Schema";
9
10import { yieldToHostTimer } from "./hostYield.ts";
11import {
12 ORCHESTRATION_THREAD_CACHE_SCHEMA_VERSION,
13 StoredOrchestrationThreadSnapshot,
14} from "./orchestrationCache.ts";
15
16// Keep the full format for web/desktop; older clients discard compact records.
17export const COMPACT_THREAD_CACHE_SCHEMA_VERSION = 6 as const;
18
19/**
20 * A thread followed live keeps every item it streams, so its projection can grow
21 * to hundreds of MB. Larger ones are not cached; reopening loads the bounded window.
22 */
23const THREAD_CACHE_MAX_JSON_LENGTH = 4 * 1024 * 1024;
24
25const ROWS_PER_CHUNK = 32;
26
27const StoredCompactThreadSnapshot = StoredOrchestrationThreadSnapshot.mapFields((fields) => ({
28 ...fields,
29 schemaVersion: Schema.Literal(COMPACT_THREAD_CACHE_SCHEMA_VERSION),
30 snapshot: fields.snapshot.mapFields((snapshotFields) => ({
31 ...snapshotFields,
32 // Required metadata prevents partial compact records from masquerading as full ones.
33 turnItemsOmitLocalVisible: Schema.Boolean,
34 turnItemCount: Schema.Int,
35 visibleTurnItemCount: Schema.Int,
36 })),
37}));
38
39class InvalidThreadCache extends Schema.TaggedError<InvalidThreadCache>()(
40 "InvalidThreadCache",
41 {},
42) {}
43
44const decodeCompatible = Schema.decodeUnknownEffect(
45 Schema.fromJsonString(
46 Schema.Union([StoredCompactThreadSnapshot, StoredOrchestrationThreadSnapshot]),
47 ),
48);
49
50const EncodedRows = Schema.Array(Schema.Unknown);
51const ProjectionWithEncodedRows = OrchestrationV2ThreadProjectionJson.mapFields((fields) => ({
52 ...fields,
53 turnItems: EncodedRows,
54 visibleTurnItems: EncodedRows,
55}));
56const encodeFullEnvelope = Schema.encodeEffect(
57 Schema.fromJsonString(
58 StoredOrchestrationThreadSnapshot.mapFields((fields) => ({
59 ...fields,
60 snapshot: fields.snapshot.mapFields((snapshotFields) => ({
61 ...snapshotFields,
62 projection: ProjectionWithEncodedRows,
63 })),
64 })),
65 ),
66);
67const encodeCompactEnvelope = Schema.encodeEffect(
68 Schema.fromJsonString(
69 StoredCompactThreadSnapshot.mapFields((fields) => ({
70 ...fields,
71 snapshot: fields.snapshot.mapFields((snapshotFields) => ({
72 ...snapshotFields,
73 projection: ProjectionWithEncodedRows,
74 })),
75 })),
76 ),
77);
78const encodeTurnItemChunk = Schema.encodeEffect(
79 OrchestrationV2ThreadProjectionJson.fields.turnItems,
80);
81const encodeVisibleRowChunk = Schema.encodeEffect(
82 OrchestrationV2ThreadProjectionJson.fields.visibleTurnItems,
83);
84
85type StoredThread = Pick<
86 typeof StoredOrchestrationThreadSnapshot.Type,
87 "environmentId" | "threadId" | "snapshot"
88>;
89type Projection = StoredThread["snapshot"]["projection"];
90
91interface EncodedRow {
92 readonly encoded: unknown;
93 readonly length: number;
94}
95
96/**
97 * Thread saves yield while encoding, so they can finish after a later save, a
98 * removal, or an environment clear. Each takes a ticket when it starts; a save
99 * writes only while no newer ticket exists for its thread or environment.
100 */
101export function makeThreadCacheWriteOrder() {
102 let latest = 0;
103 const tickets = new Map<string, number>();
104 const threadKey = (environmentId: string, threadId: string) =>
105 `${environmentId}\u0000${threadId}`;
106 return {
107 beginSave: (environmentId: string, threadId: string) => {
108 const ticket = ++latest;
109 const key = threadKey(environmentId, threadId);
110 tickets.set(key, ticket);
111 return () => tickets.get(key) === ticket && (tickets.get(environmentId) ?? 0) < ticket;
112 },
113 removeThread: (environmentId: string, threadId: string) => {
114 tickets.set(threadKey(environmentId, threadId), ++latest);
115 },
116 clearEnvironment: (environmentId: string) => {
117 tickets.set(environmentId, ++latest);
118 },
119 };
120}
121
122/**
123 * Turn items are immutable and the projection reducer keeps unchanged ones by
124 * reference, so each encoder remembers the encoding of rows it has encoded and
125 * runs the row codec only for new ones, yielding to the host between chunks.
126 * Encodes to none once the rows exceed `maxJsonLength`, before the envelope.
127 */
128export function makeStoredThreadSnapshotEncoder(options: {
129 readonly format: "full" | "compact";
130 readonly yieldToHost?: Effect.Effect<void>;
131 readonly maxJsonLength?: number;
132}) {
133 const yieldToHost = options.yieldToHost ?? yieldToHostTimer;
134 const maxJsonLength = options.maxJsonLength ?? THREAD_CACHE_MAX_JSON_LENGTH;
135 const encodedTurnItems = new WeakMap<object, EncodedRow>();
136 const encodedVisibleRows = new WeakMap<object, EncodedRow>();
137
138 return Effect.fnUntraced(function* (stored: StoredThread) {
139 let hasWorked = false;
140 let length = 0;
141 const encodeRows =
142 /* Runs once per thread save inside the traced store call; avoid a span per row list. */
143 Effect.fnUntraced(function* <Row extends object>(
144 rows: ReadonlyArray<Row>,
145 cache: WeakMap<object, EncodedRow>,
146 encodeChunk: (
147 chunk: ReadonlyArray<Row>,
148 ) => Effect.Effect<ReadonlyArray<unknown>, Schema.SchemaError>,
149 ) {
150 let misses: Array<Row> = [];
151 const encodeMisses = Effect.gen(function* () {
152 if (hasWorked) yield* yieldToHost;
153 hasWorked = true;
154 const chunk = yield* encodeChunk(misses);
155 chunk.forEach((encoded, index) => {
156 const row = { encoded, length: JSON.stringify(encoded).length };
157 cache.set(misses[index]!, row);
158 length += row.length;
159 });
160 misses = [];
161 });
162 for (const row of rows) {
163 const cached = cache.get(row);
164 if (cached === undefined) misses.push(row);
165 else length += cached.length;
166 if (misses.length === ROWS_PER_CHUNK) yield* encodeMisses;
167 if (length > maxJsonLength) return undefined;
168 }
169 if (misses.length > 0) yield* encodeMisses;
170 if (length > maxJsonLength) return undefined;
171 return rows.map((row) => cache.get(row)!.encoded);
172 });
173 const encodeProjectionRows =
174 /* Runs once per thread save inside the traced store call. */
175 Effect.fnUntraced(function* (projection: Projection) {
176 const turnItems = yield* encodeRows(
177 projection.turnItems,
178 encodedTurnItems,
179 encodeTurnItemChunk,
180 );
181 if (turnItems === undefined) return undefined;
182 const visibleTurnItems = yield* encodeRows(
183 projection.visibleTurnItems,
184 encodedVisibleRows,
185 encodeVisibleRowChunk,
186 );
187 if (visibleTurnItems === undefined) return undefined;
188 if (hasWorked) yield* yieldToHost;
189 return { ...projection, turnItems, visibleTurnItems };
190 });
191
192 const { snapshot } = stored;
193 let payload: string;
194 if (options.format === "full") {
195 const projection = yield* encodeProjectionRows(snapshot.projection);
196 if (projection === undefined) return Option.none<string>();
197 payload = yield* encodeFullEnvelope({
198 ...stored,
199 schemaVersion: ORCHESTRATION_THREAD_CACHE_SCHEMA_VERSION,
200 snapshot: { ...snapshot, projection },
201 });
202 } else {
203 const compact = omitLocalVisibleTurnItems(snapshot.projection);
204 const projection = yield* encodeProjectionRows(compact ?? snapshot.projection);
205 if (projection === undefined) return Option.none<string>();
206 payload = yield* encodeCompactEnvelope({
207 ...stored,
208 schemaVersion: COMPACT_THREAD_CACHE_SCHEMA_VERSION,
209 snapshot: {
210 ...snapshot,
211 projection,
212 turnItemsOmitLocalVisible: compact !== null,
213 turnItemCount: snapshot.projection.turnItems.length,
214 visibleTurnItemCount: snapshot.projection.visibleTurnItems.length,
215 },
216 });
217 }
218 // The rest of the projection is not budgeted row by row; a huge one still stays out.
219 return payload.length > maxJsonLength ? Option.none<string>() : Option.some(payload);
220 });
221}
222
223export const decodeStoredThreadSnapshot = Effect.fn("ThreadCache.decode")(function* (raw: string) {
224 const stored = yield* decodeCompatible(raw);
225 if (stored.schemaVersion !== COMPACT_THREAD_CACHE_SCHEMA_VERSION) return stored;
226 const { turnItemsOmitLocalVisible, turnItemCount, visibleTurnItemCount, ...snapshot } =
227 stored.snapshot;
228 const { projection } = snapshot;
229 const threadId = projection.thread.id;
230 let previousPosition = -1;
231 const visibleIds = new Set<string>();
232 const validRows = projection.visibleTurnItems.every((row) => {
233 const identity = JSON.stringify([row.sourceThreadId, row.sourceItemId]);
234 const valid =
235 row.sourceItemId === row.item.id &&
236 (row.sourceThreadId === row.item.threadId ||
237 (row.item.type === "fork" &&
238 row.item.source.type === "run" &&
239 row.sourceThreadId === row.item.source.threadId)) &&
240 (row.visibility !== "local" || row.sourceThreadId === threadId) &&
241 (row.visibility !== "inherited" || row.sourceThreadId !== threadId) &&
242 row.position > previousPosition &&
243 !visibleIds.has(identity);
244 previousPosition = row.position;
245 visibleIds.add(identity);
246 return valid;
247 });
248 const restored = turnItemsOmitLocalVisible
249 ? restoreLocalVisibleTurnItems(projection)
250 : projection;
251 const itemIds = new Set<string>();
252 const validItems = restored.turnItems.every((item) => {
253 const valid = item.threadId === threadId && !itemIds.has(item.id);
254 itemIds.add(item.id);
255 return valid;
256 });
257 if (
258 stored.threadId !== threadId ||
259 restored.turnItems.length !== turnItemCount ||
260 projection.visibleTurnItems.length !== visibleTurnItemCount ||
261 !validRows ||
262 !validItems ||
263 projection.visibleTurnItems.some(
264 (row) => row.sourceThreadId === threadId && !itemIds.has(row.sourceItemId),
265 )
266 ) {
267 return yield* new InvalidThreadCache();
268 }
269 return { ...stored, snapshot: { ...snapshot, projection: restored } };
270});