diff --git a/apps/server/src/orchestration-v2/Orchestrator.ts b/apps/server/src/orchestration-v2/Orchestrator.ts index 20e65e4b9bed1..276875c60fe89 100644 --- a/apps/server/src/orchestration-v2/Orchestrator.ts +++ b/apps/server/src/orchestration-v2/Orchestrator.ts @@ -12550,7 +12550,6 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio emit, loadProjectionForCommand, hasLiveRun, - providerTurnForRun, errors: { OrchestratorDispatchError, OrchestratorProjectionError, diff --git a/apps/server/src/orchestration-v2/ThreadRewind.test.ts b/apps/server/src/orchestration-v2/ThreadRewind.test.ts index 7abb404306709..bacb28410251b 100644 --- a/apps/server/src/orchestration-v2/ThreadRewind.test.ts +++ b/apps/server/src/orchestration-v2/ThreadRewind.test.ts @@ -1,17 +1,24 @@ import { assert, it } from "@effect/vitest"; import { CommandId, + EventId, MessageId, ProjectId, ProviderDriverKind, ProviderInstanceId, + ProviderThreadId, + ProviderTurnId, ThreadId, } from "@t3tools/contracts"; +import * as DateTime from "effect/DateTime"; import * as Effect from "effect/Effect"; import * as Layer from "effect/Layer"; +import * as Stream from "effect/Stream"; import * as SqlitePersistence from "../persistence/Sqlite.ts"; import { ClaudeProviderCapabilitiesV2 } from "./Adapters/ClaudeAdapterV2.ts"; +import * as EffectWorker from "./EffectWorker.ts"; +import * as EventSink from "./EventSink.ts"; import * as Orchestrator from "./Orchestrator.ts"; import * as ProjectionStore from "./ProjectionStore.ts"; import type * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter"; @@ -115,3 +122,180 @@ it.effect("refuses to rewind while a message waits in the queue, even a held one assert.equal(yield* orchestrator.getThreadEventSequence(threadId), before); }).pipe(Effect.provide(layerTest)), ); + +/** + * Sends `text`, then settles its run as the provider would have: `completed` + * with a Claude turn on `providerThreadId`, or `cancelled` with none (the + * user stopped it before Claude answered). + */ +const sendAndSettle = ( + threadId: ThreadId, + text: string, + outcome: + | { readonly status: "completed"; readonly turnOn?: string } + | { readonly status: "cancelled" }, +) => + Effect.gen(function* () { + const orchestrator = yield* Orchestrator.OrchestratorV2; + const eventSink = yield* EventSink.EventSinkV2; + yield* send(threadId, text, "start_immediately"); + const projection = yield* orchestrator.getThreadProjection(threadId); + const run = projection.runs.find( + (candidate) => candidate.userMessageId === messageId(threadId, text), + )!; + const attempt = projection.attempts.find((candidate) => candidate.id === run.activeAttemptId)!; + const now = yield* DateTime.now; + const turn = + outcome.status === "completed" + ? [ + { + id: EventId.make(`event:${run.id}:turn`), + type: "provider-turn.updated" as const, + threadId, + runId: run.id, + occurredAt: now, + payload: { + id: ProviderTurnId.make(`provider-turn:${run.id}`), + providerThreadId: ProviderThreadId.make(outcome.turnOn ?? attempt.providerThreadId), + nodeId: attempt.rootNodeId, + runAttemptId: attempt.id, + nativeTurnRef: null, + ordinal: run.ordinal, + status: "completed" as const, + startedAt: now, + completedAt: now, + }, + }, + ] + : []; + yield* eventSink.write({ + commandId: CommandId.make(`settle:${run.id}`), + events: [ + ...turn, + { + id: EventId.make(`event:${run.id}:done`), + type: "run.updated", + threadId, + runId: run.id, + occurredAt: now, + payload: { ...run, status: outcome.status, completedAt: now }, + }, + ], + }); + }); + +/** A Claude adapter whose session records each rollback instead of reaching a provider. */ +const makeRollbackRecordingLayer = ( + rollbacks: Array, +) => { + const recordingAdapter = { + ...adapter, + openSession: (sessionInput: { readonly providerSessionId: never }) => + Effect.map( + DateTime.now, + (now) => + ({ + instanceId, + driver: adapter.driver, + providerSessionId: sessionInput.providerSessionId, + providerSession: { + id: sessionInput.providerSessionId, + driver: adapter.driver, + providerInstanceId: instanceId, + status: "ready", + cwd: "/fallback", + model: modelSelection.model, + capabilities: ClaudeProviderCapabilitiesV2, + createdAt: now, + updatedAt: now, + lastError: null, + }, + events: Stream.empty, + rollbackThread: (input: ProviderAdapter.ProviderAdapterV2RollbackThreadInput) => + Effect.sync(() => { + rollbacks.push(input.target); + return { providerThread: input.providerThread }; + }), + }) as never, + ), + } as unknown as ProviderAdapter.ProviderAdapterV2["Service"]; + return Layer.mergeAll( + layerDatabase, + ProjectionStore.layer.pipe(Layer.provide(layerDatabase)), + ProviderReplayHarness.layerWithRegistry( + { name: "thread-rewind-rollback" }, + ProviderAdapterRegistry.layerFromAdapters([recordingAdapter]), + { databaseLayer: layerDatabase, runEffectWorker: false }, + ), + ); +}; + +it.effect( + "rewinds past cancelled turns Claude never answered, back to the last answered turn", + () => { + const rollbacks: Array = []; + return Effect.gen(function* () { + const orchestrator = yield* Orchestrator.OrchestratorV2; + const worker = yield* EffectWorker.OrchestrationEffectWorkerV2; + const threadId = ThreadId.make("thread:rewind-skips-unanswered"); + yield* createThread(threadId); + yield* sendAndSettle(threadId, "answered", { status: "completed" }); + yield* sendAndSettle(threadId, "stopped-1", { status: "cancelled" }); + yield* sendAndSettle(threadId, "stopped-2", { status: "cancelled" }); + yield* sendAndSettle(threadId, "last", { status: "completed" }); + + const commandId = CommandId.make("rewind-skips-unanswered"); + yield* orchestrator.dispatch({ + type: "thread.rewind", + commandId, + threadId, + messageId: messageId(threadId, "last"), + choice: "conversation", + }); + yield* worker.drain(); + + const projection = yield* orchestrator.getThreadProjection(threadId); + const statusOf = (text: string) => + projection.runs.find((run) => run.userMessageId === messageId(threadId, text))!; + const answered = statusOf("answered"); + // The provider rolled back to the answered run, not to a cancelled one. + assert.deepEqual( + rollbacks.map((target) => [target.type, target.appRunOrdinal]), + [["provider_turn", answered.ordinal]], + ); + assert.equal(projection.providerThreads[0]?.lastRunOrdinal, answered.ordinal); + assert.deepEqual( + ["answered", "stopped-1", "stopped-2", "last"].map((text) => statusOf(text).status), + ["completed", "rolled_back", "rolled_back", "rolled_back"], + ); + }).pipe(Effect.provide(makeRollbackRecordingLayer(rollbacks))); + }, +); + +it.effect("still refuses when the turn before was answered in a different Claude session", () => + Effect.gen(function* () { + const orchestrator = yield* Orchestrator.OrchestratorV2; + const threadId = ThreadId.make("thread:rewind-other-session"); + yield* createThread(threadId); + yield* sendAndSettle(threadId, "answered-elsewhere", { + status: "completed", + turnOn: "provider-thread:another-session", + }); + yield* sendAndSettle(threadId, "last", { status: "completed" }); + + const refusal = yield* orchestrator + .dispatch({ + type: "thread.rewind", + commandId: CommandId.make("rewind-other-session"), + threadId, + messageId: messageId(threadId, "last"), + choice: "conversation", + }) + .pipe(Effect.flip, Effect.map(userFacingDispatchErrorMessage)); + + assert.equal( + refusal, + "Cannot rewind to this message: the turn before it is not part of the current Claude session.", + ); + }).pipe(Effect.provide(layerTest)), +); diff --git a/apps/server/src/orchestration-v2/ThreadRewindCommands.ts b/apps/server/src/orchestration-v2/ThreadRewindCommands.ts index 15cd214f7f003..78bba80c102df 100644 --- a/apps/server/src/orchestration-v2/ThreadRewindCommands.ts +++ b/apps/server/src/orchestration-v2/ThreadRewindCommands.ts @@ -10,8 +10,6 @@ import { type OrchestrationV2DomainEvent, type OrchestrationV2InternalCommand, type OrchestrationV2ProviderThread, - type OrchestrationV2ProviderTurn, - type OrchestrationV2Run, type OrchestrationV2ServerCommand, type OrchestrationV2ThreadProjection, type ProviderInstanceId, @@ -41,6 +39,7 @@ import type { ProviderAdapterRegistryV2 } from "./ProviderAdapterRegistry.ts"; import { NO_FILE_SNAPSHOTS_MESSAGE, previousConversationRun, + providerTurnForRun, runFileCheckpointTurn, } from "./ThreadRewindTargets.ts"; @@ -90,10 +89,6 @@ export interface ThreadRewindCommandDeps { filter?: ProjectionRecordFilter, ) => Effect.Effect, OrchestratorProjectionError>; readonly hasLiveRun: (projection: Pick) => boolean; - readonly providerTurnForRun: ( - projection: Pick, - run: OrchestrationV2Run, - ) => OrchestrationV2ProviderTurn | undefined; /** The orchestrator's own error classes, passed in so this file never imports it at run time. */ readonly errors: { readonly OrchestratorDispatchError: typeof OrchestratorDispatchError; @@ -111,7 +106,6 @@ export function makeThreadRewindCommands(deps: ThreadRewindCommandDeps) { emit, loadProjectionForCommand, hasLiveRun, - providerTurnForRun, errors: { OrchestratorDispatchError, OrchestratorProjectionError, @@ -271,7 +265,7 @@ export function makeThreadRewindCommands(deps: ThreadRewindCommandDeps) { if (restores.code && runFileCheckpointTurn(projection, run.id) === undefined) { return yield* refuse(NO_FILE_SNAPSHOTS_MESSAGE); } - const previousRun = previousConversationRun(projection.runs, run); + const previousRun = previousConversationRun(projection, run); const previousTurn = previousRun === undefined ? undefined : providerTurnForRun(projection, previousRun); if ( diff --git a/apps/server/src/orchestration-v2/ThreadRewindService.ts b/apps/server/src/orchestration-v2/ThreadRewindService.ts index c6e372158e2a7..16a2844865747 100644 --- a/apps/server/src/orchestration-v2/ThreadRewindService.ts +++ b/apps/server/src/orchestration-v2/ThreadRewindService.ts @@ -175,7 +175,7 @@ export const layer: Layer.Layer< if (run.status === "rolled_back") { return { type: "completed" as const, notice: ALREADY_REWOUND_MESSAGE }; } - const targetOrdinal = previousConversationRun(projection.runs, run)?.ordinal ?? 0; + const targetOrdinal = previousConversationRun(projection, run)?.ordinal ?? 0; const rewindConversation = checkpointRollback.rollbackConversation({ threadId: input.threadId, source: rollbackSourceFromProjection(projection, { diff --git a/apps/server/src/orchestration-v2/ThreadRewindTargets.ts b/apps/server/src/orchestration-v2/ThreadRewindTargets.ts index 986fe2ba8f7bf..e3d1de87202a4 100644 --- a/apps/server/src/orchestration-v2/ThreadRewindTargets.ts +++ b/apps/server/src/orchestration-v2/ThreadRewindTargets.ts @@ -1,4 +1,5 @@ import { + latestProviderTurnForAttempt, type MessageId, type OrchestrationV2ProviderTurn, type OrchestrationV2Run, @@ -51,18 +52,44 @@ export function compareConversationOrder( } /** - * The last run before `run` still in the conversation: rewinding to the - * message of `run` keeps the conversation up to its end. Undefined means the - * rewind goes back to the thread start. + * The provider turn Claude took for `run`; undefined when the user stopped it + * first. The previous-run skip and the rewind refusal both read this, so they + * cannot disagree. + */ +export function providerTurnForRun( + projection: Pick, "attempts" | "providerTurns">, + run: OrchestrationV2Run, +): OrchestrationV2ProviderTurn | undefined { + if (run.activeAttemptId === null) return undefined; + return ( + latestProviderTurnForAttempt(projection.providerTurns, run.activeAttemptId) ?? + projection.providerTurns.find( + (turn) => + projection.attempts.find((attempt) => attempt.id === run.activeAttemptId) + ?.providerTurnId === turn.id, + ) + ); +} + +/** + * The last run before `run` still in the conversation and answered by Claude: + * rewinding to the message of `run` keeps the conversation up to its end. A + * run stopped before Claude answered is skipped, since it left no turn to + * resume at. Undefined means the rewind goes back to the thread start. */ export function previousConversationRun( - runs: ReadonlyArray, + projection: Pick< + ProjectionRecords<"runs" | "attempts" | "providerTurns">, + "runs" | "attempts" | "providerTurns" + >, run: OrchestrationV2Run, ): OrchestrationV2Run | undefined { - return runs + return projection.runs .filter( (candidate) => - compareConversationOrder(candidate, run) < 0 && candidate.status !== "rolled_back", + compareConversationOrder(candidate, run) < 0 && + candidate.status !== "rolled_back" && + providerTurnForRun(projection, candidate) !== undefined, ) .toSorted(compareConversationOrder) .at(-1); diff --git a/apps/server/src/project/AgentSessionImporter.test.ts b/apps/server/src/project/AgentSessionImporter.test.ts index 1e02a28f72bcb..426b450055a02 100644 --- a/apps/server/src/project/AgentSessionImporter.test.ts +++ b/apps/server/src/project/AgentSessionImporter.test.ts @@ -863,7 +863,7 @@ const rewindPoints = Effect.fn("rewindPoints")(function* (threadId: ThreadId) { .map((message) => { const run = records.runs.find((candidate) => candidate.userMessageId === message.id); if (run === undefined) return { prompt: message.text, rewind: "none" as const }; - const previous = previousConversationRun(records.runs, run); + const previous = previousConversationRun(records, run); return { prompt: message.text, resumeAfter: