fix: restore console run accounting from persisted logs

2026-10-09 02:27:27 +09:00
parent e5ea7544d4
commit f1287f27e7
19 changed files with 1023 additions and 39 deletions
@@ -193,7 +193,7 @@ export type RunResumeSource = "compaction" | "pause";

 export type RunFailureKind = "engine" | "compaction";

-export type SessionSnapshotEntryData = { "kind": "user_input", segments: Array<Segment>, } | { "kind": "message", role: SessionMessageRole, content: Array<SessionContentPart>, } | { "kind": "tool_call", call_id: string, name: string, arguments: string, } | { "kind": "tool_result", call_id: string, summary: string, content?: string | null, is_error: boolean, attachments?: Array<SessionToolAttachment>, } | { "kind": "system_item", item_kind: string, content: string, data?: unknown, } | { "kind": "run_yielded", reason: RunYieldReason, active_run_turn_count: number, } | { "kind": "run_resumed", source: RunResumeSource, active_run_turn_count: number, } | { "kind": "run_cancelled" } | { "kind": "run_error", message: string, failure?: RunFailureKind | null, };
+export type SessionSnapshotEntryData = { "kind": "user_input", segments: Array<Segment>, } | { "kind": "message", role: SessionMessageRole, content: Array<SessionContentPart>, } | { "kind": "tool_call", call_id: string, name: string, arguments: string, } | { "kind": "tool_result", call_id: string, summary: string, content?: string | null, is_error: boolean, attachments?: Array<SessionToolAttachment>, } | { "kind": "system_item", item_kind: string, content: string, data?: unknown, } | { "kind": "run_yielded", reason: RunYieldReason, active_run_turn_count: number, } | { "kind": "run_resumed", source: RunResumeSource, active_run_turn_count: number, } | { "kind": "invoke", trigger: InvokeKind, } | { "kind": "usage", input_tokens: number, cache_read_input_tokens: number, output_tokens: number, } | { "kind": "run_stats", elapsed_ms: number, requests: number, upload_tokens: number, output_tokens: number, } | { "kind": "run_cancelled" } | { "kind": "run_error", message: string, failure?: RunFailureKind | null, };

 export type SessionSnapshotEntry = {
 /**
@@ -204,7 +204,7 @@ entry_id: string,
 /**
 * Timestamp copied from the durable log record that commits this entry.
 */
-timestamp: number, provenance: SessionEntryProvenance, derived_from?: Array<string>, } & ({ "kind": "user_input", segments: Array<Segment>, } | { "kind": "message", role: SessionMessageRole, content: Array<SessionContentPart>, } | { "kind": "tool_call", call_id: string, name: string, arguments: string, } | { "kind": "tool_result", call_id: string, summary: string, content?: string | null, is_error: boolean, attachments?: Array<SessionToolAttachment>, } | { "kind": "system_item", item_kind: string, content: string, data?: unknown, } | { "kind": "run_yielded", reason: RunYieldReason, active_run_turn_count: number, } | { "kind": "run_resumed", source: RunResumeSource, active_run_turn_count: number, } | { "kind": "run_cancelled" } | { "kind": "run_error", message: string, failure?: RunFailureKind | null, });
+timestamp: number, provenance: SessionEntryProvenance, derived_from?: Array<string>, } & ({ "kind": "user_input", segments: Array<Segment>, } | { "kind": "message", role: SessionMessageRole, content: Array<SessionContentPart>, } | { "kind": "tool_call", call_id: string, name: string, arguments: string, } | { "kind": "tool_result", call_id: string, summary: string, content?: string | null, is_error: boolean, attachments?: Array<SessionToolAttachment>, } | { "kind": "system_item", item_kind: string, content: string, data?: unknown, } | { "kind": "run_yielded", reason: RunYieldReason, active_run_turn_count: number, } | { "kind": "run_resumed", source: RunResumeSource, active_run_turn_count: number, } | { "kind": "invoke", trigger: InvokeKind, } | { "kind": "usage", input_tokens: number, cache_read_input_tokens: number, output_tokens: number, } | { "kind": "run_stats", elapsed_ms: number, requests: number, upload_tokens: number, output_tokens: number, } | { "kind": "run_cancelled" } | { "kind": "run_error", message: string, failure?: RunFailureKind | null, });

 export type SessionConversationTurn = { turn_id: string, entries: Array<SessionSnapshotEntry>, };

@@ -431,7 +431,11 @@ entry_id?: string | null, segments: Array<Segment>, } } | { "event": "session_en
 /**
 * Stable durable history identity. Synthetic compatibility events omit it.
 */
-entry_id?: string | null, item: unknown, } } | { "event": "invoke_start", "data": { kind: InvokeKind, } } | { "event": "turn_start", "data": { turn: number, } } | { "event": "turn_end", "data": { turn: number, result: TurnResult, } } | { "event": "llm_call_start", "data": { llm_call: number, } } | { "event": "llm_call_end", "data": { llm_call: number, } } | { "event": "llm_retry", "data": { llm_call: number,
+entry_id?: string | null, item: unknown, } } | { "event": "invoke_start", "data": { kind: InvokeKind,
+/**
+ * Timestamp of the persisted Invoke, not the client's attach time.
+ */
+timestamp_ms?: number | null, } } | { "event": "turn_start", "data": { turn: number, } } | { "event": "turn_end", "data": { turn: number, result: TurnResult, } } | { "event": "llm_call_start", "data": { llm_call: number, } } | { "event": "llm_call_end", "data": { llm_call: number, } } | { "event": "llm_retry", "data": { llm_call: number,
 /**
 * The attempt that just failed. 1 origin.
 */
@@ -445,7 +449,11 @@ summary: string,
 * Full tool output. Absent when the tool chose to return
 * summary-only, or when the result was pruned.
 */
-output?: string | null, disposition?: ToolResultDisposition | null, is_error: boolean, } } | { "event": "usage", "data": { input_tokens: number | null, output_tokens: number | null, cache_read_input_tokens?: number | null, } } | { "event": "context_usage", "data": { usage?: ContextUsage | null, } } | { "event": "run_end", "data": { result: RunResult, } } | { "event": "error", "data": { code: ErrorCode, message: string, } } | { "event": "snapshot", "data": { session: SessionSnapshot, greeting: Greeting,
+output?: string | null, disposition?: ToolResultDisposition | null, is_error: boolean, } } | { "event": "usage", "data": {
+/**
+ * Present for usage delivered from the committed session-log lane.
+ */
+timestamp_ms?: number | null, input_tokens: number | null, output_tokens: number | null, cache_read_input_tokens?: number | null, } } | { "event": "context_usage", "data": { usage?: ContextUsage | null, } } | { "event": "run_end", "data": { result: RunResult, } } | { "event": "error", "data": { code: ErrorCode, message: string, } } | { "event": "snapshot", "data": { session: SessionSnapshot, greeting: Greeting,
 /**
 * Full revisioned live execution state. `Stopped` remains Runtime
 * catalog authority and is deliberately not represented here.
@@ -26,6 +26,32 @@ declare const Deno: {
  test(name: string, fn: () => void): void;
 };

+Deno.test("restored run stats replace live estimates in both Console modes", () => {
+  const answer: SessionSnapshotEntry = {
+    kind: "message", role: "assistant", content: [{ kind: "text", text: "answer" }],
+    entry_id: "answer-1", timestamp: 5000, provenance: "model_output", derived_from: [],
+  };
+  const stats: SessionSnapshotEntry = {
+    kind: "run_stats", entry_id: "run-stats:answer-1", timestamp: 6000,
+    provenance: "legacy_unknown", derived_from: ["answer-1"],
+    elapsed_ms: 5000, requests: 2, upload_tokens: 1500, output_tokens: 200,
+  };
+  for (const commitBeforeEnd of [true, false]) {
+    const projector = createConsoleProjector();
+    projector.append([{ eventId: "text", event: { event: "text_done", data: { text: "answer" } } }]);
+    const commit = { eventId: "commit", event: { event: "session_entry_committed", data: { entry: answer } } } as const;
+    const end = { eventId: "end", event: { event: "run_end", data: { result: "finished" } } } as const;
+    const live = projector.append(commitBeforeEnd ? [commit, end] : [end, commit]);
+    const history = projectSessionHistoryEntries([answer, stats], null);
+    for (const mode of ["normal", "overview"] as const) {
+      const lines = projectConsoleLines(mergeCommittedHistoryLines(history, live.lines), mode);
+      assertEquals(lines.map((line) => line.kind), ["assistant", "run_stats"]);
+      assertEquals(lines[1].body, "5s ・2 reqs ↑1.5k/↓200");
+      assertEquals(projectConsoleLines(history, mode)[1].body, lines[1].body);
+    }
+  }
+});
+
 function workerState(status: WorkerStatus): WorkerStateSnapshot {
  return {
    last_command_id: 0,
@@ -26,6 +26,7 @@ import {
  emptyRunActivityStats,
  formatRunElapsedCompact,
  formatRunTokens,
+  restoredRunActivity,
  type RunActivityStats,
 } from "./run-status.ts";
 import {
@@ -336,9 +337,17 @@ export function mergeCommittedHistoryLines(
  history: ConsoleLine[],
  current: ConsoleLine[],
 ): ConsoleLine[] {
+  const durableStats = new Map(
+    history.filter((line) => line.kind === "run_stats" && line.entryId)
+      .map((line) => [line.entryId!, line]),
+  );
  return aggregateReadToolLines(mergeHistorySourceLines(
    history.flatMap((line) => line.toolCallLines ?? [line]),
-    current.flatMap((line) => line.toolCallLines ?? [line]),
+    current.flatMap((line) => line.toolCallLines ?? [line]).map((line) =>
+      line.kind === "run_stats" && line.entryId
+        ? durableStats.get(line.entryId) ?? line
+        : line
+    ),
  ));
 }

@@ -866,6 +875,7 @@ function projectInternalWorkerSnapshot(
    cwd,
  );
  console.status = snapshot.status;
+  console.runActivity = restoredRunActivity(snapshot.session.entries);
  console.workerMetadata = snapshot.greeting
    ? metadataFromGreeting(snapshot.greeting)
    : null;
@@ -1076,6 +1086,13 @@ export function applyProtocolEvent(
        if (committedIndex >= 0) {
          const [committed] = next.lines.splice(committedIndex, 1);
          next.lines.splice(replacementIndex, 0, committed);
+          const stats = next.lines[replacementIndex + 1];
+          if (stats?.kind === "run_stats" && !stats.entryId) {
+            next.lines[replacementIndex + 1] = {
+              ...stats,
+              entryId: `run-stats:${committed.entryId}`,
+            };
+          }
        }
      }
      break;
@@ -1294,15 +1311,24 @@ export function applyProtocolEvent(
    case "llm_retry":
    case "llm_continuation":
      break;
-    case "run_end":
-      next.lines.push(
-        runStatsLine(
-          envelope.eventId,
-          next.runActivity,
-          envelope.observedAtMs ?? next.runActivity.startedAtMs ?? 0,
-        ),
+    case "run_end": {
+      const anchor = next.lines.findLast((line) =>
+        line.kind === "assistant" || line.kind === "user" ||
+        line.kind === "tool" || line.kind === "system"
      );
+      const summary = runStatsLine(
+        envelope.eventId,
+        next.runActivity,
+        envelope.observedAtMs ?? next.runActivity.startedAtMs ?? 0,
+      );
+      if (anchor?.entryId) summary.entryId = `run-stats:${anchor.entryId}`;
+      const previous = next.lines.findIndex((line) =>
+        summary.entryId && line.entryId === summary.entryId
+      );
+      if (previous >= 0) next.lines[previous] = summary;
+      else next.lines.push(summary);
      break;
+    }
    case "alert":
      appendAlertLine(next, envelope.eventId, event.data);
      break;
@@ -2402,6 +2428,21 @@ function applySessionEntry(
      applyTaskSystemItem(projection, item);
      break;
    }
+    case "run_stats": {
+      const elapsed = safeTokenCount(value["elapsed_ms"]);
+      const requests = safeTokenCount(value["requests"]);
+      const upload = safeTokenCount(value["upload_tokens"]);
+      const output = safeTokenCount(value["output_tokens"]);
+      if (elapsed !== null && requests !== null && upload !== null && output !== null) {
+        projection.lines.push(runStatsLine(eventId, {
+          startedAtMs: 0,
+          requests,
+          uploadTokens: upload,
+          outputTokens: output,
+        }, elapsed));
+      }
+      break;
+    }
    case "run_yielded":
    case "run_resumed":
    case "run_cancelled":
@@ -114,7 +114,32 @@ Deno.test("run activity follows TUI request and net-token accounting", () => {
  });
 });

-Deno.test("new invoke and running snapshot reset run activity", () => {
+Deno.test("reconnect rebuilds run totals from Invoke and individual persisted usage", () => {
+  const entries = [
+    { kind: "invoke", timestamp: 1_000, trigger: "user_send" },
+    { kind: "usage", timestamp: 2_000, input_tokens: 25000, cache_read_input_tokens: 20000, output_tokens: 300 },
+  ];
+  const snapshot = { event: "snapshot", data: {
+    session: { entries },
+    state: { state: { kind: "busy", state: { kind: "run", state: "running" } } },
+  } };
+  let stats = applyRunActivityEvent(emptyRunActivityStats(), snapshot, 60_000);
+  const restored = { startedAtMs: 1_000, requests: 1, uploadTokens: 5_000, outputTokens: 300, fromLog: true };
+  assertEquals(stats, restored);
+  // TurnStart is operational; each committed Usage is counted once instead.
+  stats = applyRunActivityEvent(stats, { event: "turn_start", data: { turn: 2 } }, 61_000);
+  assertEquals(stats, restored);
+  stats = applyRunActivityEvent(stats, { event: "usage", data: {
+    timestamp_ms: 62_000, input_tokens: 2000, cache_read_input_tokens: 1000, output_tokens: 100,
+  } }, 62_100);
+  assertEquals(stats, { ...restored, requests: 2, uploadTokens: 6000, outputTokens: 400 });
+  // Reapplying a snapshot replaces the prefix, rather than adding it again.
+  assertEquals(applyRunActivityEvent(stats, snapshot, 80_000), restored);
+  stats = applyRunActivityEvent(stats, { event: "invoke_start", data: { kind: "notify", timestamp_ms: 90_000 } }, 99_000);
+  assertEquals(stats, { startedAtMs: 90_000, requests: 0, uploadTokens: 0, outputTokens: 0, fromLog: true });
+});
+
+Deno.test("new invoke and legacy snapshot reset run activity", () => {
  const previous = {
    startedAtMs: 1,
    requests: 3,
@@ -1,10 +1,13 @@
 import type {
  Event as ProtocolEvent,
  InFlightCompaction,
+  SessionSnapshotEntry,
  WorkerStateSnapshot,
 } from "#lib/generated/protocol.ts";

 export type RunActivityStats = {
+  /** Request counts follow persisted usage when the producer supplies it. */
+  fromLog?: boolean;
  startedAtMs: number | null;
  requests: number;
  uploadTokens: number;
@@ -20,6 +23,25 @@ export function emptyRunActivityStats(): RunActivityStats {
  };
 }

+export function restoredRunActivity(entries: SessionSnapshotEntry[]): RunActivityStats {
+  let stats = emptyRunActivityStats();
+  for (const entry of entries) {
+    if (entry.kind === "invoke") {
+      stats = applyRunActivityEvent(stats, { event: "invoke_start", data: {
+        kind: entry.trigger, timestamp_ms: entry.timestamp,
+      } }, entry.timestamp);
+    } else if (entry.kind === "usage") {
+      stats = applyRunActivityEvent(stats, { event: "usage", data: {
+        timestamp_ms: entry.timestamp,
+        input_tokens: entry.input_tokens,
+        cache_read_input_tokens: entry.cache_read_input_tokens,
+        output_tokens: entry.output_tokens,
+      } }, entry.timestamp);
+    }
+  }
+  return stats;
+}
+
 export function applyRunActivityEvent(
  current: RunActivityStats,
  event: ProtocolEvent,
@@ -27,14 +49,22 @@ export function applyRunActivityEvent(
 ): RunActivityStats {
  switch (event.event) {
    case "invoke_start":
-      return { ...emptyRunActivityStats(), startedAtMs: observedAtMs };
-    case "snapshot":
+      return {
+        ...emptyRunActivityStats(),
+        ...(event.data.timestamp_ms != null ? { fromLog: true } : {}),
+        startedAtMs: event.data.timestamp_ms ?? observedAtMs,
+      };
+    case "snapshot": {
+      const restored = restoredRunActivity(event.data.session?.entries ?? []);
+      if (restored.fromLog) return restored;
      return event.data.state.state.kind === "busy" &&
          !(event.data.state.state.state.kind === "run" &&
            event.data.state.state.state.state === "paused")
        ? { ...emptyRunActivityStats(), startedAtMs: observedAtMs }
        : emptyRunActivityStats();
+    }
    case "turn_start":
+      if (current.fromLog) return current;
      return {
        ...current,
        startedAtMs: current.startedAtMs ?? observedAtMs,
@@ -43,9 +73,12 @@ export function applyRunActivityEvent(
    case "usage": {
      const input = event.data.input_tokens ?? 0;
      const cacheRead = event.data.cache_read_input_tokens ?? 0;
+      const fromLog = event.data.timestamp_ms != null || current.fromLog;
      return {
        ...current,
-        startedAtMs: current.startedAtMs ?? observedAtMs,
+        ...(fromLog ? { fromLog: true } : {}),
+        requests: current.requests + (fromLog ? 1 : 0),
+        startedAtMs: current.startedAtMs ?? (fromLog ? null : observedAtMs),
        uploadTokens: current.uploadTokens + Math.max(0, input - cacheRead),
        outputTokens: current.outputTokens + (event.data.output_tokens ?? 0),
      };
@@ -741,6 +741,7 @@
                    event.event.event === "segment_rotated" ||
                    event.event.event === "rewind_applied" ||
                    event.event.event === "user_message" ||
+                    event.event.event === "run_end" ||
                    event.event.event === "session_entry_committed"
                )
            ) {
@@ -265,6 +265,86 @@ async function selectRun(cm: EditorView) {
  await waitFor(() => expect(cm.dom.querySelector(".composer-typed-chip")).not.toBeNull());
 }

+test("a completed live run displays its stats below the answer", async () => {
+  const { view } = await liveComposer();
+  protocolEvent({ event: "invoke_start", data: { kind: "user_send" } });
+  protocolEvent({ event: "turn_start", data: { turn: 1 } });
+  protocolEvent({ event: "text_done", data: { text: "The final answer" } });
+  protocolEvent({ event: "usage", data: {
+    input_tokens: 2500, cache_read_input_tokens: 1000, output_tokens: 200,
+  } });
+  protocolEvent({ event: "run_end", data: { result: "finished" } });
+  protocolEvent({ event: "worker_state", data: {
+    snapshot: { last_command_id: 0, state: { kind: "idle" } },
+  } });
+  await waitFor(() => {
+    const rows = view.container.querySelectorAll(".console-line");
+    expect(rows[rows.length - 1]?.textContent).toContain("1 reqs ↑1.5k/↓200");
+    expect(view.container.querySelector(".run-stats")).not.toBeNull();
+  });
+});
+
+test("reopening Console restores completed run stats from paged history", async () => {
+  const page = historyPage([1], null);
+  const answer = page.page.turns[0].entries[1];
+  const stats = {
+    kind: "run_stats", entry_id: `run-stats:${answer.entry_id}`, timestamp: 6000,
+    provenance: "legacy_unknown", derived_from: [answer.entry_id],
+    elapsed_ms: 5000, requests: 2, upload_tokens: 1500, output_tokens: 200,
+  };
+  vi.stubGlobal("fetch", vi.fn(async (input: RequestInfo | URL) => {
+    return Response.json(String(input).includes("/session/history")
+      ? { ...page, page: { ...page.page, turns: [{ ...page.page.turns[0], entries: [...page.page.turns[0].entries, stats] }] } }
+      : { availability: "live_protocol" });
+  }));
+  for (let visit = 1; visit <= 2; visit++) {
+    const view = render(ConsolePage, { data: pageData() });
+    await waitFor(() => expect(multiplexer.subscribe).toHaveBeenCalledTimes(visit));
+    latestListener().onFrame(subscribedFrame());
+    await waitFor(() => {
+      const rows = view.container.querySelectorAll(".console-line");
+      expect(rows[rows.length - 1]?.textContent).toContain("5s ・2 reqs ↑1.5k/↓200");
+      expect(view.container.querySelectorAll(".run-stats")).toHaveLength(1);
+    });
+    view.unmount();
+  }
+});
+
+test("reopening Console keeps the Composer run clock and cumulative traffic", async () => {
+  vi.stubGlobal("fetch", vi.fn(async (input: RequestInfo | URL) =>
+    Response.json(String(input).includes("/session/history")
+      ? emptyHistoryPage() : { availability: "live_protocol" })
+  ));
+  const startedAt = Date.now() - 120_000;
+  for (let visit = 1; visit <= 2; visit++) {
+    const view = render(ConsolePage, { data: pageData() });
+    await waitFor(() => expect(multiplexer.subscribe).toHaveBeenCalledTimes(visit));
+    const snapshot = snapshotEvent();
+    if (snapshot.event !== "snapshot") throw new Error("expected snapshot");
+    snapshot.data.state = { last_command_id: 0, state: { kind: "busy", state: { kind: "run", state: "running" } } };
+    snapshot.data.session.entries = [
+      { kind: "invoke", trigger: "user_send", timestamp: startedAt, entry_id: "invoke", provenance: "legacy_unknown", derived_from: [] },
+      { kind: "usage", timestamp: startedAt + 1000, entry_id: "usage", provenance: "legacy_unknown", derived_from: [], input_tokens: 25_000, cache_read_input_tokens: 20_000, output_tokens: 300 },
+    ];
+    const frame = subscribedFrame();
+    frame.message.payload.snapshot.data.events = [snapshot];
+    latestListener().onFrame(frame);
+    await waitFor(() => {
+      const status = view.container.querySelector(".worker-run-status");
+      expect(status?.textContent).toContain("2m");
+      expect(status?.textContent).toContain("1 req");
+      expect(status?.textContent).toContain("Run ↑5.0k/↓300");
+    });
+    protocolEvent({ event: "usage", data: { timestamp_ms: Date.now(), input_tokens: 25_000, cache_read_input_tokens: 20_000, output_tokens: 100 } });
+    await waitFor(() => {
+      const status = view.container.querySelector(".worker-run-status");
+      expect(status?.textContent).toContain("2 reqs");
+      expect(status?.textContent).toContain("Run ↑10.0k/↓400");
+    });
+    view.unmount();
+  }
+});
+
 test("Console teardown cancels pending observation frames before they can read history", async () => {
  const { view, fetchMock } = await liveComposer();
  const schedule = vi.spyOn(window, "requestAnimationFrame");