1import { assert } from "@effect/vitest";
2import type { ProviderReplayTranscript } from "@t3tools/contracts";
3
4import type { OrchestratorV2ScenarioResult } from "../../OrchestratorScenario.ts";
5import {
6 assertBaseProjection,
7 assertRunOrdinals,
8 assertSemanticProjectionIntegrity,
9 assertTurnItemTypeSequence,
10 assertUserMessagesInclude,
11 assertVisibleTurnItemTypeSequence,
12 assertVisibleUserMessagesExclude,
13 assertVisibleUserMessagesInclude,
14 projectionFor,
15 THREAD_ROLLBACK_AFTER_PROMPT,
16 THREAD_ROLLBACK_FIRST_PROMPT,
17 THREAD_ROLLBACK_SECOND_PROMPT,
18} from "../shared.ts";
19
20function field(value: unknown, key: string): unknown {
21 return typeof value === "object" && value !== null ? Reflect.get(value, key) : undefined;
22}
23
24function outboundOfType(transcript: ProviderReplayTranscript, type: string) {
25 return transcript.entries.flatMap((entry) =>
26 entry.type === "expect_outbound" && field(entry.frame, "type") === type ? [entry.frame] : [],
27 );
28}
29
30/**
31 * Pi rolls back by re-rooting its session tree with `fork(entryId)` at the
32 * first user entry of the earliest discarded turn. That id is the discarded
33 * turn's strong native ref, captured from `get_entries` when the turn settled.
34 * The fork writes a new session file, and the next turn resumes that file.
35 */
36export function assertPiThreadRollbackOutput(
37 result: OrchestratorV2ScenarioResult,
38 transcript: ProviderReplayTranscript,
39) {
40 assertBaseProjection({
41 result,
42 transcript,
43 runCount: 3,
44 runStatuses: ["completed", "rolled_back", "completed"],
45 });
46 const projection = projectionFor(result, transcript.scenario);
47 assertSemanticProjectionIntegrity(projection);
48 assertRunOrdinals(projection, [1, 2, 3]);
49 const turn = ["user_message", "reasoning", "assistant_message", "checkpoint"] as const;
50 assertTurnItemTypeSequence(projection, [...turn, ...turn, ...turn]);
51 assertVisibleTurnItemTypeSequence(projection, [...turn, ...turn]);
52 assertUserMessagesInclude(projection, [
53 THREAD_ROLLBACK_FIRST_PROMPT,
54 THREAD_ROLLBACK_SECOND_PROMPT,
55 THREAD_ROLLBACK_AFTER_PROMPT,
56 ]);
57 assertVisibleUserMessagesInclude(projection, [
58 THREAD_ROLLBACK_FIRST_PROMPT,
59 THREAD_ROLLBACK_AFTER_PROMPT,
60 ]);
61 assertVisibleUserMessagesExclude(projection, [THREAD_ROLLBACK_SECOND_PROMPT]);
62
63 const forks = outboundOfType(transcript, "fork");
64 assert.lengthOf(forks, 1, "rollback must fork the Pi session tree exactly once");
65 const discardedTurn = projection.providerTurns.find((providerTurn) => providerTurn.ordinal === 2);
66 assert.equal(discardedTurn?.nativeTurnRef?.strength, "strong");
67 assert.equal(field(forks[0], "entryId"), discardedTurn?.nativeTurnRef?.nativeId);
68
69 // The fork moved the thread to a new session file; the post-rollback turn
70 // resumes that file, never the pre-rollback one.
71 const [providerThread] = projection.providerThreads;
72 assert.lengthOf(projection.providerThreads, 1);
73 assert.equal(providerThread?.nativeThreadRef?.nativeId, "/pi-sessions/session-2.jsonl");
74 assert.deepEqual(
75 outboundOfType(transcript, "switch_session").map((frame) => field(frame, "sessionPath")),
76 ["/pi-sessions/session-2.jsonl"],
77 );
78
79 // The model has only the surviving branch: the discarded exchange is gone.
80 const finalAnswer = projection.turnItems.findLast((item) => item.type === "assistant_message");
81 assert.include(finalAnswer?.text, "rollback fixture first turn complete");
82 assert.notInclude(finalAnswer?.text, "second turn");
83}