1import {
2 CheckpointId,
3 CheckpointScopeId,
4 latestProviderTurnForAttempt,
5 type OrchestrationV2DomainEvent,
6 ProviderThreadId,
7 ThreadId,
8 VcsProcessExitError,
9} from "@t3tools/contracts";
10import * as Context from "effect/Context";
11import * as DateTime from "effect/DateTime";
12import * as Effect from "effect/Effect";
13import * as Exit from "effect/Exit";
14import * as FileSystem from "effect/FileSystem";
15import * as Layer from "effect/Layer";
16import * as Option from "effect/Option";
17import * as Path from "effect/Path";
18import * as Schema from "effect/Schema";
19
20import * as GitWorkflowService from "../git/GitWorkflowService.ts";
21import type * as ServerConfig from "../config.ts";
22import type * as ServerSettings from "../serverSettings.ts";
23import { makeIsManagedWorktreePath } from "../worktreesDirectory.ts";
24import { withWorkspaceLease } from "../workspace/workspaceLease.ts";
25import {
26 isCheckpointRestoreIsolated,
27 SHARED_WORKSPACE_RESTORE_MESSAGE,
28} from "./CheckpointRestoreSafety.ts";
29import { CheckpointRestoreError, CheckpointServiceV2 } from "./CheckpointService.ts";
30import { EventSinkV2 } from "./EventSink.ts";
31import * as IdAllocator from "@t3tools/provider-core/server/IdAllocator";
32import { ProjectionStoreV2 } from "./ProjectionStore.ts";
33import { GIT_INDEX_LOCKED_DETAIL } from "../vcs/VcsProcess.ts";
34import * as ProjectStore from "./ProjectStore.ts";
35import type * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter";
36import { ProviderSessionManagerV2 } from "./ProviderSessionManager.ts";
37import { RuntimePolicyAboveCeilingError, RuntimePolicyV2 } from "./RuntimePolicy.ts";
38import * as ThreadCommandExecutor from "./ThreadCommandExecutor.ts";
39
40export const ROLLBACK_FAILED_MESSAGE =
41 "The provider could not roll back this conversation. Try again; if it keeps failing, check the provider and server logs.";
42
43export class CheckpointRollbackExecutionError extends Schema.TaggedError<CheckpointRollbackExecutionError>()(
44 "CheckpointRollbackExecutionError",
45 {
46 reason: Schema.Literals([
47 "rollback-target-invalid",
48 "active-provider-changed",
49 "provider-turn-unavailable",
50 "unexpected-failure",
51 "shared-workspace",
52 ]),
53 threadId: ThreadId,
54 providerThreadId: ProviderThreadId,
55 checkpointId: CheckpointId,
56 cause: Schema.optional(Schema.Defect()),
57 },
58) {
59 override get message(): string {
60 switch (this.reason) {
61 case "rollback-target-invalid":
62 return `Rollback target ${this.checkpointId} for provider thread ${this.providerThreadId} on thread ${this.threadId} is incomplete or invalid.`;
63 case "active-provider-changed":
64 return `Active provider changed before rollback target ${this.checkpointId} could execute on thread ${this.threadId}.`;
65 case "provider-turn-unavailable":
66 return `Provider turn for rollback target ${this.checkpointId} is unavailable on provider thread ${this.providerThreadId}.`;
67 case "shared-workspace":
68 return SHARED_WORKSPACE_RESTORE_MESSAGE;
69 case "unexpected-failure":
70 return ROLLBACK_FAILED_MESSAGE;
71 }
72 }
73}
74
75const isCheckpointRollbackExecutionError = Schema.is(CheckpointRollbackExecutionError);
76const isVcsProcessExitError = Schema.is(VcsProcessExitError);
77const isCheckpointRestoreError = Schema.is(CheckpointRestoreError);
78const isRuntimePolicyAboveCeilingError = Schema.is(RuntimePolicyAboveCeilingError);
79
80/**
81 * The message clients see when a rollback fails. Only a held Git index lock
82 * gets its own explanation; other reasons carry internal ids or raw Git
83 * output, so they fall back to the generic message.
84 */
85export function rollbackFailureClientMessage(error: CheckpointRollbackExecutionError): string {
86 if (isRuntimePolicyAboveCeilingError(error.cause)) return error.cause.message;
87 const restoreCause = isCheckpointRestoreError(error.cause) ? error.cause.cause : undefined;
88 return isVcsProcessExitError(restoreCause) && restoreCause.detail === GIT_INDEX_LOCKED_DETAIL
89 ? GIT_INDEX_LOCKED_DETAIL
90 : ROLLBACK_FAILED_MESSAGE;
91}
92
93export interface CheckpointRollbackServiceV2Shape {
94 readonly execute: (input: {
95 readonly threadId: ThreadId;
96 readonly providerThreadId: ProviderThreadId;
97 readonly checkpointId: CheckpointId;
98 readonly scopeId: CheckpointScopeId;
99 readonly restoreFiles?: boolean;
100 }) => Effect.Effect<void, CheckpointRollbackExecutionError>;
101}
102
103export class CheckpointRollbackServiceV2 extends Context.Service<
104 CheckpointRollbackServiceV2,
105 CheckpointRollbackServiceV2Shape
106>()("t3/orchestration-v2/CheckpointRollbackService/CheckpointRollbackServiceV2") {}
107
108export const layer: Layer.Layer<
109 CheckpointRollbackServiceV2,
110 never,
111 | CheckpointServiceV2
112 | GitWorkflowService.GitWorkflowService
113 | EventSinkV2
114 | IdAllocator.IdAllocatorV2
115 | ProjectionStoreV2
116 | ProviderSessionManagerV2
117 | RuntimePolicyV2
118 | FileSystem.FileSystem
119 | Path.Path
120 | ProjectStore.ProjectStoreV2
121 | ServerConfig.ServerConfig
122 | ServerSettings.ServerSettingsService
123 | ThreadCommandExecutor.ThreadCommandExecutor
124> = Layer.effect(
125 CheckpointRollbackServiceV2,
126 Effect.gen(function* () {
127 const checkpoints = yield* CheckpointServiceV2;
128 const gitWorkflow = yield* GitWorkflowService.GitWorkflowService;
129 const eventSink = yield* EventSinkV2;
130 const ids = yield* IdAllocator.IdAllocatorV2;
131 const projections = yield* ProjectionStoreV2;
132 const sessions = yield* ProviderSessionManagerV2;
133 const runtimePolicy = yield* RuntimePolicyV2;
134 const fileSystem = yield* FileSystem.FileSystem;
135 const projects = yield* ProjectStore.ProjectStoreV2;
136 const path = yield* Path.Path;
137 const threadCommands = yield* ThreadCommandExecutor.ThreadCommandExecutor;
138 const isManagedWorktreePath = yield* makeIsManagedWorktreePath;
139
140 const execute = Effect.fn("orchestrationV2.checkpointRollback.execute")(function* (input: {
141 readonly threadId: ThreadId;
142 readonly providerThreadId: ProviderThreadId;
143 readonly checkpointId: CheckpointId;
144 readonly scopeId: CheckpointScopeId;
145 readonly restoreFiles?: boolean;
146 }) {
147 const context = yield* projections.getCheckpointRollbackContext(input.threadId, input);
148 const { providerThread, checkpoint, scope, targetAttempt, runsToRollback, staleCheckpoints } =
149 context;
150 if (
151 providerThread === undefined ||
152 providerThread.providerSessionId === null ||
153 checkpoint === undefined ||
154 scope === undefined ||
155 checkpoint.scopeId !== scope.id ||
156 // Missing snapshots support conversation rewind only.
157 (checkpoint.status !== "ready" &&
158 !(input.restoreFiles === false && checkpoint.status === "missing"))
159 ) {
160 return yield* new CheckpointRollbackExecutionError({
161 reason: "rollback-target-invalid",
162 threadId: input.threadId,
163 providerThreadId: input.providerThreadId,
164 checkpointId: input.checkpointId,
165 });
166 }
167 if (
168 providerThread.id !== context.thread.activeProviderThreadId ||
169 providerThread.providerInstanceId !== context.thread.modelSelection.instanceId
170 ) {
171 return yield* new CheckpointRollbackExecutionError({
172 reason: "active-provider-changed",
173 threadId: input.threadId,
174 providerThreadId: input.providerThreadId,
175 checkpointId: input.checkpointId,
176 });
177 }
178
179 const targetOrdinal = checkpoint.appRunOrdinal ?? 0;
180 const rollbackTarget: ProviderAdapter.ProviderAdapterV2RollbackTarget =
181 targetOrdinal === 0
182 ? {
183 type: "thread_start",
184 checkpointId: checkpoint.id,
185 appRunOrdinal: 0,
186 }
187 : yield* Effect.gen(function* () {
188 const targetTurn =
189 latestProviderTurnForAttempt(context.providerTurns, targetAttempt?.id) ??
190 context.providerTurns.find((turn) => turn.id === targetAttempt?.providerTurnId);
191 if (
192 targetAttempt === undefined ||
193 targetTurn === undefined ||
194 targetTurn.providerThreadId !== providerThread.id
195 ) {
196 return yield* new CheckpointRollbackExecutionError({
197 reason: "provider-turn-unavailable",
198 threadId: input.threadId,
199 providerThreadId: input.providerThreadId,
200 checkpointId: input.checkpointId,
201 });
202 }
203 const { attempt } = yield* projections.getProviderControlContext(input.threadId, {
204 providerThreadId: providerThread.id,
205 providerTurnId: targetTurn.id,
206 attemptId: targetAttempt.id,
207 });
208 const message =
209 attempt === undefined
210 ? undefined
211 : yield* projections.getRunMessage(input.threadId, attempt.runId);
212 return {
213 type: "provider_turn" as const,
214 checkpointId: checkpoint.id,
215 appRunOrdinal: targetOrdinal,
216 providerTurn: targetTurn,
217 ...(message === undefined ? {} : { message }),
218 };
219 });
220
221 const restoreFiles = input.restoreFiles !== false;
222 const isolated = isCheckpointRestoreIsolated(context.thread, scope, {
223 fileSystem,
224 projections,
225 projects,
226 path,
227 });
228 const sharedWorkspace = () =>
229 new CheckpointRollbackExecutionError({
230 reason: "shared-workspace",
231 threadId: input.threadId,
232 providerThreadId: input.providerThreadId,
233 checkpointId: input.checkpointId,
234 });
235
236 const { worktreePath, branch } = context.thread;
237 // Recreated only where the server would have made it, as turn start does.
238 if (worktreePath != null && branch != null) {
239 // A failed removal can leave the folder without its checkout; only a file restore replaces it.
240 const leftover =
241 restoreFiles &&
242 (yield* fileSystem.exists(worktreePath)) &&
243 !(yield* fileSystem.exists(path.join(worktreePath, ".git")));
244 if (
245 (leftover || !(yield* fileSystem.exists(worktreePath))) &&
246 (yield* isManagedWorktreePath(worktreePath))
247 ) {
248 const project = yield* projects.get(context.thread.projectId);
249 if (Option.isSome(project)) {
250 const recreate = gitWorkflow.pruneWorktrees({ cwd: project.value.workspaceRoot }).pipe(
251 Effect.andThen(
252 gitWorkflow.createWorktree({
253 cwd: project.value.workspaceRoot,
254 refName: branch,
255 path: worktreePath,
256 }),
257 ),
258 );
259 if (leftover) {
260 if (!(yield* isolated)) return yield* sharedWorkspace();
261 // Set aside, not removed: a checkout that fails puts the folder and its files back.
262 yield* Effect.acquireUseRelease(
263 Effect.gen(function* () {
264 const aside = yield* fileSystem.makeTempDirectory({
265 directory: path.dirname(worktreePath),
266 prefix: ".t3-replaced-",
267 });
268 yield* fileSystem.rename(worktreePath, path.join(aside, "folder"));
269 return aside;
270 }),
271 () => recreate,
272 (aside, exit) =>
273 (Exit.isSuccess(exit)
274 ? Effect.void
275 : fileSystem
276 .remove(worktreePath, { recursive: true, force: true })
277 .pipe(
278 Effect.andThen(
279 fileSystem.rename(path.join(aside, "folder"), worktreePath),
280 ),
281 )
282 ).pipe(
283 Effect.andThen(fileSystem.remove(aside, { recursive: true, force: true })),
284 Effect.catchCause((cause) =>
285 Effect.logWarning("Edit from here left a worktree folder set aside", {
286 worktreePath,
287 setAside: path.join(aside, "folder"),
288 cause,
289 }),
290 ),
291 ),
292 );
293 } else {
294 yield* recreate;
295 }
296 }
297 }
298 }
299
300 if (restoreFiles && !(yield* isolated)) return yield* sharedWorkspace();
301
302 const modelSelection = context.thread.modelSelection;
303 const resolvedRuntimePolicy = yield* runtimePolicy.resolve({
304 thread: context.thread,
305 modelSelection,
306 });
307 const existingSession = context.providerSession;
308 const session = yield* sessions.open({
309 threadId: input.threadId,
310 providerSessionId: providerThread.providerSessionId,
311 modelSelection,
312 runtimePolicy: resolvedRuntimePolicy,
313 ...(existingSession === undefined ? {} : { resumeFromSession: existingSession }),
314 ...(providerThread.nativeThreadRef?.nativeId == null
315 ? {}
316 : { initialNativeThreadId: providerThread.nativeThreadRef.nativeId }),
317 ...(providerThread.nativeMetadata?.itemIdentityVersion === undefined
318 ? {}
319 : {
320 initialProviderItemIdentityVersion: providerThread.nativeMetadata.itemIdentityVersion,
321 }),
322 });
323
324 // Rolled-back turns remain in audit history but no longer exist in the provider conversation.
325 const rolledBackAttemptIds = new Set(context.rolledBackAttemptIds);
326 const providerThreadTurns = context.providerTurns.filter(
327 (turn) =>
328 turn.providerThreadId === providerThread.id &&
329 (turn.runAttemptId === null || !rolledBackAttemptIds.has(turn.runAttemptId)),
330 );
331
332 const snapshot =
333 runsToRollback.length === 0
334 ? { providerThread }
335 : yield* session.rollbackThread({
336 providerThread,
337 target: rollbackTarget,
338 providerThreadTurns,
339 });
340 if (input.restoreFiles !== false) yield* checkpoints.restore({ scope, checkpoint });
341 if (staleCheckpoints.length > 0) {
342 yield* checkpoints.deleteStaleRefs({ scope, checkpoints: staleCheckpoints });
343 }
344
345 // Lock against full-row thread writes that could restore the old native session after rollback.
346 yield* threadCommands.withLock(
347 input.threadId,
348 Effect.gen(function* () {
349 const now = yield* DateTime.now;
350 const makeEvent = <Event extends OrchestrationV2DomainEvent>(event: Omit<Event, "id">) =>
351 Effect.map(
352 ids.allocate.event({ threadId: event.threadId }),
353 (id) =>
354 ({
355 ...event,
356 id,
357 }) as Event,
358 );
359 const events: Array<OrchestrationV2DomainEvent> = [];
360 events.push(
361 yield* makeEvent({
362 type: "provider-thread.updated",
363 threadId: input.threadId,
364 driver: providerThread.driver,
365 providerInstanceId: providerThread.providerInstanceId,
366 occurredAt: now,
367 payload: {
368 ...snapshot.providerThread,
369 lastRunOrdinal: targetOrdinal === 0 ? null : targetOrdinal,
370 updatedAt: now,
371 },
372 }),
373 );
374 for (const staleCheckpoint of staleCheckpoints) {
375 events.push(
376 yield* makeEvent({
377 type: "checkpoint.captured",
378 threadId: input.threadId,
379 ...(staleCheckpoint.runId === null ? {} : { runId: staleCheckpoint.runId }),
380 nodeId: staleCheckpoint.nodeId,
381 providerInstanceId: providerThread.providerInstanceId,
382 occurredAt: now,
383 payload: { ...staleCheckpoint, status: "stale" },
384 }),
385 );
386 }
387 for (const run of runsToRollback) {
388 const rootNode = context.rootNodes.find((candidate) => candidate.id === run.rootNodeId);
389 events.push(
390 yield* makeEvent({
391 type: "run.updated",
392 threadId: input.threadId,
393 runId: run.id,
394 ...(rootNode === undefined ? {} : { nodeId: rootNode.id }),
395 providerInstanceId: run.providerInstanceId,
396 occurredAt: now,
397 payload: { ...run, status: "rolled_back", completedAt: now },
398 }),
399 );
400 if (rootNode !== undefined) {
401 events.push(
402 yield* makeEvent({
403 type: "node.updated",
404 threadId: input.threadId,
405 runId: run.id,
406 nodeId: rootNode.id,
407 providerInstanceId: run.providerInstanceId,
408 occurredAt: now,
409 payload: { ...rootNode, status: "rolled_back", completedAt: now },
410 }),
411 );
412 }
413 }
414 yield* eventSink.write({ events });
415 }),
416 );
417 });
418
419 return CheckpointRollbackServiceV2.of({
420 // The inner execute owns the rollback span; this wrapper selects the workspace lease.
421 execute: Effect.fnUntraced(
422 function* (input) {
423 // Storage cleanup removes worktrees under the same lease.
424 const thread = yield* projections.getThread(input.threadId);
425 yield* thread.worktreePath === null
426 ? execute(input)
427 : withWorkspaceLease(path.resolve(thread.worktreePath), execute(input));
428 },
429 (effect, input) =>
430 effect.pipe(
431 Effect.mapError((cause) =>
432 isCheckpointRollbackExecutionError(cause)
433 ? cause
434 : new CheckpointRollbackExecutionError({
435 reason: "unexpected-failure",
436 threadId: input.threadId,
437 providerThreadId: input.providerThreadId,
438 checkpointId: input.checkpointId,
439 cause,
440 }),
441 ),
442 ),
443 ),
444 });
445 }),
446);