1import { assert } from "@effect/vitest";
2import type { ProviderReplayTranscript } from "@t3tools/contracts";
3
4import type { OrchestratorV2ScenarioResult } from "../../OrchestratorScenario.ts";
5import { projectionFor } from "../shared.ts";
6import { assertSimpleOutput } from "./codex_output.ts";
7
8function isRecord(value: unknown): value is Record<string, unknown> {
9 return typeof value === "object" && value !== null && !Array.isArray(value);
10}
11
12function field(value: unknown, key: string): unknown {
13 return isRecord(value) ? value[key] : undefined;
14}
15
16/** Data of the last recorded successful response to `command`. */
17function lastPiResponseData(transcript: ProviderReplayTranscript, command: string): unknown {
18 const frame = transcript.entries
19 .flatMap((entry) => (entry.type === "emit_inbound" ? [entry.frame] : []))
20 .findLast(
21 (candidate) =>
22 field(candidate, "type") === "response" && field(candidate, "command") === command,
23 );
24 assert.isDefined(frame, `transcript must record a ${command} response`);
25 return field(frame, "data");
26}
27
28/**
29 * Pi reports usage only through `get_session_stats`, read once per settled
30 * turn: each completed provider turn carries the stats response recorded for
31 * it, in order, not the last live streaming total.
32 */
33export function assertPiSettledTokenUsage(
34 result: OrchestratorV2ScenarioResult,
35 transcript: ProviderReplayTranscript,
36) {
37 const stats = transcript.entries.flatMap((entry) =>
38 entry.type === "emit_inbound" &&
39 field(entry.frame, "type") === "response" &&
40 field(entry.frame, "command") === "get_session_stats"
41 ? [field(entry.frame, "data")]
42 : [],
43 );
44 const turns = projectionFor(result, transcript.scenario).providerTurns;
45 assert.lengthOf(turns, stats.length);
46 for (const [index, turn] of turns.entries()) {
47 const contextUsage = field(stats[index], "contextUsage");
48 const tokens = field(stats[index], "tokens");
49 const { updatedAt: _updatedAt, ...tokenUsage } = turn.tokenUsage ?? { updatedAt: "" };
50 assert.deepEqual(tokenUsage, {
51 usedTokens: field(contextUsage, "tokens"),
52 maxTokens: field(contextUsage, "contextWindow"),
53 inputTokens: field(tokens, "input"),
54 cachedInputTokens: field(tokens, "cacheRead"),
55 outputTokens: field(tokens, "output"),
56 });
57 }
58}
59
60/**
61 * Pi-specific additions to the shared simple contract: the streamed thinking
62 * block becomes a reasoning item, the settled turn carries the context usage
63 * from `get_session_stats`, and the turn's native ref is the session-tree id
64 * of its user message (the point rollback and fork re-root at).
65 */
66export function assertPiSimpleOutput(
67 result: OrchestratorV2ScenarioResult,
68 transcript: ProviderReplayTranscript,
69) {
70 assertSimpleOutput(result, transcript);
71 const projection = projectionFor(result, transcript.scenario);
72
73 const recordedThinking = transcript.entries.flatMap((entry) => {
74 const update = field(
75 entry.type === "emit_inbound" ? entry.frame : undefined,
76 "assistantMessageEvent",
77 );
78 return field(update, "type") === "thinking_end" ? [field(update, "content")] : [];
79 });
80 assert.lengthOf(recordedThinking, 1, "the recording must stream one thinking block");
81 assert.deepEqual(
82 projection.turnItems.flatMap((item) => (item.type === "reasoning" ? [item.text] : [])),
83 recordedThinking,
84 "Pi thinking deltas must project as one reasoning item with the recorded text",
85 );
86
87 assertPiSettledTokenUsage(result, transcript);
88 const [turn] = projection.providerTurns;
89
90 // Pi attaches cumulative usage to every streaming update, zeros until the
91 // provider reports. The running turn moves the meter once per new total.
92 const streamedTotals = transcript.entries.flatMap((entry) =>
93 entry.type === "emit_inbound" && field(entry.frame, "type") === "message_update"
94 ? [field(field(entry.frame, "usage"), "totalTokens")]
95 : [],
96 );
97 const liveUsage = result.domainEvents.flatMap((event) =>
98 event.type === "provider-turn.updated" &&
99 event.payload.status === "running" &&
100 event.payload.tokenUsage !== undefined
101 ? [event.payload.tokenUsage.usedTokens]
102 : [],
103 );
104 assert.deepEqual(liveUsage, [...new Set(streamedTotals.filter((total) => total !== 0))]);
105
106 const entries = field(lastPiResponseData(transcript, "get_entries"), "entries");
107 const userEntryId = (Array.isArray(entries) ? entries : [])
108 .filter(
109 (entry) =>
110 field(entry, "type") === "message" && field(field(entry, "message"), "role") === "user",
111 )
112 .map((entry) => field(entry, "id"))
113 .at(0);
114 assert.isString(userEntryId);
115 assert.equal(turn?.nativeTurnRef?.nativeId, userEntryId);
116 assert.equal(turn?.nativeTurnRef?.strength, "strong");
117}