| 1 | import assert from "node:assert/strict"; |
| 2 | import { initialState, reducer, type State } from "../lib/useController"; |
| 3 | import { sessionIdentityStableKey } from "../lib/sessionIdentity"; |
| 4 | |
| 5 | const originalNow = Date.now; |
| 6 | let now = 1_000; |
| 7 | Date.now = () => now; |
| 8 | const done = (s: State) => reducer(s, { type: "event", e: { kind: "turn_done" } }); |
| 9 | const idle = (s: State, backgroundJobs = 0) => reducer(s, { |
| 10 | type: "backend_status", running: false, pendingPrompt: false, |
| 11 | cancellable: false, cancelRequested: false, backgroundJobs, |
| 12 | }); |
| 13 | const metrics = (s: State) => [s.turnDoneAt, s.lastTurnDoneAt, s.lastTurnOutputTokens, |
| 14 | s.lastTurnModelMs, s.lastTurnWaitAccumMs, s.lastTurnOutputEstimated]; |
| 15 | try { |
| 16 | const active: State = { ...initialState, running: true, turnActive: true, |
| 17 | turnStartAt: 1_000, backgroundJobs: 1, turnOutputTokens: 20, turnTokens: 100, |
| 18 | turnModelActiveMs: 2_000, live: { id: "a", text: "x".repeat(16), reasoning: "", reasoningComplete: true } }; |
| 19 | now = 21_000; |
| 20 | const completed = done(active); |
| 21 | assert.equal(completed.turnDoneAt - completed.turnStartAt, 20_000); |
| 22 | assert.equal(completed.lastTurnOutputTokens, 24); |
| 23 | now = 90_000; |
| 24 | const jobsEnded = idle(completed); |
| 25 | assert.equal(jobsEnded.backgroundJobs, 0); |
| 26 | assert.deepEqual(metrics(jobsEnded), metrics(completed), "background completion preserves the final snapshot"); |
| 27 | assert.deepEqual(metrics(done(jobsEnded)), metrics(completed), "duplicate turn_done is idempotent"); |
| 28 | now = 21_000; |
| 29 | const fallback = idle(active, 1); |
| 30 | now = 90_000; |
| 31 | assert.deepEqual(metrics(done(fallback)), metrics(fallback), "late turn_done preserves a backend-settled snapshot"); |
| 32 | |
| 33 | now = 21_000; |
| 34 | const waiting = done({ ...active, promptWaitStartedAt: 11_000, turnModelActiveAt: 10_000 }); |
| 35 | assert.equal(waiting.lastTurnWaitAccumMs, 10_000, "open wait closes at completion"); |
| 36 | assert.equal(waiting.promptWaitStartedAt, undefined); |
| 37 | assert.equal(waiting.lastTurnModelMs, 13_000, "model activity closes at the same timestamp"); |
| 38 | assert.deepEqual(metrics(idle(waiting)), metrics(waiting)); |
| 39 | |
| 40 | const planGate = done({ ...active, promptWaitStartedAt: 11_000, |
| 41 | approval: { id: "plan", tool: "exit_plan_mode", subject: "Plan ready" } }); |
| 42 | assert.equal(planGate.running, true, "plan approval remains actionable after turn_done"); |
| 43 | assert.equal(planGate.promptWaitStartedAt, 21_000); |
| 44 | now = 90_000; |
| 45 | const gateClosed = idle(planGate); |
| 46 | assert.deepEqual(metrics(gateClosed), metrics(planGate), "later plan approval wait cannot rewrite the settled snapshot"); |
| 47 | const lateUsage = reducer(completed, { type: "event", e: { kind: "usage", usage: { |
| 48 | promptTokens: 100, completionTokens: 30, totalTokens: 130, cacheHitTokens: 0, |
| 49 | cacheMissTokens: 100, sessionCacheHitTokens: 0, sessionCacheMissTokens: 100, |
| 50 | } } }); |
| 51 | assert.equal(lateUsage.turnDoneAt, completed.turnDoneAt, "late usage cannot change the completion timestamp"); |
| 52 | |
| 53 | now = 100_000; |
| 54 | for (const next of [ |
| 55 | reducer(completed, { type: "user", text: "next", seq: 1, submissionId: "next-turn" }), |
| 56 | reducer(completed, { type: "event", e: { kind: "turn_started", turnStartedAt: now } }), |
| 57 | reducer(completed, { type: "backend_status", running: true, cancellable: true, turnStartedAt: now }), |
| 58 | ]) { |
| 59 | assert.equal(next.turnDoneAt, 0, "new turn clears completion through every entrypoint"); |
| 60 | assert.equal(next.turnOutputTokens, 0); |
| 61 | now += 1_000; |
| 62 | assert.notEqual(done(next).turnDoneAt, completed.turnDoneAt); |
| 63 | } |
| 64 | const retry = reducer(active, { type: "event", e: { kind: "retrying", retryAttempt: 1, retryMax: 3 } }); |
| 65 | assert.equal(retry.turnStartAt, active.turnStartAt); |
| 66 | assert.equal(retry.turnOutputTokens, 20, "same-turn retry retains counters"); |
| 67 | assert.deepEqual(metrics(completed), [21_000, 21_000, 24, 2_000, 0, true], "other tab transitions cannot mutate a completed state"); |
| 68 | |
| 69 | // An MCP interaction is a user wait like an approval or an ask. Clearing the |
| 70 | // approval while one is outstanding must not close the shared interval, or the |
| 71 | // MCP wait never reaches turnWaitAccumMs and the turn clock overcounts it. |
| 72 | now = 5_000; |
| 73 | const bothPrompts: State = { ...initialState, running: true, turnActive: true, |
| 74 | turnStartAt: 1_000, promptWaitStartedAt: 1_000, |
| 75 | meta: { label: "", ready: true, eventChannel: "agent:event", cwd: "", session: { hostId: "local", sessionId: "session-a" }, sessionGeneration: 1 }, |
| 76 | approval: { id: "a1", tool: "write_file", subject: "Run command" }, |
| 77 | mcpInteraction: { id: "m1", server: "srv", mode: "form", message: "Fill the form" } }; |
| 78 | const approvalCleared = reducer(bothPrompts, { type: "clearApproval" }); |
| 79 | assert.equal(approvalCleared.promptWaitStartedAt, 1_000, |
| 80 | "clearing an approval leaves a concurrent MCP wait open"); |
| 81 | assert.equal(approvalCleared.pendingPrompt, true, |
| 82 | "an outstanding MCP interaction still counts as a pending prompt"); |
| 83 | now = 9_000; |
| 84 | const mcpAnswered = reducer(approvalCleared, { |
| 85 | type: "expire_prompt", target: { tabId: "tab-a", sessionKey: sessionIdentityStableKey(bothPrompts.meta), hostId: "local", sessionId: "session-a", sessionGeneration: 1, |
| 86 | promptId: "m1", kind: "mcp", instanceKey: "mcp-1" }, epoch: approvalCleared.promptEpoch }); |
| 87 | assert.equal(mcpAnswered.promptWaitStartedAt, undefined, "the last prompt closes the interval"); |
| 88 | assert.equal(mcpAnswered.turnWaitAccumMs, 8_000, "the whole MCP wait is charged exactly once"); |
| 89 | |
| 90 | now = 5_000; |
| 91 | const drain: State = { ...bothPrompts }; |
| 92 | const drained = reducer(drain, { type: "approval_drained", ids: ["a1"], epoch: drain.promptEpoch }); |
| 93 | assert.equal(drained.promptWaitStartedAt, 1_000, |
| 94 | "an auto-drained approval also leaves a concurrent MCP wait open"); |
| 95 | console.log("completed-turn telemetry: all assertions passed"); |
| 96 | } finally { |
| 97 | Date.now = originalNow; |
| 98 | } |
| 99 |