Fix logout race: clear partitioned cookies, recheck session on 401

.changeset/logout-partitioned-cookie-clear.md

@@ -0,0 +1,5 @@
+---
+"@agent-native/core": patch
+---
+
+Clear framework auth cookies from the CHIPS partition they were set in, so logout cannot leave a live session cookie behind, and re-resolve the client session when a request comes back 401 instead of painting a generic load error.

packages/core/src/client/use-action.spec.ts

@@ -5,6 +5,11 @@ const analyticsMocks = vi.hoisted(() => ({
 }));
 vi.mock("./analytics.js", () => analyticsMocks);
 
+const sessionMocks = vi.hoisted(() => ({
+  recheckSessionAfterUnauthorized: vi.fn(),
+}));
+vi.mock("./use-session.js", () => sessionMocks);
+
 import {
   ACTION_KEEPALIVE_BODY_BUDGET_BYTES,
   actionErrorMessage,
@@ -226,6 +231,49 @@ describe("callAction", () => {
     );
   });
 
+  it("re-resolves the session when an action is refused as unauthenticated", async () => {
+    // 401 means the server stopped recognising this browser. Without telling
+    // the session gate, the shell stays mounted on its last "authenticated"
+    // read and this failure surfaces as a generic load error instead of a
+    // redirect to sign-in - the screen reported after the logout race.
+    vi.stubGlobal(
+      "fetch",
+      vi
+        .fn()
+        .mockResolvedValue(
+          jsonResponse({ error: "Unauthorized" }, { status: 401 }),
+        ),
+    );
+
+    await expect(
+      callAction("list-designs", {}, { method: "GET" }),
+    ).rejects.toMatchObject({ status: 401 });
+
+    expect(sessionMocks.recheckSessionAfterUnauthorized).toHaveBeenCalledTimes(
+      1,
+    );
+  });
+
+  it("leaves the session alone when an action is refused as forbidden", async () => {
+    // 403 is an authenticated caller being refused one thing. Re-reading the
+    // session here would be noise, and treating it as signed-out would sign a
+    // working session out.
+    vi.stubGlobal(
+      "fetch",
+      vi
+        .fn()
+        .mockResolvedValue(
+          jsonResponse({ error: "Forbidden" }, { status: 403 }),
+        ),
+    );
+
+    await expect(
+      callAction("list-designs", {}, { method: "GET" }),
+    ).rejects.toMatchObject({ status: 403 });
+
+    expect(sessionMocks.recheckSessionAfterUnauthorized).not.toHaveBeenCalled();
+  });
+
   it("calls mutating actions through the framework action transport", async () => {
     const fetchMock = vi
       .fn()

packages/core/src/client/use-action.ts

@@ -43,6 +43,7 @@
   reloadForClientCompatibilityMismatch,
 } from "./build-compatibility.js";
 import { ensureEmbedAuthFetchInterceptor } from "./embed-auth.js";
+import { recheckSessionAfterUnauthorized } from "./use-session.js";
 
 const ACTION_PREFIX = agentNativePath("/_agent-native/actions");
 
@@ -191,7 +192,7 @@
 /** Resolves to the union of registered action names, or `string` if no registry exists. */
 type ActionName = keyof ActionRegistry extends never
   ? string
   : (keyof ActionRegistry & string) | (string & {});
 
 /** Resolves the return type of an action, or `any` if not in the registry. */
 type ActionResult<T extends string> = T extends keyof ActionRegistry
@@ -264,7 +265,7 @@
     qs.append(key, JSON.stringify(value));
     return;
   }
   qs.append(key, String(value));
 }
 
 export interface ActionFetchOptions {
@@ -493,6 +494,14 @@
   }
 
   if (!res.ok) {
+    // The server does not recognise this browser any more. Nothing else
+    // tells the session gate that, so without this the shell stays mounted
+    // on a stale authenticated answer and the failure reaches the user as a
+    // generic load error instead of a redirect to sign-in. 403 is
+    // deliberately excluded: that is an authenticated caller being refused
+    // one thing.
+    if (res.status === 401) recheckSessionAfterUnauthorized();
+
     // Text the action itself wrote for the caller, as opposed to transport
     // noise. Only a JSON `error`/`message` qualifies: an HTML error page or a
     // bare status line is not something a UI should ever put in a toast.

packages/core/src/client/use-session.spec.tsx

@@ -10,7 +10,11 @@ const analyticsMocks = vi.hoisted(() => ({
 }));
 vi.mock("./analytics.js", () => analyticsMocks);
 
-import { notifySessionInvalidated, useSession } from "./use-session.js";
+import {
+  notifySessionInvalidated,
+  recheckSessionAfterUnauthorized,
+  useSession,
+} from "./use-session.js";
 
 /**
  * A fresh copy of the session module. `signingOut` is one-way for the life of a
@@ -489,6 +493,89 @@ describe("useSession", () => {
     expect(container.textContent).toBe("signing-out");
   });
 
+  it("re-resolves the session after an authenticated request comes back 401", async () => {
+    // The half-authenticated state behind the reported logout race: the last
+    // completed read said "authenticated", so the shell stays mounted and every
+    // data query paints its own generic load error. Nothing else tells the gate
+    // the server stopped recognising this browser.
+    const fetchMock = vi
+      .fn()
+      .mockResolvedValueOnce(
+        jsonResponse({ userId: "user-401", email: "stale@example.com" }),
+      )
+      .mockResolvedValueOnce(jsonResponse({ error: "Not authenticated" }));
+    vi.stubGlobal("fetch", fetchMock);
+
+    await renderConsumers(["first"]);
+    expect(container.textContent).toBe("stale@example.com");
+
+    await act(async () => {
+      recheckSessionAfterUnauthorized();
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    });
+
+    expect(fetchMock).toHaveBeenCalledTimes(2);
+    expect(container.textContent).toBe("signed-out");
+  });
+
+  it("throttles the 401 re-check so one failing screen cannot storm the session endpoint", async () => {
+    // A listing page fails many queries at once. Each invalidation schedules a
+    // fresh read, so an unthrottled re-check turns one expired cookie into a
+    // request per failing query.
+    const fetchMock = vi.fn(async () =>
+      jsonResponse({ userId: "user-storm", email: "storm@example.com" }),
+    );
+    vi.stubGlobal("fetch", fetchMock);
+
+    await renderConsumers(["first"]);
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+
+    await act(async () => {
+      recheckSessionAfterUnauthorized();
+      recheckSessionAfterUnauthorized();
+      recheckSessionAfterUnauthorized();
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    });
+
+    expect(fetchMock).toHaveBeenCalledTimes(2);
+  });
+
+  it("ignores a 401 re-check once sign-out has started", async () => {
+    // Sign-out revokes the session and then navigates, so the 401s it produces
+    // are expected. Asking again here is exactly how a late reply used to
+    // resurrect the session the document had already given up.
+    const {
+      beginSignOut: begin,
+      recheckSessionAfterUnauthorized: recheck,
+      useSession: useFreshSession,
+    } = await freshSessionModule();
+    function Probe() {
+      return <div data-testid="status">{useFreshSession().status}</div>;
+    }
+    const fetchMock = vi.fn(async () =>
+      jsonResponse({ userId: "user-out", email: "leaving@example.com" }),
+    );
+    vi.stubGlobal("fetch", fetchMock);
+
+    await act(async () => {
+      root.render(<Probe />);
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    });
+    await act(async () => {
+      begin();
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    });
+    fetchMock.mockClear();
+
+    await act(async () => {
+      recheck();
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    });
+
+    expect(fetchMock).not.toHaveBeenCalled();
+    expect(container.textContent).toBe("signing-out");
+  });
+
   it("revalidates a cached session when the browser regains focus", async () => {
     const fetchMock = vi
       .fn()

packages/core/src/client/use-session.ts

@@ -161,6 +161,40 @@ export function isSigningOut(): boolean {
   return signingOut;
 }
 
+const UNAUTHORIZED_RECHECK_MIN_INTERVAL_MS = 5_000;
+let lastUnauthorizedRecheckAt = 0;
+
+/**
+ * Re-resolve the session because an authenticated request came back 401.
+ *
+ * `status` is only written when a session fetch resolves, so a 401 anywhere
+ * else leaves every mounted consumer holding the previous `"authenticated"`
+ * answer. The app shell stays mounted over a session the server no longer
+ * recognises, and each data query paints its own generic load error instead of
+ * the visitor being sent to sign in — which is what a stale cookie surviving
+ * logout looks like on screen. Re-reading the session lets the gate reach the
+ * truth and redirect.
+ *
+ * This asks the server rather than forcing `"unauthenticated"` from here: a
+ * 401 can also come from one request a live session is not allowed to make,
+ * and assuming otherwise would sign that visitor out of a working session.
+ * Throttled because one screen can fail many requests at once, and each
+ * invalidation schedules a fresh read.
+ */
+export function recheckSessionAfterUnauthorized(): void {
+  if (typeof window === "undefined") return;
+  // Already leaving. The sign-out flow owns the navigation from here, and
+  // re-reading could only reintroduce the session this document gave up.
+  if (signingOut) return;
+  const now = Date.now();
+  if (now - lastUnauthorizedRecheckAt < UNAUTHORIZED_RECHECK_MIN_INTERVAL_MS) {
+    return;
+  }
+  lastUnauthorizedRecheckAt = now;
+  installSessionInvalidationListeners();
+  invalidateSessionCache();
+}
+
 /**
  * Enter the terminal `"signing-out"` state for this document.
  *