fix(server): Pi rollback works past a turn you stopped

apps/server/src/orchestration-v2/Adapters/PiAdapterV2.ts

@@ -351,6 +351,13 @@ interface ActivePiTurn {
   activeCompaction: PiCompactionState | null;
   activeProviderRetry: PiProviderRetryState | null;
   failure: ReturnType<typeof makeProviderFailure> | null;
+  /** Session-tree refs read just before Stop terminates Pi, when no read is possible later. */
+  stopTreeRefs?: PiTurnTreeRefs | null;
+}
+
+interface PiTurnTreeRefs {
+  readonly turnStartEntryId: string | null;
+  readonly leafId: string | null;
 }
 
 interface PendingPiPrompt {
@@ -1367,13 +1374,18 @@ export function makePiAdapterV2(options: PiAdapterV2Options): ProviderAdapterV2S
        * point `fork` rolls back to); the leaf becomes the conversation head.
        * Pure bookkeeping: failures degrade to the synthetic refs.
        */
-      const captureTurnTreeRefs = Effect.fnUntraced(function* () {
+      const captureTurnTreeRefs = Effect.fnUntraced(function* (
+        timeoutMs = PI_REQUEST_TIMEOUT_MS,
+      ): Effect.fn.Return<PiTurnTreeRefs | null> {
         const cursorWasStale = leafCursorStale;
         const cursor = cursorWasStale ? null : lastKnownLeaf;
-        const data = yield* request({
-          type: "get_entries",
-          ...(cursor === null ? {} : { since: cursor }),
-        }).pipe(Effect.orElseSucceed(() => undefined));
+        const data = yield* request(
+          {
+            type: "get_entries",
+            ...(cursor === null ? {} : { since: cursor }),
+          },
+          timeoutMs,
+        ).pipe(Effect.orElseSucceed(() => undefined));
         if (data === undefined) {
           // Pi may have advanced past `lastKnownLeaf` while this failed, so the
           // cursor can no longer be trusted to bound a single turn.
@@ -1431,7 +1443,8 @@ export function makePiAdapterV2(options: PiAdapterV2Options): ProviderAdapterV2S
           }
         }
         yield* cancelPendingPrompts(completedAt);
-        const treeRefs = yield* captureTurnTreeRefs();
+        const treeRefs =
+          turn.stopTreeRefs !== undefined ? turn.stopTreeRefs : yield* captureTurnTreeRefs();
         const tokenUsage = readUsage
           ? yield* readTokenUsage(turn.latestCompactionAfterTokens, completedAt)
           : undefined;
@@ -2490,7 +2503,18 @@ export function makePiAdapterV2(options: PiAdapterV2Options): ProviderAdapterV2S
               if (interruptInput.requestRuntimeRestart === true && !turn.settleWhenIdle) {
                 yield* request({ type: "abort" }, 2_000).pipe(Effect.ignore);
               }
-              yield* connection.terminate;
+              // Terminating fails every later request, so read the stopped
+              // turn's session-tree refs first: rolling back past this turn
+              // forks at its user entry. Holding the event permit also lets a
+              // finalize that is already reading them finish before the kill.
+              yield* sessionEventPermit.withPermits(1)(
+                Effect.gen(function* () {
+                  if (threadState?.activeTurn === turn && turn.stopTreeRefs === undefined) {
+                    turn.stopTreeRefs = yield* captureTurnTreeRefs(2_000);
+                  }
+                  yield* connection.terminate;
+                }),
+              );
               return;
             }
             yield* request({ type: "abort" }).pipe(

apps/server/src/orchestration-v2/testkit/fixtures/index.ts

@@ -87,6 +87,8 @@ import { assertClaudeThreadRollbackOutput } from "./thread_rollback/claude_outpu
 import { assertThreadRollbackOutput } from "./thread_rollback/codex_output.ts";
 import { threadRollbackInput } from "./thread_rollback/input.ts";
 import { assertPiThreadRollbackOutput } from "./thread_rollback/pi_output.ts";
+import { threadRollbackAfterStopInput } from "./thread_rollback_after_stop/input.ts";
+import { assertPiThreadRollbackAfterStopOutput } from "./thread_rollback_after_stop/pi_output.ts";
 import { assertTodoListOutput } from "./todo_list/codex_output.ts";
 import { assertTodoListCursorOutput } from "./todo_list/cursor_output.ts";
 import { assertTodoListGrokOutput } from "./todo_list/grok_output.ts";
@@ -1109,6 +1111,21 @@ export const ORCHESTRATOR_REPLAY_FIXTURES: ReadonlyArray<OrchestratorReplayFixtu
       },
     ],
   },
+  {
+    name: "thread_rollback_after_stop",
+    buildInput: threadRollbackAfterStopInput,
+    providers: [
+      {
+        driver: ProviderDriverKind.make("pi"),
+        transcriptFile: new URL(
+          "./thread_rollback_after_stop/pi_transcript.ndjson",
+          import.meta.url,
+        ),
+        modelSelection: PI_MODEL_SELECTION,
+        assertOutput: assertPiThreadRollbackAfterStopOutput,
+      },
+    ],
+  },
 ];
 
 // TODO(claude-v2/context-transfer): add provider-switch handoff and return fixtures when portable

apps/server/src/orchestration-v2/testkit/fixtures/thread_rollback_after_stop/input.ts

@@ -0,0 +1,29 @@
+import {
+  THREAD_ROLLBACK_AFTER_PROMPT,
+  THREAD_ROLLBACK_FIRST_PROMPT,
+  THREAD_ROLLBACK_SECOND_PROMPT,
+  TURN_INTERRUPT_MID_TOOL_PROMPT,
+  type OrchestratorFixtureInput,
+} from "../shared.ts";
+
+/**
+ * Rolls back past a turn the user stopped mid-tool. Stop restarts the
+ * provider runtime, so the stopped turn's native rollback point must be
+ * captured before the process goes away.
+ */
+export function threadRollbackAfterStopInput(): OrchestratorFixtureInput {
+  return {
+    steps: [
+      { type: "message", text: THREAD_ROLLBACK_FIRST_PROMPT },
+      { type: "message", text: TURN_INTERRUPT_MID_TOOL_PROMPT },
+      { type: "interrupt", targetRunIndex: 2, waitForTurnItemType: "command_execution" },
+      { type: "message", text: THREAD_ROLLBACK_SECOND_PROMPT },
+      {
+        type: "rollback",
+        checkpointScopeSuffix: "root",
+        checkpointSuffix: "1",
+      },
+      { type: "message", text: THREAD_ROLLBACK_AFTER_PROMPT },
+    ],
+  };
+}

apps/server/src/orchestration-v2/testkit/fixtures/thread_rollback_after_stop/pi_output.ts

@@ -0,0 +1,77 @@
+import { assert } from "@effect/vitest";
+import type { ProviderReplayTranscript } from "@t3tools/contracts";
+
+import type { OrchestratorV2ScenarioResult } from "../../OrchestratorScenario.ts";
+import {
+  assertBaseProjection,
+  assertRunOrdinals,
+  assertSemanticProjectionIntegrity,
+  assertVisibleUserMessagesExclude,
+  assertVisibleUserMessagesInclude,
+  projectionFor,
+  THREAD_ROLLBACK_AFTER_PROMPT,
+  THREAD_ROLLBACK_FIRST_PROMPT,
+  THREAD_ROLLBACK_SECOND_PROMPT,
+} from "../shared.ts";
+
+function field(value: unknown, key: string): unknown {
+  return typeof value === "object" && value !== null ? Reflect.get(value, key) : undefined;
+}
+
+function outboundOfType(transcript: ProviderReplayTranscript, type: string) {
+  return transcript.entries.flatMap((entry) =>
+    entry.type === "expect_outbound" && field(entry.frame, "type") === type ? [entry.frame] : [],
+  );
+}
+
+/**
+ * Stop terminates Pi (Stop-with-restart), so the stopped turn's session-tree
+ * user entry has to be read before the kill. Rolling back to turn 1 then forks
+ * at that entry, which discards the stopped turn and everything after it.
+ */
+export function assertPiThreadRollbackAfterStopOutput(
+  result: OrchestratorV2ScenarioResult,
+  transcript: ProviderReplayTranscript,
+) {
+  assertBaseProjection({ result, transcript, runCount: 4 });
+  const projection = projectionFor(result, transcript.scenario);
+  assertSemanticProjectionIntegrity(projection);
+  assertRunOrdinals(projection, [1, 2, 3, 4]);
+  const runStatus = (ordinal: number) =>
+    projection.runs.find((run) => run.ordinal === ordinal)?.status;
+  assert.equal(runStatus(1), "completed");
+  assert.equal(runStatus(3), "rolled_back");
+  assert.equal(runStatus(4), "completed");
+
+  // The stopped turn must carry the strong ref rollback forks at.
+  const stoppedTurn = projection.providerTurns.find((providerTurn) => providerTurn.ordinal === 2);
+  assert.equal(stoppedTurn?.status, "interrupted");
+  assert.equal(stoppedTurn?.nativeTurnRef?.strength, "strong");
+
+  const stopIndex = transcript.entries.findIndex(
+    (entry) => entry.type === "expect_outbound" && field(entry.frame, "type") === "abort",
+  );
+  const captureIndex = transcript.entries.findIndex(
+    (entry, index) =>
+      index > stopIndex &&
+      entry.type === "expect_outbound" &&
+      field(entry.frame, "type") === "get_entries",
+  );
+  assert.isAbove(captureIndex, stopIndex, "Stop must read the session tree before terminating Pi");
+
+  const forks = outboundOfType(transcript, "fork");
+  assert.lengthOf(forks, 1, "rollback must fork the Pi session tree exactly once");
+  assert.equal(field(forks[0], "entryId"), stoppedTurn?.nativeTurnRef?.nativeId);
+
+  assertVisibleUserMessagesInclude(projection, [
+    THREAD_ROLLBACK_FIRST_PROMPT,
+    THREAD_ROLLBACK_AFTER_PROMPT,
+  ]);
+  assertVisibleUserMessagesExclude(projection, [THREAD_ROLLBACK_SECOND_PROMPT]);
+
+  // Pi's surviving branch holds only turn 1.
+  const finalAnswer = projection.turnItems.findLast((item) => item.type === "assistant_message");
+  assert.include(finalAnswer?.text, "rollback fixture first turn complete");
+  assert.notInclude(finalAnswer?.text, "second turn");
+  assert.notInclude(finalAnswer?.text, "interrupt fixture");
+}