fix: make session resume atomic across restarts

2026-07-15 00:42:35 +02:00
parent 2b4797f133
commit 08b1f4909e
12 changed files with 1082 additions and 107 deletions
@@ -134,6 +134,60 @@ test("an already-active same-session Resume preserves its EventSource and store"
  expect(useSessionStore.getState().activityLog).toEqual(before);
 });

+test("a false then true pair of concurrent Resume completions preserves the cold stream", async () => {
+  let resumeCalls = 0;
+  let releaseCold!: () => void;
+  let releaseAlreadyActive!: () => void;
+  let markColdStarted!: () => void;
+  let markAlreadyActiveStarted!: () => void;
+  const coldStarted = new Promise<void>((resolve) => { markColdStarted = resolve; });
+  const alreadyActiveStarted = new Promise<void>((resolve) => { markAlreadyActiveStarted = resolve; });
+  const coldReleased = new Promise<void>((resolve) => { releaseCold = resolve; });
+  const alreadyActiveReleased = new Promise<void>((resolve) => { releaseAlreadyActive = resolve; });
+  server.use(
+    http.post("http://localhost:8787/sessions/:id/resume", async () => {
+      resumeCalls += 1;
+      if (resumeCalls === 1) {
+        markColdStarted();
+        await coldReleased;
+        return resumeResult("s1", false);
+      }
+      markAlreadyActiveStarted();
+      await alreadyActiveReleased;
+      return resumeResult("s1", true);
+    }),
+    http.get("http://localhost:8787/sessions/:id", () =>
+      HttpResponse.json({ id: "s1", status: "open", phase: 1 })),
+  );
+  wrap();
+  await userEvent.click(await screen.findByText("Attiva uno"));
+  const resumeButton = await screen.findByRole("button", { name: /resume/i });
+
+  await userEvent.click(resumeButton);
+  await coldStarted;
+  await userEvent.click(resumeButton);
+  await alreadyActiveStarted;
+  expect(resumeCalls).toBe(2);
+
+  releaseCold();
+  await waitFor(() => expect(FakeEventSource.instances).toHaveLength(1));
+  const stream = FakeEventSource.instances[0];
+  act(() => stream.emitNamed(
+    "ui_request",
+    { type: "ui_request", ui_request: { id: "cold-gate", widget: "select" } },
+    "1",
+  ));
+  const before = useSessionStore.getState().activityLog.map((entry) => ({ ...entry }));
+
+  releaseAlreadyActive();
+  await waitFor(() => expect(screen.queryByText("Domanda originale")).not.toBeInTheDocument());
+
+  expect(FakeEventSource.instances).toHaveLength(1);
+  expect(stream.closed).toBe(false);
+  expect(useSessionStore.getState().pendingWidget?.id).toBe("cold-gate");
+  expect(useSessionStore.getState().activityLog).toEqual(before);
+});
+
 test("cold same-session Resume keeps the old stream until success then receives post-clear events once", async () => {
  server.use(
    http.post("http://localhost:8787/sessions/:id/resume", () =>
@@ -146,7 +200,7 @@ test("cold same-session Resume keeps the old stream until success then receives
  await userEvent.click(await screen.findByRole("button", { name: /resume/i }));
  await waitFor(() => expect(FakeEventSource.instances).toHaveLength(1));
  const first = FakeEventSource.instances[0];
-  act(() => first.emitNamed("info", { type: "info", text: "Before cold Resume" }, "7"));
+  act(() => first.emitNamed("info", { type: "info", text: "Before cold Resume" }, "900"));

  let releaseResume!: () => void;
  let markStarted!: () => void;
@@ -168,7 +222,7 @@ test("cold same-session Resume keeps the old stream until success then receives
  expect(useSessionStore.getState().activityLog).toContainEqual({
    kind: "status", phase: "F1", text: "Before cold Resume", level: "info",
  });
-  act(() => first.emitNamed("text_delta", { type: "text_delta", text: "Still attached" }, "8"));
+  act(() => first.emitNamed("text_delta", { type: "text_delta", text: "Still attached" }, "901"));
  expect(useSessionStore.getState().transcript).toEqual([
    { role: "assistant", text: "Still attached" },
  ]);
@@ -177,7 +231,7 @@ test("cold same-session Resume keeps the old stream until success then receives
  await waitFor(() => expect(FakeEventSource.instances).toHaveLength(2));
  expect(first.closed).toBe(true);
  const replacement = FakeEventSource.instances[1];
-  expect(replacement.url).toContain("/sessions/s1/events?lastEventId=8");
+  expect(replacement.url).toBe("http://localhost:8787/sessions/s1/events");
  expect(useSessionStore.getState().transcript).toEqual([]);
  expect(useSessionStore.getState().activityLog).toEqual([
    { kind: "lifecycle", phase: null, text: "Resuming session" },
@@ -186,7 +240,7 @@ test("cold same-session Resume keeps the old stream until success then receives
  act(() => replacement.emitNamed(
    "text_delta",
    { type: "text_delta", text: "Post-resume delivery" },
-    "9",
+    "1",
  ));
  expect(useSessionStore.getState().transcript).toEqual([
    { role: "assistant", text: "Post-resume delivery" },
@@ -199,8 +253,8 @@ test("cold same-session Resume keeps the old stream until success then receives
    ui_request: { id: "post-resume-gate", widget: "select", title: "Review once" },
  };
  act(() => {
-    replacement.emitNamed("ui_request", gate, "10");
-    replacement.emitNamed("ui_request", gate, "11");
+    replacement.emitNamed("ui_request", gate, "2");
+    replacement.emitNamed("ui_request", gate, "3");
  });
  expect(useSessionStore.getState().pendingWidget?.id).toBe("post-resume-gate");
  expect(useSessionStore.getState().activityLog.filter((entry) => entry.kind === "gate"))
@@ -32,7 +32,8 @@ import { useEffect, useMemo, useRef, useState } from "react";
 */
 export function AppShell() {
  const [activeSessionId, setActiveSessionId] = useState<string | null>(null);
-  const [streamGeneration, setStreamGeneration] = useState(0);
+  const activeSessionIdRef = useRef<string | null>(null);
+  const [streamCursorResetEpoch, setStreamCursorResetEpoch] = useState(0);
  const [creatingSession, setCreatingSession] = useState(false);
  const [awaitingQuestion, setAwaitingQuestion] = useState(false);
  const { data: sessions = [] } = useQuery<SessionSummary[]>({
@@ -65,6 +66,12 @@ export function AppShell() {
  const selectedSessions = sessions.filter((session) => selectedSessionIds.has(session.id));
  const allSessionsSelected = sessions.length > 0 && selectedSessions.length === sessions.length;

+  function selectActiveSession(id: string | null) {
+    // Keep async Resume completions synchronized before React commits the state update.
+    activeSessionIdRef.current = id;
+    setActiveSessionId(id);
+  }
+
  // A background refresh can remove a session (for example from another browser).
  // Keep the local selection aligned with the authoritative list.
  useEffect(() => {
@@ -100,9 +107,9 @@ export function AppShell() {
    });
  }
  async function doResume(id: string) {
-    const reconnectSameSession = activeSessionId === id;
    try {
      const result = await resumeSession(id);
+      const reconnectSameSession = activeSessionIdRef.current === id;
      setPanelSession(null);
      setAwaitingQuestion(false);

@@ -115,11 +122,11 @@ export function AppShell() {
        recordLifecycle("Resuming session");
        setAgentActive(true);
      }
-      setActiveSessionId(id);
-      // The backend has now completed clear/rebind. Recreate a same-id source with its cursor;
-      // switching ids naturally creates a new source and resets the cursor in the stream hook.
+      selectActiveSession(id);
+      // A cold runtime starts a fresh SSE id sequence. Recreate a same-id source only after
+      // Resume succeeds, and explicitly discard the old runtime's cursor.
      if (!result.alreadyActive && reconnectSameSession) {
-        setStreamGeneration((value) => value + 1);
+        setStreamCursorResetEpoch((value) => value + 1);
      }

      // Paint the persisted re-entry phase while the replacement stream starts replaying.
@@ -182,7 +189,7 @@ export function AppShell() {
        targets.filter((_, index) => results[index].status === "fulfilled").map((session) => session.id),
      );
      if (deletedIds.has(panelSession?.id ?? "")) setPanelSession(null);
-      if (deletedIds.has(activeSessionId ?? "")) { resetSession(); setActiveSessionId(null); }
+      if (deletedIds.has(activeSessionId ?? "")) { resetSession(); selectActiveSession(null); }
      setSelectedSessionIds((current) => new Set([...current].filter((id) => !deletedIds.has(id))));
      refresh();
      if (deletedIds.size !== targets.length) {
@@ -226,7 +233,7 @@ export function AppShell() {
  // session sits idle or a gate awaits the reviewer (pendingWidget).
  const running = working && !finalized;

-  useSessionStream(activeSessionId, streamGeneration);
+  useSessionStream(activeSessionId, 0, streamCursorResetEpoch);

  // A backend "session_exit" system event (e.g. the replay server emitting it
  // when the reviewer picks "Esci") asks us to leave the live session view and
@@ -240,7 +247,7 @@ export function AppShell() {
      // Never let a streamed event terminate the managed Pi child. Only the
      // explicit “Stop & save” action is allowed to call /close.
      resetSession();
-      setActiveSessionId(null);
+      selectActiveSession(null);
      setAwaitingQuestion(false);
    }
    // The final workflow turn ends with the session already finalized on disk:
@@ -254,7 +261,7 @@ export function AppShell() {
    resetSession();
    setAwaitingQuestion(true);
    setCreatingSession(false);
-    setActiveSessionId(null);
+    selectActiveSession(null);
    // Best effort only: session creation keeps the authoritative readiness gate.
    // Composer focus is deliberately independent of this network request.
    void prewarmRuntime().catch(() => undefined);
@@ -269,7 +276,7 @@ export function AppShell() {
  function finishSessionCreation(id: string) {
    // React batches these updates, preserving the provisional view and timer
    // while useSessionStream opens the durable session's SSE channel.
-    setActiveSessionId(id);
+    selectActiveSession(id);
    setCreatingSession(false);
    setAwaitingQuestion(false);
    refresh();
@@ -287,7 +294,7 @@ export function AppShell() {
      await closeSession(activeSessionId);
    } finally {
      resetSession();
-      setActiveSessionId(null);
+      selectActiveSession(null);
      setAwaitingQuestion(false);
    }
  }
@@ -1,5 +1,5 @@
 import { renderHook } from "@testing-library/react";
-import { act } from "react";
+import { act, useLayoutEffect } from "react";
 import { FakeEventSource } from "../test/fakeEventSource";
 import { useSessionStream } from "./useSessionStream";
 import { useSessionStore } from "../store/sessionStore";
@@ -100,6 +100,58 @@ test("a same-session generation reconnect includes the last consumed SSE id", ()
    .toEqual([{ kind: "assistant", phase: null, text: "hello world" }]);
 });

+test("a same-session reset epoch drops a high cursor and accepts fresh low-id events", () => {
+  const { rerender } = renderHook(
+    ({ resetEpoch }) => useSessionStream("s1", 0, resetEpoch),
+    { initialProps: { resetEpoch: 0 } },
+  );
+  const first = FakeEventSource.instances[0];
+  act(() => first.emitNamed("info", { type: "info", text: "old runtime" }, "900"));
+
+  rerender({ resetEpoch: 1 });
+
+  expect(first.closed).toBe(true);
+  expect(FakeEventSource.instances).toHaveLength(2);
+  expect(FakeEventSource.instances[1].url).toBe(
+    "http://localhost:8787/sessions/s1/events",
+  );
+
+  act(() =>
+    FakeEventSource.instances[1].emitNamed(
+      "ui_request",
+      { type: "ui_request", ui_request: { id: "fresh-gate", widget: "select" } },
+      "1",
+    ),
+  );
+  expect(useSessionStore.getState().pendingWidget?.id).toBe("fresh-gate");
+});
+
+test("a reset epoch also accepts an in-process preserved high id and resumes from it", () => {
+  const { rerender } = renderHook(
+    ({ generation, resetEpoch }) => useSessionStream("s1", generation, resetEpoch),
+    { initialProps: { generation: 0, resetEpoch: 0 } },
+  );
+  const first = FakeEventSource.instances[0];
+  act(() => first.emitNamed("info", { type: "info", text: "old runtime" }, "900"));
+
+  rerender({ generation: 0, resetEpoch: 1 });
+  const replacement = FakeEventSource.instances[1];
+  expect(replacement.url).toBe("http://localhost:8787/sessions/s1/events");
+  act(() =>
+    replacement.emitNamed(
+      "ui_request",
+      { type: "ui_request", ui_request: { id: "preserved-high-gate", widget: "select" } },
+      "902",
+    ),
+  );
+
+  expect(useSessionStore.getState().pendingWidget?.id).toBe("preserved-high-gate");
+  rerender({ generation: 1, resetEpoch: 1 });
+  expect(FakeEventSource.instances[2].url).toBe(
+    "http://localhost:8787/sessions/s1/events?lastEventId=902",
+  );
+});
+
 test("changing the session resets the manual reconnect cursor", () => {
  const { rerender } = renderHook(
    ({ sessionId }) => useSessionStream(sessionId),
@@ -113,3 +165,55 @@ test("changing the session resets the manual reconnect cursor", () => {
  expect(first.closed).toBe(true);
  expect(FakeEventSource.instances[1].url).toBe("http://localhost:8787/sessions/s2/events");
 });
+
+test("a queued event from a replaced source cannot mutate the new session or its cursor", () => {
+  const { rerender } = renderHook(
+    ({ sessionId, generation }) => useSessionStream(sessionId, generation, 0),
+    { initialProps: { sessionId: "s1" as string | null, generation: 0 } },
+  );
+  const first = FakeEventSource.instances[0];
+  act(() => first.emitNamed("info", { type: "info", text: "session one" }, "12"));
+
+  rerender({ sessionId: "s2", generation: 0 });
+  expect(first.closed).toBe(true);
+  expect(FakeEventSource.instances[1].url).toBe("http://localhost:8787/sessions/s2/events");
+  useSessionStore.getState().resetSession();
+
+  act(() =>
+    first.emitQueuedNamed(
+      "text_delta",
+      { type: "text_delta", text: "stale session one" },
+      "999",
+    ),
+  );
+  expect(useSessionStore.getState().transcript).toEqual([]);
+
+  rerender({ sessionId: "s2", generation: 1 });
+  expect(FakeEventSource.instances[2].url).toBe("http://localhost:8787/sessions/s2/events");
+});
+
+test("the old source is invalid before later layout effects can deliver a queued event", () => {
+  let first: FakeEventSource;
+  const { rerender } = renderHook(
+    ({ sessionId, emitDuringLayout }) => {
+      useSessionStream(sessionId);
+      useLayoutEffect(() => {
+        if (emitDuringLayout) {
+          first.emitQueuedNamed(
+            "text_delta",
+            { type: "text_delta", text: "commit-window stale text" },
+            "777",
+          );
+        }
+      }, [emitDuringLayout, sessionId]);
+    },
+    { initialProps: { sessionId: "s1" as string | null, emitDuringLayout: false } },
+  );
+  first = FakeEventSource.instances[0];
+  useSessionStore.getState().resetSession();
+
+  rerender({ sessionId: "s2", emitDuringLayout: true });
+
+  expect(first.closed).toBe(true);
+  expect(useSessionStore.getState().transcript).toEqual([]);
+});
@@ -1,29 +1,51 @@
-import { useEffect, useRef, useState } from "react";
+import { useLayoutEffect, useRef, useState } from "react";
 import { BASE } from "../api/client";
 import { joinBackendPath } from "../api/runtime-config";
 import { useSessionStore } from "../store/sessionStore";
 import type { StreamEvent } from "../api/types";

-export function useSessionStream(sessionId: string | null, generation = 0) {
+export function useSessionStream(
+  sessionId: string | null,
+  generation = 0,
+  cursorResetEpoch = 0,
+) {
  const [connected, setConnected] = useState(false);
  const applyEvent = useSessionStore((s) => s.applyEvent);
-  const cursor = useRef({ sessionId: null as string | null, lastEventId: "" });
+  const cursor = useRef({
+    sessionId: null as string | null,
+    cursorResetEpoch,
+    lastEventId: "",
+  });
+  const activeSource = useRef<object | null>(null);

-  if (cursor.current.sessionId !== sessionId) {
-    cursor.current = { sessionId, lastEventId: "" };
-  }
-
-  useEffect(() => {
-    if (!sessionId) return;
+  useLayoutEffect(() => {
+    if (
+      cursor.current.sessionId !== sessionId
+      || cursor.current.cursorResetEpoch !== cursorResetEpoch
+    ) {
+      cursor.current = { sessionId, cursorResetEpoch, lastEventId: "" };
+    }
+    if (!sessionId) {
+      activeSource.current = null;
+      setConnected(false);
+      return;
+    }

    const query = cursor.current.lastEventId
      ? `?lastEventId=${encodeURIComponent(cursor.current.lastEventId)}`
      : "";
    const es = new EventSource(joinBackendPath(BASE, `/sessions/${sessionId}/events${query}`));
-    es.onopen = () => setConnected(true);
-    es.onerror = () => setConnected(false);
+    const identity = { source: es, sessionId, cursorResetEpoch };
+    activeSource.current = identity;
+    es.onopen = () => {
+      if (activeSource.current === identity) setConnected(true);
+    };
+    es.onerror = () => {
+      if (activeSource.current === identity) setConnected(false);
+    };

    const handle = (ev: MessageEvent) => {
+      if (activeSource.current !== identity) return;
      if (ev.lastEventId) cursor.current.lastEventId = ev.lastEventId;
      try {
        applyEvent(JSON.parse(ev.data) as StreamEvent);
@@ -47,10 +69,14 @@ export function useSessionStream(sessionId: string | null, generation = 0) {

    return () => {
      for (const name of namedEvents) es.removeEventListener(name, handle);
+      es.onmessage = null;
      es.close();
-      setConnected(false);
+      if (activeSource.current === identity) {
+        activeSource.current = null;
+        setConnected(false);
+      }
    };
-  }, [sessionId, generation, applyEvent]);
+  }, [sessionId, generation, cursorResetEpoch, applyEvent]);

  return { connected };
 }
@@ -5,6 +5,10 @@ export class FakeEventSource {
  onerror: (() => void) | null = null;
  closed = false;
  private listeners = new Map<string, Set<(e: { data: string; lastEventId: string }) => void>>();
+  private historicalListeners = new Map<
+    string,
+    Set<(e: { data: string; lastEventId: string }) => void>
+  >();
  constructor(public url: string) {
    FakeEventSource.instances.push(this);
  }
@@ -17,9 +21,16 @@ export class FakeEventSource {
    const data = JSON.stringify(obj);
    for (const h of this.listeners.get(type) ?? []) h({ data, lastEventId });
  }
+  /** Simulate an event already queued before removeEventListener/close completed. */
+  emitQueuedNamed(type: string, obj: unknown, lastEventId = "") {
+    const data = JSON.stringify(obj);
+    for (const h of this.historicalListeners.get(type) ?? []) h({ data, lastEventId });
+  }
  addEventListener(type: string, handler: (e: { data: string; lastEventId: string }) => void) {
    if (!this.listeners.has(type)) this.listeners.set(type, new Set());
+    if (!this.historicalListeners.has(type)) this.historicalListeners.set(type, new Set());
    this.listeners.get(type)!.add(handler);
+    this.historicalListeners.get(type)!.add(handler);
  }
  removeEventListener(type: string, handler: (e: { data: string; lastEventId: string }) => void) {
    this.listeners.get(type)?.delete(handler);