1import {
2 CheckpointId,
3 CheckpointScopeId,
4 latestProviderTurnForAttempt,
5 type OrchestrationV2DomainEvent,
6 ProviderThreadId,
7 ThreadId,
8} from "@t3tools/contracts";
9import * as Context from "effect/Context";
10import * as DateTime from "effect/DateTime";
11import * as Effect from "effect/Effect";
12import * as FileSystem from "effect/FileSystem";
13import * as Layer from "effect/Layer";
14import * as Option from "effect/Option";
15import * as Path from "effect/Path";
16import * as Schema from "effect/Schema";
17
18import * as GitWorkflowService from "../git/GitWorkflowService.ts";
19import { withWorkspaceLease } from "../workspace/workspaceLease.ts";
20import {
21 isCheckpointRestoreIsolated,
22 SHARED_WORKSPACE_RESTORE_MESSAGE,
23} from "./CheckpointRestoreSafety.ts";
24import { CheckpointServiceV2 } from "./CheckpointService.ts";
25import { EventSinkV2 } from "./EventSink.ts";
26import * as IdAllocator from "@t3tools/provider-core/server/IdAllocator";
27import { ProjectionStoreV2 } from "./ProjectionStore.ts";
28import * as ProjectStore from "./ProjectStore.ts";
29import type * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter";
30import { ProviderSessionManagerV2 } from "./ProviderSessionManager.ts";
31import { RuntimePolicyV2 } from "./RuntimePolicy.ts";
32import * as ThreadCommandExecutor from "./ThreadCommandExecutor.ts";
33
34export const ROLLBACK_FAILED_MESSAGE =
35 "The provider could not roll back this conversation. Try again; if it keeps failing, check the provider and server logs.";
36
37export class CheckpointRollbackExecutionError extends Schema.TaggedError<CheckpointRollbackExecutionError>()(
38 "CheckpointRollbackExecutionError",
39 {
40 reason: Schema.Literals([
41 "rollback-target-invalid",
42 "active-provider-changed",
43 "provider-turn-unavailable",
44 "unexpected-failure",
45 "shared-workspace",
46 ]),
47 threadId: ThreadId,
48 providerThreadId: ProviderThreadId,
49 checkpointId: CheckpointId,
50 cause: Schema.optional(Schema.Defect()),
51 },
52) {
53 override get message(): string {
54 switch (this.reason) {
55 case "rollback-target-invalid":
56 return `Rollback target ${this.checkpointId} for provider thread ${this.providerThreadId} on thread ${this.threadId} is incomplete or invalid.`;
57 case "active-provider-changed":
58 return `Active provider changed before rollback target ${this.checkpointId} could execute on thread ${this.threadId}.`;
59 case "provider-turn-unavailable":
60 return `Provider turn for rollback target ${this.checkpointId} is unavailable on provider thread ${this.providerThreadId}.`;
61 case "shared-workspace":
62 return SHARED_WORKSPACE_RESTORE_MESSAGE;
63 case "unexpected-failure":
64 return ROLLBACK_FAILED_MESSAGE;
65 }
66 }
67}
68
69const isCheckpointRollbackExecutionError = Schema.is(CheckpointRollbackExecutionError);
70
71export interface CheckpointRollbackServiceV2Shape {
72 readonly execute: (input: {
73 readonly threadId: ThreadId;
74 readonly providerThreadId: ProviderThreadId;
75 readonly checkpointId: CheckpointId;
76 readonly scopeId: CheckpointScopeId;
77 readonly restoreFiles?: boolean;
78 }) => Effect.Effect<void, CheckpointRollbackExecutionError>;
79}
80
81export class CheckpointRollbackServiceV2 extends Context.Service<
82 CheckpointRollbackServiceV2,
83 CheckpointRollbackServiceV2Shape
84>()("t3/orchestration-v2/CheckpointRollbackService/CheckpointRollbackServiceV2") {}
85
86export const layer: Layer.Layer<
87 CheckpointRollbackServiceV2,
88 never,
89 | CheckpointServiceV2
90 | GitWorkflowService.GitWorkflowService
91 | EventSinkV2
92 | IdAllocator.IdAllocatorV2
93 | ProjectionStoreV2
94 | ProviderSessionManagerV2
95 | RuntimePolicyV2
96 | FileSystem.FileSystem
97 | Path.Path
98 | ProjectStore.ProjectStoreV2
99 | ThreadCommandExecutor.ThreadCommandExecutor
100> = Layer.effect(
101 CheckpointRollbackServiceV2,
102 Effect.gen(function* () {
103 const checkpoints = yield* CheckpointServiceV2;
104 const gitWorkflow = yield* GitWorkflowService.GitWorkflowService;
105 const eventSink = yield* EventSinkV2;
106 const ids = yield* IdAllocator.IdAllocatorV2;
107 const projections = yield* ProjectionStoreV2;
108 const sessions = yield* ProviderSessionManagerV2;
109 const runtimePolicy = yield* RuntimePolicyV2;
110 const fileSystem = yield* FileSystem.FileSystem;
111 const projects = yield* ProjectStore.ProjectStoreV2;
112 const path = yield* Path.Path;
113 const threadCommands = yield* ThreadCommandExecutor.ThreadCommandExecutor;
114
115 const execute = Effect.fn("orchestrationV2.checkpointRollback.execute")(function* (input: {
116 readonly threadId: ThreadId;
117 readonly providerThreadId: ProviderThreadId;
118 readonly checkpointId: CheckpointId;
119 readonly scopeId: CheckpointScopeId;
120 readonly restoreFiles?: boolean;
121 }) {
122 const projection = yield* projections.getThreadRecords(input.threadId, [
123 "providerThreads",
124 "providerSessions",
125 "checkpoints",
126 "checkpointScopes",
127 "runs",
128 "attempts",
129 "nodes",
130 "providerTurns",
131 ]);
132 const providerThread = projection.providerThreads.find(
133 (candidate) => candidate.id === input.providerThreadId,
134 );
135 const checkpoint = projection.checkpoints.find(
136 (candidate) => candidate.id === input.checkpointId,
137 );
138 const scope = projection.checkpointScopes.find((candidate) => candidate.id === input.scopeId);
139 if (
140 providerThread === undefined ||
141 providerThread.providerSessionId === null ||
142 checkpoint === undefined ||
143 scope === undefined ||
144 checkpoint.scopeId !== scope.id ||
145 checkpoint.status !== "ready"
146 ) {
147 return yield* new CheckpointRollbackExecutionError({
148 reason: "rollback-target-invalid",
149 threadId: input.threadId,
150 providerThreadId: input.providerThreadId,
151 checkpointId: input.checkpointId,
152 });
153 }
154 if (
155 providerThread.id !== projection.thread.activeProviderThreadId ||
156 providerThread.providerInstanceId !== projection.thread.modelSelection.instanceId
157 ) {
158 return yield* new CheckpointRollbackExecutionError({
159 reason: "active-provider-changed",
160 threadId: input.threadId,
161 providerThreadId: input.providerThreadId,
162 checkpointId: input.checkpointId,
163 });
164 }
165
166 const targetOrdinal = checkpoint.appRunOrdinal ?? 0;
167 const rollbackTarget: ProviderAdapter.ProviderAdapterV2RollbackTarget =
168 targetOrdinal === 0
169 ? {
170 type: "thread_start",
171 checkpointId: checkpoint.id,
172 appRunOrdinal: 0,
173 }
174 : yield* Effect.gen(function* () {
175 const targetRun = projection.runs.find((run) => run.ordinal === targetOrdinal);
176 const targetAttempt = projection.attempts.find(
177 (attempt) => attempt.id === targetRun?.activeAttemptId,
178 );
179 // A goal run can span several native turns; roll back to its last.
180 const targetTurn =
181 latestProviderTurnForAttempt(projection.providerTurns, targetAttempt?.id) ??
182 projection.providerTurns.find((turn) => turn.id === targetAttempt?.providerTurnId);
183 if (targetTurn === undefined || targetTurn.providerThreadId !== providerThread.id) {
184 return yield* new CheckpointRollbackExecutionError({
185 reason: "provider-turn-unavailable",
186 threadId: input.threadId,
187 providerThreadId: input.providerThreadId,
188 checkpointId: input.checkpointId,
189 });
190 }
191 return {
192 type: "provider_turn" as const,
193 checkpointId: checkpoint.id,
194 appRunOrdinal: targetOrdinal,
195 providerTurn: targetTurn,
196 };
197 });
198
199 const { worktreePath, branch } = projection.thread;
200 if (worktreePath != null && branch != null && !(yield* fileSystem.exists(worktreePath))) {
201 const project = yield* projects.get(projection.thread.projectId);
202 if (Option.isSome(project)) {
203 yield* gitWorkflow.pruneWorktrees({ cwd: project.value.workspaceRoot });
204 yield* gitWorkflow.createWorktree({
205 cwd: project.value.workspaceRoot,
206 refName: branch,
207 path: worktreePath,
208 });
209 }
210 }
211
212 if (
213 input.restoreFiles !== false &&
214 !(yield* isCheckpointRestoreIsolated(projection.thread, scope, {
215 fileSystem,
216 projections,
217 projects,
218 path,
219 }))
220 ) {
221 return yield* new CheckpointRollbackExecutionError({
222 reason: "shared-workspace",
223 threadId: input.threadId,
224 providerThreadId: input.providerThreadId,
225 checkpointId: input.checkpointId,
226 });
227 }
228
229 const modelSelection = projection.thread.modelSelection;
230 const resolvedRuntimePolicy = yield* runtimePolicy.resolve({
231 thread: projection.thread,
232 modelSelection,
233 });
234 const existingSession = projection.providerSessions.find(
235 (candidate) => candidate.id === providerThread.providerSessionId,
236 );
237 const session = yield* sessions.open({
238 threadId: input.threadId,
239 providerSessionId: providerThread.providerSessionId,
240 modelSelection,
241 runtimePolicy: resolvedRuntimePolicy,
242 ...(existingSession === undefined ? {} : { resumeFromSession: existingSession }),
243 ...(providerThread.nativeThreadRef?.nativeId == null
244 ? {}
245 : { initialNativeThreadId: providerThread.nativeThreadRef.nativeId }),
246 ...(providerThread.nativeMetadata?.itemIdentityVersion === undefined
247 ? {}
248 : {
249 initialProviderItemIdentityVersion: providerThread.nativeMetadata.itemIdentityVersion,
250 }),
251 });
252
253 // Stopped and failed runs after the target leave the provider
254 // conversation too, so they must not stay visible.
255 const runsToRollback = projection.runs.filter(
256 (run) =>
257 run.ordinal > targetOrdinal &&
258 (run.status === "completed" ||
259 run.status === "interrupted" ||
260 run.status === "failed" ||
261 run.status === "cancelled"),
262 );
263 // Rolled-back turns stay in the audit history, but no longer exist in
264 // the provider conversation and must not be counted by a later rewind.
265 const rolledBackRunIds = new Set(
266 projection.runs.filter((run) => run.status === "rolled_back").map((run) => run.id),
267 );
268 const rolledBackAttemptIds = new Set(
269 projection.attempts
270 .filter((attempt) => rolledBackRunIds.has(attempt.runId))
271 .map((attempt) => attempt.id),
272 );
273 const providerThreadTurns = projection.providerTurns.filter(
274 (turn) =>
275 turn.providerThreadId === providerThread.id &&
276 (turn.runAttemptId === null || !rolledBackAttemptIds.has(turn.runAttemptId)),
277 );
278
279 const snapshot =
280 runsToRollback.length === 0
281 ? { providerThread }
282 : yield* session.rollbackThread({
283 providerThread,
284 target: rollbackTarget,
285 providerThreadTurns,
286 });
287 if (input.restoreFiles !== false) yield* checkpoints.restore({ scope, checkpoint });
288 const staleCheckpoints = projection.checkpoints.filter(
289 (candidate) =>
290 candidate.scopeId === scope.id &&
291 candidate.appRunOrdinal !== null &&
292 candidate.appRunOrdinal > targetOrdinal &&
293 candidate.status === "ready",
294 );
295 if (staleCheckpoints.length > 0) {
296 yield* checkpoints.deleteStaleRefs({ scope, checkpoints: staleCheckpoints });
297 }
298
299 // Thread commands write full provider thread rows under this lock. Without
300 // it, a message sent during the rollback can plan against the old native
301 // session and commit after this, undoing the provider's fork.
302 yield* threadCommands.withLock(
303 input.threadId,
304 Effect.gen(function* () {
305 const now = yield* DateTime.now;
306 const makeEvent = <Event extends OrchestrationV2DomainEvent>(event: Omit<Event, "id">) =>
307 Effect.map(
308 ids.allocate.event({ threadId: event.threadId }),
309 (id) =>
310 ({
311 ...event,
312 id,
313 }) as Event,
314 );
315 const events: Array<OrchestrationV2DomainEvent> = [];
316 events.push(
317 yield* makeEvent({
318 type: "provider-thread.updated",
319 threadId: input.threadId,
320 driver: providerThread.driver,
321 providerInstanceId: providerThread.providerInstanceId,
322 occurredAt: now,
323 payload: {
324 ...snapshot.providerThread,
325 lastRunOrdinal: targetOrdinal === 0 ? null : targetOrdinal,
326 updatedAt: now,
327 },
328 }),
329 );
330 for (const staleCheckpoint of staleCheckpoints) {
331 events.push(
332 yield* makeEvent({
333 type: "checkpoint.captured",
334 threadId: input.threadId,
335 ...(staleCheckpoint.runId === null ? {} : { runId: staleCheckpoint.runId }),
336 nodeId: staleCheckpoint.nodeId,
337 providerInstanceId: providerThread.providerInstanceId,
338 occurredAt: now,
339 payload: { ...staleCheckpoint, status: "stale" },
340 }),
341 );
342 }
343 for (const run of runsToRollback) {
344 const rootNode = projection.nodes.find((candidate) => candidate.id === run.rootNodeId);
345 events.push(
346 yield* makeEvent({
347 type: "run.updated",
348 threadId: input.threadId,
349 runId: run.id,
350 ...(rootNode === undefined ? {} : { nodeId: rootNode.id }),
351 providerInstanceId: run.providerInstanceId,
352 occurredAt: now,
353 payload: { ...run, status: "rolled_back", completedAt: now },
354 }),
355 );
356 if (rootNode !== undefined) {
357 events.push(
358 yield* makeEvent({
359 type: "node.updated",
360 threadId: input.threadId,
361 runId: run.id,
362 nodeId: rootNode.id,
363 providerInstanceId: run.providerInstanceId,
364 occurredAt: now,
365 payload: { ...rootNode, status: "rolled_back", completedAt: now },
366 }),
367 );
368 }
369 }
370 yield* eventSink.write({ events });
371 }),
372 );
373 });
374
375 return CheckpointRollbackServiceV2.of({
376 execute: (input) =>
377 Effect.gen(function* () {
378 // Storage cleanup removes worktrees under the same lease.
379 const { thread } = yield* projections.getThreadRecords(input.threadId, []);
380 yield* thread.worktreePath === null
381 ? execute(input)
382 : withWorkspaceLease(path.resolve(thread.worktreePath), execute(input));
383 }).pipe(
384 Effect.mapError((cause) =>
385 isCheckpointRollbackExecutionError(cause)
386 ? cause
387 : new CheckpointRollbackExecutionError({
388 reason: "unexpected-failure",
389 threadId: input.threadId,
390 providerThreadId: input.providerThreadId,
391 checkpointId: input.checkpointId,
392 cause,
393 }),
394 ),
395 ),
396 });
397 }),
398);