1import {
2 OrchestrationV2CheckpointFileSummary,
3 type CheckpointId,
4 type MessageId,
5 type NodeId,
6 type OrchestrationV2Checkpoint,
7 type OrchestrationV2Subagent,
8 type OrchestrationV2ThreadProjection,
9 type OrchestrationV2TurnItem,
10} from "@t3tools/contracts";
11import * as Schema from "effect/Schema";
12
13/**
14 * For clients that opt in, the server blanks a turn item's body that its record already carries
15 * (message text, checkpoint files, a subagent's prompt and result) and marks the field. The client
16 * copies the value back from the record. Only a body equal to the record's is blanked, so the
17 * restore is exact.
18 */
19export interface DedupedBodySource {
20 readonly messageText: (id: MessageId) => string | undefined;
21 readonly checkpointFiles: (id: CheckpointId) => OrchestrationV2Checkpoint["files"] | undefined;
22 readonly subagent: (id: NodeId) => WithSubagentBodies | undefined;
23}
24
25type SubagentBodies = Pick<OrchestrationV2Subagent, "prompt" | "result">;
26type WithSubagentBodies = SubagentBodies & {
27 readonly promptOmitted?: true;
28 readonly resultOmitted?: true;
29};
30
31const sameCheckpointFiles = Schema.toEquivalence(
32 Schema.Array(OrchestrationV2CheckpointFileSummary),
33);
34
35/** For subagent items and for a repeated `subagent.updated`, whose `previous` was sent earlier. */
36export function dedupeSubagentBodies<T extends WithSubagentBodies>(
37 value: T,
38 previous: SubagentBodies | undefined,
39): T {
40 const prompt = value.prompt !== "" && previous?.prompt === value.prompt;
41 const result = value.result !== null && previous?.result === value.result;
42 if (!prompt && !result) return value;
43 return {
44 ...value,
45 ...(prompt ? { prompt: "", promptOmitted: true as const } : {}),
46 ...(result ? { result: null, resultOmitted: true as const } : {}),
47 };
48}
49
50/**
51 * A marked body whose record is missing, or whose record field is itself still marked, stays
52 * blank and marked; the server only blanks bodies it sent, so either means a server bug.
53 */
54export function restoreSubagentBodies<T extends WithSubagentBodies>(
55 value: T,
56 previous: WithSubagentBodies | undefined,
57): T {
58 if (previous === undefined) return value;
59 const prompt = value.promptOmitted === true && previous.promptOmitted !== true;
60 const result = value.resultOmitted === true && previous.resultOmitted !== true;
61 if (!prompt && !result) return value;
62 const { promptOmitted, resultOmitted, ...rest } = value;
63 // Dropping both optional markers from `T` and keeping the unrestored one leaves a `T`.
64 return {
65 ...rest,
66 ...(prompt ? { prompt: previous.prompt } : promptOmitted === true ? { promptOmitted } : {}),
67 ...(result ? { result: previous.result } : resultOmitted === true ? { resultOmitted } : {}),
68 } as T;
69}
70
71export function dedupeTurnItemBodies(
72 item: OrchestrationV2TurnItem,
73 source: DedupedBodySource,
74): OrchestrationV2TurnItem {
75 switch (item.type) {
76 case "user_message":
77 case "assistant_message":
78 return item.text !== "" && source.messageText(item.messageId) === item.text
79 ? { ...item, text: "", textOmitted: true }
80 : item;
81 case "checkpoint": {
82 const files = source.checkpointFiles(item.checkpointId);
83 return item.files.length > 0 && files !== undefined && sameCheckpointFiles(files, item.files)
84 ? { ...item, files: [], filesOmitted: true }
85 : item;
86 }
87 case "subagent":
88 return dedupeSubagentBodies(item, source.subagent(item.subagentId));
89 default:
90 return item;
91 }
92}
93
94export function restoreTurnItemBodies(
95 item: OrchestrationV2TurnItem,
96 source: DedupedBodySource,
97): OrchestrationV2TurnItem {
98 switch (item.type) {
99 case "user_message":
100 case "assistant_message": {
101 const text = item.textOmitted === true ? source.messageText(item.messageId) : undefined;
102 if (text === undefined) return item;
103 const { textOmitted: _textOmitted, ...rest } = item;
104 return { ...rest, text };
105 }
106 case "checkpoint": {
107 const files =
108 item.filesOmitted === true ? source.checkpointFiles(item.checkpointId) : undefined;
109 if (files === undefined) return item;
110 const { filesOmitted: _filesOmitted, ...rest } = item;
111 return { ...rest, files };
112 }
113 case "subagent":
114 return item.promptOmitted === true || item.resultOmitted === true
115 ? restoreSubagentBodies(item, source.subagent(item.subagentId))
116 : item;
117 default:
118 return item;
119 }
120}
121
122/** Builds each index on first lookup, so a snapshot without marked items builds none. */
123function snapshotBodySource(projection: OrchestrationV2ThreadProjection): DedupedBodySource {
124 let messages: Map<MessageId, string> | undefined;
125 let checkpoints: Map<CheckpointId, OrchestrationV2Checkpoint["files"]> | undefined;
126 let subagents: Map<NodeId, OrchestrationV2Subagent> | undefined;
127 return {
128 messageText: (id) =>
129 (messages ??= new Map(projection.messages.map((message) => [message.id, message.text]))).get(
130 id,
131 ),
132 checkpointFiles: (id) =>
133 (checkpoints ??= new Map(
134 projection.checkpoints.map((checkpoint) => [checkpoint.id, checkpoint.files]),
135 )).get(id),
136 subagent: (id) =>
137 (subagents ??= new Map(projection.subagents.map((subagent) => [subagent.id, subagent]))).get(
138 id,
139 ),
140 };
141}
142
143/**
144 * Inherited rows belong to another thread, whose records this projection lacks. Maps each item
145 * once, so an item shared by `visibleTurnItems` and `turnItems` stays one object.
146 */
147function mapLocalTurnItems(
148 projection: OrchestrationV2ThreadProjection,
149 map: (item: OrchestrationV2TurnItem) => OrchestrationV2TurnItem,
150): OrchestrationV2ThreadProjection {
151 const mapped = new Map<OrchestrationV2TurnItem, OrchestrationV2TurnItem>();
152 let changed = false;
153 const once = (item: OrchestrationV2TurnItem) => {
154 let next = mapped.get(item);
155 if (next === undefined) {
156 next = map(item);
157 mapped.set(item, next);
158 changed ||= next !== item;
159 }
160 return next;
161 };
162 const visibleTurnItems = projection.visibleTurnItems.map((row) => {
163 if (row.visibility !== "local" && row.sourceThreadId !== projection.thread.id) return row;
164 const item = once(row.item);
165 return item === row.item ? row : { ...row, item };
166 });
167 const turnItems = projection.turnItems.map(once);
168 return changed ? { ...projection, visibleTurnItems, turnItems } : projection;
169}
170
171export function dedupeSnapshotBodies(
172 projection: OrchestrationV2ThreadProjection,
173): OrchestrationV2ThreadProjection {
174 const source = snapshotBodySource(projection);
175 return mapLocalTurnItems(projection, (item) => dedupeTurnItemBodies(item, source));
176}
177
178/** Call on every decoded snapshot before reducers or caches see it. */
179export function restoreSnapshotBodies(
180 projection: OrchestrationV2ThreadProjection,
181): OrchestrationV2ThreadProjection {
182 const source = snapshotBodySource(projection);
183 return mapLocalTurnItems(projection, (item) => restoreTurnItemBodies(item, source));
184}