1import {
2 OrchestrationGetTurnDiffResult,
3 type OrchestrationGetFullThreadDiffInput,
4 type OrchestrationGetFullThreadDiffResult,
5 type OrchestrationGetTurnDiffInput,
6 type OrchestrationGetTurnDiffResult as OrchestrationGetTurnDiffResultType,
7 type ThreadId,
8} from "@t3tools/contracts";
9import * as Context from "effect/Context";
10import * as Digest from "effect/crypto/Digest";
11import * as Effect from "effect/Effect";
12import * as Layer from "effect/Layer";
13import * as Schema from "effect/Schema";
14
15import { CHECKPOINT_DIFF_MAX_OUTPUT_BYTES } from "../vcs/VcsDriver.ts";
16import { checkpointRefForScopeOrdinal } from "../orchestration-v2/CheckpointService.ts";
17import * as ThreadManagement from "../orchestration-v2/ThreadManagementService.ts";
18import {
19 CheckpointDiffResultInvalidError,
20 CheckpointRefUnavailableError,
21 CheckpointThreadNotFoundError,
22 CheckpointTurnRangeUnavailableError,
23 CheckpointWorkspacePathMissingError,
24 type CheckpointServiceError,
25} from "./Errors.ts";
26import * as CheckpointStore from "./CheckpointStore.ts";
27import { isGitImport, parseTurnDiffFilesFromNumstat, type TurnDiffFileSummary } from "./Diffs.ts";
28
29// Windows limits a command line to 32,767 characters. Leave room for the rest of git's arguments.
30const MAX_PATHSPEC_CHARS = 24_000;
31const DIFF_TRUNCATED_MARKER = "\n\n[truncated]";
32
33function batchFilePaths(files: ReadonlyArray<TurnDiffFileSummary>) {
34 const batches: string[][] = [];
35 let batch: string[] = [];
36 let length = 0;
37 for (const file of files) {
38 const paths = file.previousPath === undefined ? [file.path] : [file.previousPath, file.path];
39 const size = paths.reduce((total, path) => total + path.length + 16, 0);
40 if (batch.length > 0 && length + size > MAX_PATHSPEC_CHARS) {
41 batches.push(batch);
42 batch = [];
43 length = 0;
44 }
45 batch.push(...paths);
46 length += size;
47 }
48 if (batch.length > 0) batches.push(batch);
49 return batches;
50}
51
52export class CheckpointDiffQuery extends Context.Service<
53 CheckpointDiffQuery,
54 {
55 readonly getTurnDiff: (
56 input: OrchestrationGetTurnDiffInput,
57 ) => Effect.Effect<OrchestrationGetTurnDiffResultType, CheckpointServiceError>;
58
59 readonly getFullThreadDiff: (
60 input: OrchestrationGetFullThreadDiffInput,
61 ) => Effect.Effect<OrchestrationGetFullThreadDiffResult, CheckpointServiceError>;
62 }
63>()("t3/checkpointing/CheckpointDiffQuery") {}
64
65const isTurnDiffResult = Schema.is(OrchestrationGetTurnDiffResult);
66
67function buildTurnDiffResult(
68 input: {
69 readonly threadId: ThreadId;
70 readonly fromTurnCount: number;
71 readonly toTurnCount: number;
72 },
73 diff: string,
74): OrchestrationGetTurnDiffResultType {
75 return {
76 threadId: input.threadId,
77 fromTurnCount: input.fromTurnCount,
78 toTurnCount: input.toTurnCount,
79 diff,
80 };
81}
82
83/** @public */
84export const make = Effect.gen(function* () {
85 const threads = yield* ThreadManagement.ThreadManagementService;
86 const checkpointStore = yield* CheckpointStore.CheckpointStore;
87 const digest = yield* Digest.Digest;
88
89 const getTurnDiff: CheckpointDiffQuery["Service"]["getTurnDiff"] = Effect.fn("getTurnDiff")(
90 function* (input) {
91 const operation = "CheckpointDiffQuery.getTurnDiff";
92 const ignoreWhitespace = input.ignoreWhitespace ?? true;
93 yield* Effect.annotateCurrentSpan({
94 "checkpoint.thread_id": input.threadId,
95 "checkpoint.from_turn_count": input.fromTurnCount,
96 "checkpoint.to_turn_count": input.toTurnCount,
97 "checkpoint.ignore_whitespace": ignoreWhitespace,
98 });
99
100 if (input.fromTurnCount === input.toTurnCount) {
101 const emptyDiff = buildTurnDiffResult(input, "");
102 if (!isTurnDiffResult(emptyDiff)) {
103 return yield* new CheckpointDiffResultInvalidError({
104 operation,
105 threadId: input.threadId,
106 });
107 }
108 return emptyDiff;
109 }
110
111 const projection = yield* threads.getCheckpointContext(input.threadId).pipe(
112 Effect.mapError(
113 () =>
114 new CheckpointThreadNotFoundError({
115 operation,
116 threadId: input.threadId,
117 }),
118 ),
119 Effect.withSpan("checkpoint.turnDiff.lookupContext"),
120 );
121 const completedRunIds = new Set(
122 projection.runs.filter((run) => run.status === "completed").map((run) => run.id),
123 );
124 const readyCheckpoints = projection.checkpoints.filter(
125 (checkpoint) =>
126 checkpoint.status === "ready" &&
127 checkpoint.appRunOrdinal !== null &&
128 checkpoint.runId !== null &&
129 completedRunIds.has(checkpoint.runId),
130 );
131 const maxTurnCount = readyCheckpoints.reduce(
132 (max, checkpoint) => Math.max(max, checkpoint.appRunOrdinal ?? 0),
133 0,
134 );
135 if (input.toTurnCount > maxTurnCount) {
136 return yield* new CheckpointTurnRangeUnavailableError({
137 operation,
138 threadId: input.threadId,
139 requestedTurnCount: input.toTurnCount,
140 availableTurnCount: maxTurnCount,
141 });
142 }
143
144 const toCheckpoint = readyCheckpoints.find(
145 (checkpoint) => checkpoint.appRunOrdinal === input.toTurnCount,
146 );
147 if (toCheckpoint === undefined) {
148 return yield* new CheckpointRefUnavailableError({
149 operation,
150 threadId: input.threadId,
151 turnCount: input.toTurnCount,
152 checkpoint: "to",
153 });
154 }
155
156 const toScope = projection.checkpointScopes.find(
157 (scope) => scope.id === toCheckpoint.scopeId,
158 );
159 if (toScope === undefined) {
160 return yield* new CheckpointWorkspacePathMissingError({
161 operation,
162 threadId: input.threadId,
163 });
164 }
165
166 // Capture compares against the preceding scope ordinal, even when that
167 // checkpoint belongs to a stopped run or has no run.
168 const parentCheckpoint =
169 input.fromTurnCount === input.toTurnCount - 1
170 ? projection.checkpoints.find(
171 (checkpoint) =>
172 checkpoint.status === "ready" &&
173 checkpoint.scopeId === toCheckpoint.scopeId &&
174 checkpoint.ordinalWithinScope === toCheckpoint.ordinalWithinScope - 1,
175 )
176 : undefined;
177
178 const firstScope =
179 input.fromTurnCount === 0
180 ? projection.checkpointScopes.find((scope) => scope.kind === "root_run")
181 : undefined;
182 const fromCheckpointRef =
183 input.fromTurnCount === 0
184 ? firstScope === undefined
185 ? undefined
186 : yield* checkpointRefForScopeOrdinal({
187 scopeId: firstScope.id,
188 ordinalWithinScope: 0,
189 }).pipe(Effect.provideService(Digest.Digest, digest))
190 : (
191 parentCheckpoint ??
192 readyCheckpoints.find(
193 (checkpoint) => checkpoint.appRunOrdinal === input.fromTurnCount,
194 )
195 )?.ref;
196 if (fromCheckpointRef === undefined) {
197 return yield* new CheckpointRefUnavailableError({
198 operation,
199 threadId: input.threadId,
200 turnCount: input.fromTurnCount,
201 checkpoint: "from",
202 });
203 }
204
205 const comparison = {
206 cwd: toScope.cwd,
207 fromCheckpointRef,
208 toCheckpointRef: toCheckpoint.ref,
209 fallbackFromToHead: false,
210 ignoreWhitespace,
211 };
212 const authoredPaths = yield* checkpointStore
213 .listAuthoredPaths(comparison)
214 .pipe(Effect.orElseSucceed(() => null));
215 const files =
216 authoredPaths === null
217 ? []
218 : yield* checkpointStore.diffCheckpoints({ ...comparison, format: "numstat" }).pipe(
219 Effect.map(parseTurnDiffFilesFromNumstat),
220 Effect.orElseSucceed(() => []),
221 );
222 const retainedFiles = files.filter((file) => !isGitImport(file, authoredPaths));
223 // Select retained paths before generating a patch, so imported bulk cannot exhaust its output limit.
224 const pathBatches =
225 retainedFiles.length === files.length ? [undefined] : batchFilePaths(retainedFiles);
226 let diff = "";
227 let remainingBytes = CHECKPOINT_DIFF_MAX_OUTPUT_BYTES;
228 for (const filePaths of pathBatches) {
229 const patch = yield* checkpointStore
230 .diffCheckpoints({
231 ...comparison,
232 ...(filePaths === undefined ? {} : { filePaths }),
233 // One extra byte distinguishes an exact fit from a truncated patch.
234 maxOutputBytes: remainingBytes + 1,
235 })
236 .pipe(Effect.withSpan("checkpoint.turnDiff.diffCheckpoints"));
237 const patchBytes = Buffer.byteLength(patch);
238 diff += patch;
239 if (patchBytes > remainingBytes) {
240 // Reserve marker space only on overflow, which may trim an earlier batch.
241 // Streaming decode drops an incomplete final UTF-8 character at the byte boundary.
242 diff =
243 new TextDecoder().decode(
244 Buffer.from(diff).subarray(
245 0,
246 CHECKPOINT_DIFF_MAX_OUTPUT_BYTES - Buffer.byteLength(DIFF_TRUNCATED_MARKER),
247 ),
248 { stream: true },
249 ) + DIFF_TRUNCATED_MARKER;
250 break;
251 }
252 remainingBytes -= patchBytes;
253 }
254
255 const turnDiff = buildTurnDiffResult(input, diff);
256 if (!isTurnDiffResult(turnDiff)) {
257 return yield* new CheckpointDiffResultInvalidError({
258 operation,
259 threadId: input.threadId,
260 });
261 }
262
263 return turnDiff;
264 },
265 );
266
267 const getFullThreadDiff: CheckpointDiffQuery["Service"]["getFullThreadDiff"] = Effect.fn(
268 "CheckpointDiffQuery.getFullThreadDiff",
269 )(function* (input) {
270 const operation = "CheckpointDiffQuery.getFullThreadDiff";
271 const ignoreWhitespace = input.ignoreWhitespace ?? true;
272 yield* Effect.annotateCurrentSpan({
273 "checkpoint.thread_id": input.threadId,
274 "checkpoint.from_turn_count": 0,
275 "checkpoint.to_turn_count": input.toTurnCount,
276 "checkpoint.ignore_whitespace": ignoreWhitespace,
277 "checkpoint.diff_kind": "full-thread",
278 });
279
280 if (input.toTurnCount === 0) {
281 const emptyDiff = buildTurnDiffResult(
282 {
283 threadId: input.threadId,
284 fromTurnCount: 0,
285 toTurnCount: 0,
286 },
287 "",
288 );
289 if (!isTurnDiffResult(emptyDiff)) {
290 return yield* new CheckpointDiffResultInvalidError({
291 operation,
292 threadId: input.threadId,
293 });
294 }
295 return emptyDiff satisfies OrchestrationGetFullThreadDiffResult;
296 }
297
298 const turnDiff = yield* getTurnDiff({
299 threadId: input.threadId,
300 fromTurnCount: 0,
301 toTurnCount: input.toTurnCount,
302 ignoreWhitespace,
303 });
304 if (!isTurnDiffResult(turnDiff)) {
305 return yield* new CheckpointDiffResultInvalidError({
306 operation,
307 threadId: input.threadId,
308 });
309 }
310
311 return turnDiff satisfies OrchestrationGetFullThreadDiffResult;
312 });
313
314 return CheckpointDiffQuery.of({
315 getTurnDiff,
316 getFullThreadDiff,
317 });
318});
319
320export const layer = Layer.effect(CheckpointDiffQuery, make);