🤖 fix: keep init logs visible when returning to a workspace (#3139)

## Summary

Keep SSH/Coder init logs visible when returning to a workspace by making
running `workspace-init` replay visible immediately and by letting
running init-only transcripts bypass the generic workspace-loading
placeholder.

## Background

The earlier fix only handled deferred transcript rendering. The more
structural issue was that `WorkspaceStore` buffered init replay until
`caught-up`, while `WorkspaceShell` still treated init-only workspaces
as "loading" because there were no persisted chat messages yet. That
meant switching away from and back to a workspace mid-init could still
hide the running init hook and its SSH setup output behind a loading
screen.

## Implementation

- immediately apply replayed `init-start` / `init-output` / `init-end`
events during pre-`caught-up` buffering, just like other
reconnect-visible startup state
- treat actively running `workspace-init` rows as authoritative
transcript content when deriving `loading` / `isHydratingTranscript`
- keep the deferred-render guard for running init rows so reconnects do
not reuse stale snapshots
- add regression coverage for switching away from and back to a
workspace while SSH init output is still streaming

## Validation

- `bun test src/browser/stores/WorkspaceStore.test.ts`
- `bun test src/browser/utils/messages/messageUtils.test.ts`
- `make static-check`

## Risks

Low to moderate. The change is scoped to reconnect/catch-up behavior in
the renderer, but it touches the buffering rules that decide which
pre-stream state is visible before `caught-up`.

---

_Generated with `mux` • Model: `openai:gpt-5.4` • Thinking: `xhigh` •
Cost: `$5.06`_

<!-- mux-attribution: model=openai:gpt-5.4 thinking=xhigh costs=5.06 -->

Author: ammar-agent

src/browser/stores/WorkspaceStore.test.ts

@@ -1,4 +1,5 @@
 import { describe, expect, it, beforeEach, afterEach, mock, type Mock } from "bun:test";
+import type { DisplayedMessage } from "@/common/types/message";
 import type { FrontendWorkspaceMetadata } from "@/common/types/workspace";
 import type { StreamStartEvent, ToolCallStartEvent } from "@/common/types/stream";
 import type { WorkspaceActivitySnapshot, WorkspaceChatMessage } from "@/common/orpc/types";
@@ -1662,6 +1663,133 @@ describe("WorkspaceStore", () => {
       expect(stayedVisibleAfterCaughtUp).toBe(true);
     });
 
+    it("shows replayed init output before caught-up when switching back to a workspace", async () => {
+      const workspaceId = "workspace-init-replay";
+      const otherWorkspaceId = "workspace-init-other";
+      const firstLine = "Preparing workspace...";
+      const replayedLine = "Syncing repository over SSH...";
+      let subscriptionCount = 0;
+      let releaseSecondInitOutput: (() => void) | undefined;
+      let releaseSecondCaughtUp: (() => void) | undefined;
+
+      const getInitMessage = (): {
+        state: ReturnType<WorkspaceStore["getWorkspaceState"]>;
+        initMessage: Extract<DisplayedMessage, { type: "workspace-init" }> | undefined;
+      } => {
+        const state = store.getWorkspaceState(workspaceId);
+        const initMessage = state.messages.find(
+          (message): message is Extract<DisplayedMessage, { type: "workspace-init" }> =>
+            message.type === "workspace-init"
+        );
+        return { state, initMessage };
+      };
+
+      mockOnChat.mockImplementation(async function* (
+        input?: { workspaceId: string; mode?: unknown },
+        options?: { signal?: AbortSignal }
+      ): AsyncGenerator<WorkspaceChatMessage, void, unknown> {
+        if (input?.workspaceId !== workspaceId) {
+          await waitForAbortSignal(options?.signal);
+          return;
+        }
+
+        subscriptionCount += 1;
+
+        if (subscriptionCount === 1) {
+          yield { type: "caught-up" };
+          await Promise.resolve();
+          yield {
+            type: "init-start",
+            hookPath: "/tmp/project/.mux/init",
+            timestamp: 1_000,
+          };
+          await Promise.resolve();
+          yield {
+            type: "init-output",
+            line: firstLine,
+            isError: false,
+            timestamp: 1_001,
+          };
+          await waitForAbortSignal(options?.signal);
+          return;
+        }
+
+        yield {
+          type: "init-start",
+          hookPath: "/tmp/project/.mux/init",
+          timestamp: 2_000,
+        };
+        await new Promise<void>((resolve) => {
+          releaseSecondInitOutput = resolve;
+        });
+        yield {
+          type: "init-output",
+          line: replayedLine,
+          isError: false,
+          timestamp: 2_001,
+        };
+        await new Promise<void>((resolve) => {
+          releaseSecondCaughtUp = resolve;
+        });
+        yield { type: "caught-up", replay: "full" };
+        await waitForAbortSignal(options?.signal);
+      });
+
+      createAndAddWorkspace(store, workspaceId);
+
+      const sawInitialInit = await waitUntil(() => {
+        const { state, initMessage } = getInitMessage();
+        return (
+          state.loading === false &&
+          initMessage?.status === "running" &&
+          initMessage.lines[0]?.line === firstLine
+        );
+      });
+      expect(sawInitialInit).toBe(true);
+
+      createAndAddWorkspace(store, otherWorkspaceId);
+      store.setActiveWorkspaceId(workspaceId);
+
+      const sawEmptyReconnectInit = await waitUntil(() => {
+        const { state, initMessage } = getInitMessage();
+        return (
+          subscriptionCount >= 2 &&
+          state.loading === false &&
+          state.isHydratingTranscript === false &&
+          initMessage?.status === "running" &&
+          initMessage.lines.length === 0
+        );
+      });
+      expect(sawEmptyReconnectInit).toBe(true);
+
+      releaseSecondInitOutput?.();
+
+      const replayedInitBeforeCaughtUp = await waitUntil(() => {
+        const { state, initMessage } = getInitMessage();
+        return (
+          subscriptionCount >= 2 &&
+          state.loading === false &&
+          state.isHydratingTranscript === false &&
+          initMessage?.status === "running" &&
+          initMessage.lines.length === 1 &&
+          initMessage.lines[0]?.line === replayedLine
+        );
+      });
+      expect(replayedInitBeforeCaughtUp).toBe(true);
+
+      releaseSecondCaughtUp?.();
+
+      const stayedVisibleAfterCaughtUp = await waitUntil(() => {
+        const { state, initMessage } = getInitMessage();
+        return (
+          !state.loading &&
+          initMessage?.status === "running" &&
+          initMessage.lines[0]?.line === replayedLine
+        );
+      });
+      expect(stayedVisibleAfterCaughtUp).toBe(true);
+    });
+
     it("active workspace still shows starting during legitimate startup gap", async () => {
       const workspaceId = "stream-starting-active-workspace";
 

src/browser/stores/WorkspaceStore.ts

@@ -32,6 +32,9 @@ import {
   isStreamError,
   isStreamLifecycle,
   isDeleteMessage,
+  isInitEnd,
+  isInitOutput,
+  isInitStart,
   isBashOutputEvent,
   isTaskCreatedEvent,
   isMuxMessage,
@@ -1525,6 +1528,10 @@ export class WorkspaceStore {
       const aggregator = this.assertGet(workspaceId);
 
       const hasMessages = aggregator.hasMessages();
+      const displayedMessages = aggregator.getDisplayedMessages();
+      const hasRunningInitMessage = displayedMessages.some(
+        (message) => message.type === "workspace-init" && message.status === "running"
+      );
       const transient = this.assertChatTransientState(workspaceId);
       const historyPagination =
         this.historyPagination.get(workspaceId) ?? createInitialHistoryPaginationState();
@@ -1574,8 +1581,14 @@ export class WorkspaceStore {
         (streamLifecycle?.phase === "preparing" ||
           (!hasAuthoritativeStreamLifecycle && pendingStreamStartTime !== null)) &&
         !canInterrupt;
+      // Only actively running init output should bypass transcript hydration. Completed init
+      // rows are still replayed, but they should not suppress the normal catch-up placeholder
+      // for stale cached transcript content on reconnect.
       const isHydratingTranscript =
-        isActiveWorkspace && transient.isHydratingTranscript && !transient.caughtUp;
+        isActiveWorkspace &&
+        transient.isHydratingTranscript &&
+        !transient.caughtUp &&
+        !hasRunningInitMessage;
       const aggregatorTodos = aggregator.getCurrentTodos();
       const displayStatus = useAggregatorState ? undefined : (activity?.displayStatus ?? undefined);
       const todoStatus = useAggregatorState
@@ -1596,13 +1609,13 @@ export class WorkspaceStore {
 
       return {
         name: metadata?.name ?? workspaceId, // Fall back to ID if metadata missing
-        messages: aggregator.getDisplayedMessages(),
+        messages: displayedMessages,
         queuedMessage: transient.queuedMessage,
         canInterrupt,
         isCompacting: aggregator.isCompacting(),
         isStreamStarting,
         awaitingUserQuestion: aggregator.hasAwaitingUserQuestion(),
-        loading: !hasMessages && !transient.caughtUp,
+        loading: !hasMessages && !hasRunningInitMessage && !transient.caughtUp,
         isHydratingTranscript,
         hasOlderHistory: historyPagination.hasOlder,
         loadingOlderHistory: historyPagination.loading,
@@ -3577,9 +3590,24 @@ export class WorkspaceStore {
     }
 
     if (!transient.caughtUp && this.isBufferedEvent(data)) {
-      if (isStreamLifecycle(data) || isStreamAbort(data) || isRuntimeStatus(data)) {
+      if (
+        isStreamLifecycle(data) ||
+        isStreamAbort(data) ||
+        isRuntimeStatus(data) ||
+        isInitStart(data) ||
+        isInitOutput(data) ||
+        isInitEnd(data)
+      ) {
+        // SSH/Coder init replay can be the only transcript content for a new workspace.
+        // Apply it immediately so switching back mid-init shows the latest backend output
+        // instead of only a generic loading placeholder until caught-up lands.
         applyWorkspaceChatEventToAggregator(aggregator, data, { allowSideEffects: false });
-        this.states.bump(workspaceId);
+        if (isInitOutput(data)) {
+          aggregator.flushPendingInitOutput();
+          this.scheduleIdleStateBump(workspaceId);
+        } else {
+          this.states.bump(workspaceId);
+        }
       }
 
       transient.pendingStreamEvents.push(data);

src/browser/utils/messages/StreamingMessageAggregator.ts

@@ -1595,6 +1595,20 @@ export class StreamingMessageAggregator {
     }
   }
 
+  /**
+   * Replay-visible init output needs a synchronous cache flush. WorkspaceStore can
+   * bump subscribers before the normal 100ms init-output throttle fires, and in
+   * reconnects that would otherwise leave the newest line hidden until caught-up.
+   */
+  flushPendingInitOutput(): void {
+    if (this.initOutputThrottleTimer) {
+      clearTimeout(this.initOutputThrottleTimer);
+      this.initOutputThrottleTimer = null;
+    }
+
+    this.invalidateCache();
+  }
+
   clearPendingStreamStart(): void {
     this.setPendingStreamStartTime(null);
   }

src/browser/utils/messages/messageUtils.test.ts

@@ -113,6 +113,25 @@ describe("shouldBypassDeferredMessages", () => {
     result: { success: true, output: "hi", exitCode: 0, wall_duration_ms: 5 },
   };
 
+  const runningInit: DisplayedMessage = {
+    type: "workspace-init",
+    id: "workspace-init",
+    historySequence: -1,
+    status: "running",
+    hookPath: "/tmp/project/.mux/init",
+    lines: [{ line: "Installing dependencies...", isError: false }],
+    exitCode: null,
+    timestamp: 1,
+    durationMs: null,
+  };
+
+  const completedInit: DisplayedMessage = {
+    ...runningInit,
+    status: "success",
+    exitCode: 0,
+    durationMs: 2_000,
+  };
+
   it("returns true when immediate snapshot has active rows", () => {
     expect(shouldBypassDeferredMessages([executingBash], [executingBash])).toBe(true);
   });
@@ -141,6 +160,16 @@ describe("shouldBypassDeferredMessages", () => {
     );
   });
 
+  it("returns true when init output is still running", () => {
+    expect(shouldBypassDeferredMessages([runningInit], [runningInit])).toBe(true);
+  });
+
+  it("returns true when the deferred snapshot still shows a running init hook", () => {
+    // Regression scenario: reconnect replay completed the init hook, but the deferred
+    // snapshot is still holding on to the older running row from before catch-up.
+    expect(shouldBypassDeferredMessages([completedInit], [runningInit])).toBe(true);
+  });
+
   it("returns false when both snapshots are settled and in sync", () => {
     expect(shouldBypassDeferredMessages([completedBash], [completedBash])).toBe(false);
   });

src/browser/utils/messages/messageUtils.ts

@@ -115,13 +115,15 @@ export function shouldShowInterruptedBarrier(
 
 /**
  * Returns whether ChatPane should bypass useDeferredValue and render the immediate
- * message list. We bypass deferral while assistant content is streaming OR while
- * any tool call is still executing (e.g. live bash output).
+ * message list. We bypass deferral while assistant content is streaming, while the
+ * init hook is still running, OR while any tool call is still executing (for example,
+ * live bash output).
  *
  * We also bypass when the deferred snapshot appears stale (it still has active
  * streaming/executing rows after the immediate snapshot is idle), or when both
  * snapshots have diverged in row identity/order. Showing stale deferred rows can
- * cause hidden-marker placement and tool-state flash at stream completion.
+ * cause hidden-marker placement, hide live init logs after a workspace switch, and
+ * flash tool state at stream completion.
  */
 export function shouldBypassDeferredMessages(
   messages: DisplayedMessage[],
@@ -130,7 +132,12 @@ export function shouldBypassDeferredMessages(
   const hasActiveRows = (rows: DisplayedMessage[]) =>
     rows.some(
       (m) =>
-        ("isStreaming" in m && m.isStreaming) || (m.type === "tool" && m.status === "executing")
+        ("isStreaming" in m && m.isStreaming) ||
+        (m.type === "tool" && m.status === "executing") ||
+        // Keep SSH/Coder init output on the immediate path when a user returns to a
+        // workspace mid-setup; otherwise the deferred snapshot can keep showing an
+        // older empty/stale init row because workspace-init uses a stable display ID.
+        (m.type === "workspace-init" && m.status === "running")
     );
 
   if (messages.length !== deferredMessages.length) {