fix(ui): replay deferred session.message reloads after failed runs

ui/src/ui/app-gateway.node.test.ts

@@ -622,6 +622,38 @@ describe("connectGateway", () => {
     },
   );
 
+  it.each(["aborted", "error"] as const)(
+    "replays deferred session.message reloads after %s clears the active run",
+    (terminalState) => {
+      const { host, client } = connectHostGateway();
+      host.chatRunId = "main-run-3";
+      loadChatHistoryMock.mockClear();
+
+      client.emitEvent({
+        event: "session.message",
+        payload: {
+          sessionKey: "main",
+        },
+      });
+
+      expect(loadChatHistoryMock).not.toHaveBeenCalled();
+
+      client.emitEvent({
+        event: "chat",
+        payload: {
+          runId: "main-run-3",
+          sessionKey: "main",
+          state: terminalState,
+          errorMessage: terminalState === "error" ? "chat failed" : undefined,
+        },
+      });
+
+      expect(host.chatRunId).toBeNull();
+      expect(loadChatHistoryMock).toHaveBeenCalledTimes(1);
+      expect(loadChatHistoryMock).toHaveBeenCalledWith(host);
+    },
+  );
+
   it("clears tracked BTW terminal runs after reconnect hello", () => {
     const host = createHost();
 

ui/src/ui/app-gateway.ts

@@ -97,6 +97,10 @@ type GatewayHost = {
   updateAvailable: UpdateAvailable | null;
 };
 
+type GatewayHostWithDeferredSessionMessageReload = GatewayHost & {
+  pendingSessionMessageReloadSessionKey?: string | null;
+};
+
 type SessionDefaultsSnapshot = {
   defaultAgentId?: string;
   mainKey?: string;
@@ -409,8 +413,26 @@ function handleChatGatewayEvent(host: GatewayHost, payload: ChatEventPayload | u
   }
   const state = handleChatEvent(host as unknown as ChatState, payload);
   const historyReloaded = handleTerminalChatEvent(host, payload, state);
+  const deferredReloadHost = host as GatewayHostWithDeferredSessionMessageReload;
+  const deferredSessionKey = deferredReloadHost.pendingSessionMessageReloadSessionKey?.trim();
+  const payloadSessionKey = payload?.sessionKey?.trim();
+  const shouldReplayDeferredSessionMessageReload = Boolean(
+    deferredSessionKey &&
+    payloadSessionKey &&
+    deferredSessionKey === payloadSessionKey &&
+    isTerminalChatState(state) &&
+    payloadSessionKey === host.sessionKey &&
+    !host.chatRunId,
+  );
+  if (deferredSessionKey && payloadSessionKey && deferredSessionKey === payloadSessionKey) {
+    deferredReloadHost.pendingSessionMessageReloadSessionKey = null;
+  }
   if (state === "final" && !historyReloaded && shouldReloadHistoryForFinalEvent(payload)) {
     void loadChatHistory(host as unknown as ChatState);
+    return;
+  }
+  if (shouldReplayDeferredSessionMessageReload && !historyReloaded) {
+    void loadChatHistory(host as unknown as ChatState);
   }
 }
 
@@ -418,6 +440,7 @@ function handleSessionMessageGatewayEvent(
   host: GatewayHost,
   payload: { sessionKey?: string } | undefined,
 ) {
+  const deferredReloadHost = host as GatewayHostWithDeferredSessionMessageReload;
   const sessionKey = payload?.sessionKey?.trim();
   if (!sessionKey || sessionKey !== host.sessionKey) {
     return;
@@ -428,8 +451,10 @@ function handleSessionMessageGatewayEvent(
   // chatStream, which delays the user message card from appearing until the
   // first LLM delta arrives.
   if (host.chatRunId) {
+    deferredReloadHost.pendingSessionMessageReloadSessionKey = sessionKey;
     return;
   }
+  deferredReloadHost.pendingSessionMessageReloadSessionKey = null;
   void loadChatHistory(host as unknown as ChatState);
 }