mstro-app 0.4.52 → 0.5.1

package/server/services/websocket/git-worktree-handlers.ts

@@ -1,13 +1,26 @@
 // Copyright (c) 2025-present Mstro, Inc. All rights reserved.
-// Licensed under the MIT License. See LICENSE file for details.
 
 import { existsSync, readFileSync, writeFileSync } from 'node:fs';
 import { dirname, join } from 'node:path';
 import { resolvePmDir } from '../plan/parser.js';
 import type { Workspace } from '../plan/types.js';
 import { executeGitCommand, handleGitStatus, spawnWithOutput } from './git-handlers.js';
+import { handleGitLog } from './git-log-handlers.js';
+import { parseGitStatus } from './git-utils.js';
 import type { HandlerContext } from './handler-context.js';
-import type { WebSocketMessage, WorktreeInfo, WSContext } from './types.js';
+import type { GitFileStatus, WebSocketMessage, WorktreeInfo, WSContext } from './types.js';
+
+/**
+ * A file in the merge target's working tree whose uncommitted state would be
+ * lost if the merge proceeds. Used to drive the merge dialog's blocker state
+ * before the merge attempt and to translate the post-flight git error if a
+ * race lets one slip through.
+ */
+interface MergeBlocker {
+  path: string;
+  status: GitFileStatus['status'];
+  staged: boolean;
+}
 
 function persistBoardWorktree(workingDir: string, boardId: string, worktreePath: string | null, branch: string | null): void {
   const pmDir = resolvePmDir(workingDir);
@@ -43,6 +56,8 @@ export async function handleGitWorktreeMessage(ctx: HandlerContext, ws: WSContex
     gitWorktreeMerge: () => handleGitWorktreeMerge(ctx, ws, msg, tabId, gitDir),
     gitMergeAbort: () => handleGitMergeAbort(ctx, ws, tabId, gitDir),
     gitMergeComplete: () => handleGitMergeComplete(ctx, ws, msg, tabId, gitDir),
+    gitMergeStashPop: () => handleGitMergeStashPop(ctx, ws, msg, tabId, gitDir),
+    gitMergeDiscardBlockers: () => handleGitMergeDiscardBlockers(ctx, ws, msg, tabId, gitDir),
   };
   await handlers[msg.type]?.();
 }
@@ -214,7 +229,7 @@ export async function handleTabWorktreeSwitch(ctx: HandlerContext, ws: WSContext
         persistBoardWorktree(workingDir, resolvedTabId, null, null);
       }
       ctx.send(ws, { type: 'tabWorktreeSwitched', tabId: resolvedTabId, data: { tabId: resolvedTabId, worktreePath: null, branch: null } });
-      handleGitStatus(ctx, ws, resolvedTabId, workingDir);
+      refreshScopeAfterWorktreeSwitch(ctx, ws, resolvedTabId, workingDir);
       return;
     }
 
@@ -229,12 +244,38 @@ export async function handleTabWorktreeSwitch(ctx: HandlerContext, ws: WSContext
     }
 
     ctx.send(ws, { type: 'tabWorktreeSwitched', tabId: resolvedTabId, data: { tabId: resolvedTabId, worktreePath, branch } });
-    handleGitStatus(ctx, ws, resolvedTabId, worktreePath);
+    refreshScopeAfterWorktreeSwitch(ctx, ws, resolvedTabId, worktreePath);
   } catch (error: unknown) {
     ctx.send(ws, { type: 'gitError', tabId, data: { error: error instanceof Error ? error.message : String(error) } });
   }
 }
 
+/**
+ * After a worktree switch, re-fetch everything that's worktree-specific so
+ * the client sees a complete, consistent view of the newly-selected workspace
+ * from a single `tabWorktreeSwitched` event. Keeping the refresh on the
+ * server side means the client has one signal to react to instead of having
+ * to orchestrate status/log/... fetches itself.
+ *
+ * Branches and the worktree list are NOT re-fetched: they're repo-wide, not
+ * worktree-specific.
+ *
+ * Fire-and-forget: the switch itself has already been acknowledged via
+ * `tabWorktreeSwitched`. `handleGitStatus` and `handleGitLog` each own their
+ * own error handling (they send `gitError` scoped to the correct tabId).
+ * Awaiting here would let any unexpected throw escape into the caller's
+ * outer try/catch and produce a misleading `gitError` on the original
+ * dispatch tabId after success has already been signalled.
+ */
+function refreshScopeAfterWorktreeSwitch(ctx: HandlerContext, ws: WSContext, tabId: string, gitDir: string): void {
+  (async () => {
+    await handleGitStatus(ctx, ws, tabId, gitDir);
+    await handleGitLog(ctx, ws, { type: 'gitLog', tabId, data: { limit: 20 } }, tabId, gitDir);
+  })().catch((error: unknown) => {
+    console.error('[handleTabWorktreeSwitch] scope refresh failed:', error);
+  });
+}
+
 async function pushWithUpstreamRetry(
   worktreePath: string,
   pushRemote: string,
@@ -312,6 +353,64 @@ async function handleGitWorktreeCreatePR(ctx: HandlerContext, ws: WSContext, msg
   }
 }
 
+/**
+ * Files in the target worktree's working tree whose dirty state intersects
+ * the set of files the merge would touch. These are the rows that drive the
+ * merge dialog's "uncommitted changes on <target>" blocker UI — surfacing
+ * them preflight prevents the user from ever hitting git's raw "would be
+ * overwritten" error.
+ *
+ * Detection strategy:
+ *   - `git status --porcelain=v1` on the target's worktree → all dirty paths
+ *   - `git diff --name-only target..source` → all paths the merge would touch
+ *   - intersection = blockers
+ *
+ * Untracked files are included via the porcelain output (they collide with
+ * incoming additions of the same path). Files that are dirty but outside the
+ * merge's diff are NOT blockers — git would carry them through cleanly.
+ */
+async function detectMergeBlockers(
+  targetWorktreePath: string,
+  sourceBranch: string,
+  targetBranch: string,
+): Promise<MergeBlocker[]> {
+  const statusResult = await executeGitCommand(['status', '--porcelain=v1'], targetWorktreePath);
+  if (statusResult.exitCode !== 0) return [];
+  const { staged, unstaged, untracked } = parseGitStatus(statusResult.stdout);
+
+  const diffResult = await executeGitCommand(
+    ['diff', '--name-only', `${targetBranch}..${sourceBranch}`],
+    targetWorktreePath,
+  );
+  if (diffResult.exitCode !== 0) return [];
+  const mergeTouches = new Set(
+    diffResult.stdout.split('\n').map(s => s.trim()).filter(Boolean),
+  );
+  if (mergeTouches.size === 0) return [];
+
+  // Prefer the worktree (unstaged) status when both index and worktree are
+  // dirty — that's what the user would lose. Fall back to staged, then
+  // untracked. Order in the result list groups deletions last so the UI can
+  // surface them where users expect.
+  const seen = new Map<string, MergeBlocker>();
+  const consider = (entries: GitFileStatus[]) => {
+    for (const entry of entries) {
+      if (!mergeTouches.has(entry.path)) continue;
+      if (seen.has(entry.path)) continue;
+      seen.set(entry.path, { path: entry.path, status: entry.status, staged: entry.staged });
+    }
+  };
+  consider(unstaged);
+  consider(staged);
+  consider(untracked);
+
+  return Array.from(seen.values()).sort((a, b) => {
+    if (a.status === 'D' && b.status !== 'D') return 1;
+    if (a.status !== 'D' && b.status === 'D') return -1;
+    return a.path.localeCompare(b.path);
+  });
+}
+
 async function handleGitMergePreview(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, workingDir: string): Promise<void> {
   try {
     const { sourceBranch, targetBranch } = msg.data || {};
@@ -347,10 +446,30 @@ async function handleGitMergePreview(ctx: HandlerContext, ws: WSContext, msg: We
         return { hash: hash.trim(), message: rest.join('|').trim() };
       });
 
+    // The merge always lands on the main worktree currently checked out to
+    // the target branch. If main isn't on `targetBranch` we can't preflight
+    // blockers (the dirty state wouldn't be the relevant one), so skip the
+    // check and let the post-flight handler surface the branch-mismatch
+    // error like before.
+    const mainPath = await resolveMainWorktreePath(workingDir);
+    const mainBranchResult = await executeGitCommand(['rev-parse', '--abbrev-ref', 'HEAD'], mainPath);
+    const mainOnTarget = mainBranchResult.stdout.trim() === targetBranch;
+    const blockers = mainOnTarget
+      ? await detectMergeBlockers(mainPath, sourceBranch, targetBranch)
+      : [];
+
     ctx.send(ws, {
       type: 'gitMergePreviewResult',
       tabId,
-      data: { clean, conflicts, stat, commits, ahead: commits.length },
+      data: {
+        clean,
+        conflicts,
+        stat,
+        commits,
+        ahead: commits.length,
+        targetWorktreePath: mainPath,
+        targetWorktreeBlockers: blockers,
+      },
     });
   } catch (error: unknown) {
     ctx.send(ws, { type: 'gitError', tabId, data: { error: error instanceof Error ? error.message : String(error) } });
@@ -455,9 +574,103 @@ export function findWorktreePathForBranch(porcelainOutput: string, branchName: s
   return null;
 }
 
+/**
+ * Parse a git stderr buffer for the "your local changes / untracked working
+ * tree files would be overwritten by merge" message. Returned paths feed the
+ * web client's blocker UI when a merge slips through preflight (race against
+ * a teammate's auto-formatter, stale status, etc.).
+ */
+function parseOverwritePaths(stderr: string): string[] {
+  const lines = stderr.split('\n');
+  const headerIdx = lines.findIndex(l =>
+    /your local changes to the following files would be overwritten/i.test(l) ||
+    /untracked working tree files would be overwritten/i.test(l),
+  );
+  if (headerIdx < 0) return [];
+  const paths: string[] = [];
+  for (let i = headerIdx + 1; i < lines.length; i++) {
+    const line = lines[i];
+    const trimmed = line.trim();
+    if (!trimmed) continue;
+    if (/^(merge with strategy|aborting|please commit)/i.test(trimmed)) break;
+    // Git indents each path with a tab; tolerate spaces too.
+    if (/^[\t ]/.test(line)) {
+      paths.push(trimmed);
+    }
+  }
+  return paths;
+}
+
+const STASH_MESSAGE_TAG = 'mstro:pre-merge';
+
+async function pushMergeStash(targetWorktreePath: string, sourceBranch: string, targetBranch: string): Promise<{ ref: string; sha: string; message: string } | null> {
+  const message = `${STASH_MESSAGE_TAG} ${targetBranch} <- ${sourceBranch} @ ${new Date().toISOString()}`;
+  const result = await executeGitCommand(['stash', 'push', '-u', '-m', message], targetWorktreePath);
+  if (result.exitCode !== 0) return null;
+  // "No local changes to save" exits 0 but does not push a stash; detect by
+  // checking whether stash@{0}'s message matches what we just wrote.
+  const top = await executeGitCommand(['stash', 'list', '-n', '1', '--format=%H%x09%gs'], targetWorktreePath);
+  if (top.exitCode !== 0) return null;
+  const [sha, ...rest] = top.stdout.trim().split('\t');
+  const stashedMsg = rest.join('\t');
+  if (!sha || !stashedMsg.includes(STASH_MESSAGE_TAG)) return null;
+  return { ref: 'stash@{0}', sha, message };
+}
+
+async function popMergeStashBySha(targetWorktreePath: string, sha: string): Promise<{ exitCode: number; error?: string }> {
+  const list = await executeGitCommand(['stash', 'list', '--format=%gd%x09%H'], targetWorktreePath);
+  if (list.exitCode !== 0) return { exitCode: list.exitCode, error: list.stderr || 'Failed to read stash list' };
+  const ref = list.stdout.split('\n').map(l => l.split('\t')).find(([, h]) => h === sha)?.[0];
+  if (!ref) return { exitCode: 1, error: 'Stash no longer exists — it may have already been popped or dropped.' };
+  const pop = await executeGitCommand(['stash', 'pop', ref], targetWorktreePath);
+  return { exitCode: pop.exitCode, error: pop.exitCode !== 0 ? pop.stderr : undefined };
+}
+
+/**
+ * Decide what shape of merge-failure payload to send the client. Splits the
+ * three error modes — index conflict, target-worktree blocker, generic — out
+ * of the main handler so each is a self-contained branch.
+ */
+async function buildMergeFailurePayload(
+  mainPath: string,
+  sourceBranch: string,
+  targetBranch: string,
+  mergeError: string | undefined,
+): Promise<Record<string, unknown>> {
+  const conflictFiles = await detectMergeConflicts(mainPath);
+  if (conflictFiles.length > 0) {
+    return { success: false, conflictFiles, targetWorktreePath: mainPath };
+  }
+  const overwritten = parseOverwritePaths(mergeError || '');
+  if (overwritten.length > 0) {
+    const blockers = await detectMergeBlockers(mainPath, sourceBranch, targetBranch);
+    const targetWorktreeBlockers: MergeBlocker[] = blockers.length > 0
+      ? blockers
+      : overwritten.map(p => ({ path: p, status: 'M' as const, staged: false }));
+    return { success: false, targetWorktreePath: mainPath, targetWorktreeBlockers };
+  }
+  return { success: false, error: mergeError || 'Merge failed', targetWorktreePath: mainPath };
+}
+
+function buildMergeSuccessPayload(
+  mainPath: string,
+  mergeCommit: string,
+  warnings: string[],
+  stash: { ref: string; sha: string; message: string } | null,
+): Record<string, unknown> {
+  const data: Record<string, unknown> = { success: true, mergeCommit, targetWorktreePath: mainPath };
+  if (warnings.length > 0) data.warnings = warnings;
+  if (stash) {
+    data.stashRef = stash.ref;
+    data.stashSha = stash.sha;
+    data.stashMessage = stash.message;
+  }
+  return data;
+}
+
 async function handleGitWorktreeMerge(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, workingDir: string): Promise<void> {
   try {
-    const { sourceBranch, targetBranch, strategy, commitMessage, deleteWorktree, deleteBranch } = msg.data || {};
+    const { sourceBranch, targetBranch, strategy, commitMessage, deleteWorktree, deleteBranch, stashFirst } = msg.data || {};
     if (!sourceBranch || !targetBranch || !strategy) {
       ctx.send(ws, { type: 'gitError', tabId, data: { error: 'Source branch, target branch, and strategy are required' } });
       return;
@@ -467,40 +680,97 @@ async function handleGitWorktreeMerge(ctx: HandlerContext, ws: WSContext, msg: W
 
     const mainBranchResult = await executeGitCommand(['rev-parse', '--abbrev-ref', 'HEAD'], mainPath);
     if (mainBranchResult.stdout.trim() !== targetBranch) {
-      ctx.send(ws, { type: 'gitWorktreeMergeResult', tabId, data: { success: false, error: `Switch the main worktree to "${targetBranch}" before merging` } });
+      ctx.send(ws, { type: 'gitWorktreeMergeResult', tabId, data: { success: false, error: `Switch the main worktree to "${targetBranch}" before merging`, targetWorktreePath: mainPath } });
       return;
     }
 
-    const headBefore = await executeGitCommand(['rev-parse', 'HEAD'], mainPath);
+    // Stash first if requested. A no-op result (nothing to stash, e.g. user
+    // already cleaned up between preflight and merge) is fine — we still
+    // proceed with the merge.
+    const stash = stashFirst ? await pushMergeStash(mainPath, sourceBranch, targetBranch) : null;
 
+    const headBefore = await executeGitCommand(['rev-parse', 'HEAD'], mainPath);
     const mergeResult = await executeMergeStrategy(strategy, sourceBranch, commitMessage, mainPath);
+
     if (mergeResult.exitCode !== 0) {
-      const conflictFiles = await detectMergeConflicts(mainPath);
-      const data = conflictFiles.length > 0
-        ? { success: false, conflictFiles }
-        : { success: false, error: mergeResult.error || 'Merge failed' };
+      // Best-effort restore of the user's working tree before reporting.
+      // If pop conflicts, the stash stays around and the user can recover it
+      // manually — we still surface the original merge error.
+      if (stash) await popMergeStashBySha(mainPath, stash.sha).catch(() => undefined);
+      const data = await buildMergeFailurePayload(mainPath, sourceBranch, targetBranch, mergeResult.error);
       ctx.send(ws, { type: 'gitWorktreeMergeResult', tabId, data });
       return;
     }
 
     const headAfter = await executeGitCommand(['rev-parse', 'HEAD'], mainPath);
     if (headBefore.stdout.trim() === headAfter.stdout.trim()) {
-      ctx.send(ws, { type: 'gitWorktreeMergeResult', tabId, data: { success: false, error: `Already up to date — "${sourceBranch}" has no new commits to merge into "${targetBranch}"` } });
+      if (stash) await popMergeStashBySha(mainPath, stash.sha).catch(() => undefined);
+      ctx.send(ws, { type: 'gitWorktreeMergeResult', tabId, data: { success: false, error: `Already up to date — "${sourceBranch}" has no new commits to merge into "${targetBranch}"`, targetWorktreePath: mainPath } });
       return;
     }
 
     const commitHashResult = await executeGitCommand(['rev-parse', '--short', 'HEAD'], mainPath);
     const { warnings, removedWorktreePath } = await cleanupAfterMerge(mainPath, sourceBranch, strategy, !!deleteWorktree, !!deleteBranch);
+    if (removedWorktreePath) cleanupWorktreeReferences(ctx, workingDir, removedWorktreePath);
 
-    if (removedWorktreePath) {
-      cleanupWorktreeReferences(ctx, workingDir, removedWorktreePath);
+    const data = buildMergeSuccessPayload(mainPath, commitHashResult.stdout.trim(), warnings, stash);
+    ctx.send(ws, { type: 'gitWorktreeMergeResult', tabId, data });
+  } catch (error: unknown) {
+    ctx.send(ws, { type: 'gitError', tabId, data: { error: error instanceof Error ? error.message : String(error) } });
+  }
+}
+
+async function handleGitMergeStashPop(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, workingDir: string): Promise<void> {
+  try {
+    const { stashSha, targetWorktreePath } = msg.data || {};
+    if (!stashSha) {
+      ctx.send(ws, { type: 'gitMergeStashPopped', tabId, data: { success: false, error: 'Missing stash reference' } });
+      return;
     }
+    const path = targetWorktreePath || (await resolveMainWorktreePath(workingDir));
+    const result = await popMergeStashBySha(path, stashSha);
+    if (result.exitCode !== 0) {
+      ctx.send(ws, { type: 'gitMergeStashPopped', tabId, data: { success: false, error: result.error || 'Failed to pop stash', targetWorktreePath: path } });
+      return;
+    }
+    ctx.send(ws, { type: 'gitMergeStashPopped', tabId, data: { success: true, targetWorktreePath: path } });
+    // Refresh status so the file explorer / git view reflects the restored changes.
+    handleGitStatus(ctx, ws, tabId, path);
+  } catch (error: unknown) {
+    ctx.send(ws, { type: 'gitError', tabId, data: { error: error instanceof Error ? error.message : String(error) } });
+  }
+}
 
-    const data: Record<string, unknown> = { success: true, mergeCommit: commitHashResult.stdout.trim() };
-    if (warnings.length > 0) {
-      data.warnings = warnings;
+async function handleGitMergeDiscardBlockers(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, workingDir: string): Promise<void> {
+  try {
+    const { paths, targetWorktreePath } = msg.data || {};
+    if (!Array.isArray(paths) || paths.length === 0) {
+      ctx.send(ws, { type: 'gitMergeBlockersDiscarded', tabId, data: { success: false, error: 'No paths to discard' } });
+      return;
     }
-    ctx.send(ws, { type: 'gitWorktreeMergeResult', tabId, data });
+    const path = targetWorktreePath || (await resolveMainWorktreePath(workingDir));
+
+    // Restore tracked files (modifications, deletions, stages) from HEAD.
+    // `git checkout HEAD -- <paths>` overwrites both index and worktree, which
+    // is exactly the "discard" semantics the UI promised the user.
+    const checkoutResult = await executeGitCommand(['checkout', 'HEAD', '--', ...paths], path);
+    // Untracked files won't be touched by `checkout HEAD --` (git refuses with
+    // "did not match any file(s) known to git" when ALL paths are untracked,
+    // but mixed lists succeed for the tracked subset). Run `git clean -f` on
+    // the originally-supplied list to remove any untracked leftovers — git
+    // ignores anything that's already gone, so this is safe to run after.
+    const cleanResult = await executeGitCommand(['clean', '-f', '--', ...paths], path);
+
+    if (checkoutResult.exitCode !== 0 && cleanResult.exitCode !== 0) {
+      ctx.send(ws, {
+        type: 'gitMergeBlockersDiscarded',
+        tabId,
+        data: { success: false, error: checkoutResult.stderr || cleanResult.stderr || 'Failed to discard changes', targetWorktreePath: path },
+      });
+      return;
+    }
+    ctx.send(ws, { type: 'gitMergeBlockersDiscarded', tabId, data: { success: true, targetWorktreePath: path, paths } });
+    handleGitStatus(ctx, ws, tabId, path);
   } catch (error: unknown) {
     ctx.send(ws, { type: 'gitError', tabId, data: { error: error instanceof Error ? error.message : String(error) } });
   }

package/server/services/websocket/handler-context.ts

@@ -1,13 +1,14 @@
 // Copyright (c) 2025-present Mstro, Inc. All rights reserved.
-// Licensed under the MIT License. See LICENSE file for details.
 
 import type { ChildProcess } from 'node:child_process';
 import type { ImprovisationSessionManager } from '../../cli/improvisation-session-manager.js';
 import type { AutocompleteService } from './autocomplete.js';
 import type { FileUploadHandler } from './file-upload-handler.js';
 import type { GitHeadWatcher } from './git-head-watcher.js';
+import type { MsgIdTracker } from './msg-id-tracker.js';
 import type { SessionRegistry } from './session-registry.js';
 import type { SkillsWatcher } from './skill-watcher.js';
+import type { TabEventBufferRegistry } from './tab-event-buffer.js';
 import type { WebSocketResponse, WSContext } from './types.js';
 
 export interface UsageReport {
@@ -37,6 +38,19 @@ export interface HandlerContext {
   fileUploadHandler: FileUploadHandler | null;
   gitHeadWatcher: GitHeadWatcher | null;
   skillsWatcher: SkillsWatcher | null;
+  /**
+   * Per-tab replay buffer for tab-scoped broadcasts. Populated by
+   * `broadcastTabEvent` (see `tab-broadcast.ts`) so a web client rejoining a
+   * tab after a reconnect can request replay of anything it missed during
+   * the transport gap. See `tab-event-buffer.ts`.
+   */
+  tabEventBuffers: TabEventBufferRegistry;
+  /**
+   * Idempotency tracker for `execute` `msgId`s. Lets the web replay the
+   * same prompt across reconnects without causing double execution — the
+   * CLI still acks, but skips running the prompt a second time.
+   */
+  msgIdTracker: MsgIdTracker;
 
   // Registry access
   getRegistry(workingDir: string): SessionRegistry;

package/server/services/websocket/handler.ts

@@ -1,5 +1,4 @@
 // Copyright (c) 2025-present Mstro, Inc. All rights reserved.
-// Licensed under the MIT License. See LICENSE file for details.
 
 /**
  * WebSocket Handler for Improvisation Sessions
@@ -14,20 +13,24 @@ import { homedir } from 'node:os';
 import { dirname, join } from 'node:path';
 import type { ImprovisationSessionManager } from '../../cli/improvisation-session-manager.js';
 import { captureException } from '../sentry.js';
+import { getPTYManager } from '../terminal/pty-manager.js';
 import { AutocompleteService } from './autocomplete.js';
+import { FileDownloadHandler } from './file-download-handler.js';
 import { handleFileExplorerMessage, handleFileMessage } from './file-explorer-handlers.js';
 import { FileUploadHandler } from './file-upload-handler.js';
 import { handleGitMessage } from './git-handlers.js';
 import { GitHeadWatcher } from './git-head-watcher.js';
 import type { HandlerContext, UsageReporter } from './handler-context.js';
+import { MsgIdTracker } from './msg-id-tracker.js';
 import { handlePlanMessage } from './plan-handlers.js';
 import { handleQualityMessage } from './quality-handlers.js';
-import { handleHistoryMessage, handleSessionMessage, initializeTab, resumeHistoricalSession } from './session-handlers.js';
+import { handleHistoryMessage, handleSessionMessage, initializeTab, restoreWorktreeFromRegistry, resumeHistoricalSession } from './session-handlers.js';
 import { SessionRegistry } from './session-registry.js';
 import { generateNotificationSummary, handleGetSettings, handleUpdateSettings } from './settings-handlers.js';
 import { handleListSkills } from './skill-handlers.js';
 import { SkillsWatcher } from './skill-watcher.js';
-import { handleCreateTab, handleGetActiveTabs, handleMarkTabViewed, handleRemoveTab, handleReorderTabs, handleSyncPromptText, handleSyncTabMeta } from './tab-handlers.js';
+import { TabEventBufferRegistry } from './tab-event-buffer.js';
+import { handleCreateTab, handleGetActiveTabs, handleMarkTabViewed, handleRemoveTab, handleReorderTabs, handleSyncTabMeta } from './tab-handlers.js';
 import { cleanupTerminalSubscribers, handleTerminalMessage } from './terminal-handlers.js';
 import type { FrecencyData, WebSocketMessage, WebSocketResponse, WSContext } from './types.js';
 
@@ -47,8 +50,11 @@ export class WebSocketImproviseHandler implements HandlerContext {
   terminalListenerCleanups: Map<string, () => void> = new Map();
   terminalSubscribers: Map<string, Set<WSContext>> = new Map();
   fileUploadHandler: FileUploadHandler | null = null;
+  fileDownloadHandler: FileDownloadHandler | null = null;
   gitHeadWatcher: GitHeadWatcher | null = null;
   skillsWatcher: SkillsWatcher | null = null;
+  tabEventBuffers: TabEventBufferRegistry = new TabEventBufferRegistry();
+  msgIdTracker: MsgIdTracker = new MsgIdTracker();
 
   constructor() {
     this.frecencyPath = join(homedir(), '.mstro', 'autocomplete-frecency.json');
@@ -146,7 +152,7 @@ export class WebSocketImproviseHandler implements HandlerContext {
   }
 
   /** Dispatch table mapping message types to domain handlers. Built once, looked up per message. */
-  private static readonly DISPATCH: Record<string, 'session' | 'history' | 'file' | 'terminal' | 'fileExplorer' | 'git' | 'quality' | 'plan' | 'fileUpload'> = {
+  private static readonly DISPATCH: Record<string, 'session' | 'history' | 'file' | 'terminal' | 'fileExplorer' | 'git' | 'quality' | 'plan' | 'fileUpload' | 'fileDownload'> = {
     // Session
     execute: 'session', cancel: 'session', getHistory: 'session', new: 'session', approve: 'session', reject: 'session',
     // History
@@ -165,6 +171,8 @@ export class WebSocketImproviseHandler implements HandlerContext {
     planInit: 'plan', planGetState: 'plan', planListIssues: 'plan', planGetIssue: 'plan', planGetSprint: 'plan', planGetMilestone: 'plan', planCreateIssue: 'plan', planUpdateIssue: 'plan', planDeleteIssue: 'plan', planScaffold: 'plan', planPrompt: 'plan', planExecute: 'plan', planExecuteEpic: 'plan', planPause: 'plan', planStop: 'plan', planResume: 'plan', planCreateBoard: 'plan', planUpdateBoard: 'plan', planArchiveBoard: 'plan', planGetBoard: 'plan', planGetBoardState: 'plan', planReorderBoards: 'plan', planSetActiveBoard: 'plan', planGetBoardArtifacts: 'plan', planCreateSprint: 'plan', planActivateSprint: 'plan', planCompleteSprint: 'plan', planGetSprintArtifacts: 'plan', chatToBoard: 'plan',
     // File upload
     fileUploadStart: 'fileUpload', fileUploadChunk: 'fileUpload', fileUploadComplete: 'fileUpload', fileUploadCancel: 'fileUpload',
+    // File download (chunked streaming for large binaries)
+    fileDownloadStart: 'fileDownload', fileDownloadCancel: 'fileDownload',
   };
 
   private async dispatchMessage(ws: WSContext, msg: WebSocketMessage, tabId: string, workingDir: string, permission?: 'view'): Promise<void> {
@@ -174,10 +182,10 @@ export class WebSocketImproviseHandler implements HandlerContext {
         this.send(ws, { type: 'pong', tabId });
         return;
       case 'initTab':
-        return void await initializeTab(this, ws, tabId, workingDir, msg.data?.tabName);
+        return void await initializeTab(this, ws, tabId, workingDir, msg.data?.tabName, msg.data);
       case 'resumeSession':
         if (!msg.data?.historicalSessionId) throw new Error('Historical session ID is required');
-        return void await resumeHistoricalSession(this, ws, tabId, workingDir, msg.data.historicalSessionId);
+        return void await resumeHistoricalSession(this, ws, tabId, workingDir, msg.data.historicalSessionId, msg.data);
       case 'requestNotificationSummary':
         if (!msg.data?.prompt || !msg.data?.output) throw new Error('Prompt and output are required for notification summary');
         return void await generateNotificationSummary(this, ws, tabId, msg.data.prompt, msg.data.output, workingDir);
@@ -189,8 +197,6 @@ export class WebSocketImproviseHandler implements HandlerContext {
         return handleReorderTabs(this, ws, workingDir, msg.data?.tabOrder);
       case 'syncTabMeta':
         return handleSyncTabMeta(this, ws, msg, tabId, workingDir);
-      case 'syncPromptText':
-        return handleSyncPromptText(this, ws, msg, tabId);
       case 'removeTab':
         return handleRemoveTab(this, ws, tabId, workingDir);
       case 'markTabViewed':
@@ -208,6 +214,14 @@ export class WebSocketImproviseHandler implements HandlerContext {
     const domain = WebSocketImproviseHandler.DISPATCH[msg.type];
     if (!domain) throw new Error(`Unknown message type: ${msg.type}`);
 
+    // Hydrate worktree state from the registry before any domain handler
+    // reads it, so git/file/autocomplete ops route to the tab's worktree
+    // even when they arrive before the initTab handshake completes. The
+    // registry is authoritative; the in-memory Map is just a cache.
+    if (msg.tabId && !this.gitDirectories.has(tabId)) {
+      restoreWorktreeFromRegistry(this, this.getRegistry(workingDir), tabId);
+    }
+
     // Resolve effective working directory: use worktree path if tab is on a worktree
     const effectiveDir = this.gitDirectories.get(tabId) || workingDir;
 
@@ -220,20 +234,38 @@ export class WebSocketImproviseHandler implements HandlerContext {
       case 'git':        return handleGitMessage(this, ws, msg, tabId, workingDir);
       case 'quality':    return handleQualityMessage(this, ws, msg, tabId, workingDir, permission);
       case 'plan':       return handlePlanMessage(this, ws, msg, tabId, workingDir, permission);
-      case 'fileUpload': return this.handleFileUploadMessage(ws, msg, tabId, workingDir, permission);
+      case 'fileUpload': return this.handleFileUploadMessage(ws, msg, tabId, effectiveDir, permission);
+      case 'fileDownload': return this.handleFileDownloadMessage(ws, msg, tabId, effectiveDir);
+    }
+  }
+
+  private handleFileDownloadMessage(ws: WSContext, msg: WebSocketMessage, tabId: string, workingDir: string): void {
+    if (!this.fileDownloadHandler) {
+      this.fileDownloadHandler = new FileDownloadHandler(workingDir);
+    }
+    const handler = this.fileDownloadHandler;
+    const send = this.send.bind(this);
+
+    switch (msg.type) {
+      case 'fileDownloadStart':
+        handler.handleDownloadStart(ws, send, tabId, msg.data);
+        break;
+      case 'fileDownloadCancel':
+        handler.handleDownloadCancel(ws, send, tabId, msg.data);
+        break;
     }
   }
 
   private handleFileUploadMessage(ws: WSContext, msg: WebSocketMessage, tabId: string, workingDir: string, permission?: 'view'): void {
     if (!this.fileUploadHandler) {
-      this.fileUploadHandler = new FileUploadHandler(workingDir, this);
+      this.fileUploadHandler = new FileUploadHandler(this);
     }
     const handler = this.fileUploadHandler;
     const send = this.send.bind(this);
 
     switch (msg.type) {
       case 'fileUploadStart':
-        handler.handleUploadStart(ws, send, tabId, msg.data, permission);
+        handler.handleUploadStart(ws, send, tabId, msg.data, workingDir, permission);
         break;
       case 'fileUploadChunk':
         handler.handleUploadChunk(ws, send, tabId, msg.data);
@@ -262,6 +294,10 @@ export class WebSocketImproviseHandler implements HandlerContext {
         this.fileUploadHandler.destroy();
         this.fileUploadHandler = null;
       }
+      if (this.fileDownloadHandler) {
+        this.fileDownloadHandler.destroy();
+        this.fileDownloadHandler = null;
+      }
       if (this.gitHeadWatcher) {
         this.gitHeadWatcher.stop();
         this.gitHeadWatcher = null;
@@ -270,18 +306,45 @@ export class WebSocketImproviseHandler implements HandlerContext {
         this.skillsWatcher.stop();
         this.skillsWatcher = null;
       }
+
+      // Close orphan PTYs when no web client is watching any more.
+      //
+      // Prior behavior: PTYs survived until the user typed `exit` or the
+      // CLI process died. A page refresh / browser tab close / view-switch
+      // would leak the pty, and after a few of these the user would have
+      // tens of zombie shells competing for I/O bandwidth, which manifests
+      // as the app feeling "unresponsive" — interactive operations starve
+      // because the relay socket is saturated streaming output for ptys
+      // that no UI is rendering.
+      //
+      // The active-session preservation in `cleanupConnectionResources`
+      // above is intentionally separate: improvise sessions can produce
+      // useful work while the browser is closed (Claude Code keeps running
+      // in the background); a shell process can't, by construction. So we
+      // preserve the former and reap the latter.
+      const ptyManager = getPTYManager();
+      const activeTerminals = ptyManager.getActiveTerminals();
+      if (activeTerminals.length > 0) {
+        console.log(`[handler] No web subscribers — closing ${activeTerminals.length} orphan PTY${activeTerminals.length === 1 ? '' : 's'}`);
+        ptyManager.closeAll();
+      }
     }
   }
 
   private cleanupConnectionResources(tabMap: Map<string, string>): void {
-    // Destroy sessions owned by this connection
+    // Preserve actively-executing sessions across web reconnects. The runner
+    // is still producing output, and the new web connection will re-attach
+    // via session-handlers.ts::getSession (or initializeTab → reattachSession)
+    // which rebinds listeners and replays executionEventLog. Destroying the
+    // session here would orphan the runner and silently drop all streamed
+    // output for the rest of the prompt.
     const sessionIds = new Set(tabMap.values());
     for (const sessionId of sessionIds) {
       const session = this.sessions.get(sessionId);
-      if (session) {
-        session.destroy();
-        this.sessions.delete(sessionId);
-      }
+      if (!session) continue;
+      if (session.isExecuting) continue;
+      session.destroy();
+      this.sessions.delete(sessionId);
     }
     // Kill search processes owned by this connection's tabs
     for (const tabId of tabMap.keys()) {

package/server/services/websocket/index.ts

@@ -1,5 +1,4 @@
 // Copyright (c) 2025-present Mstro, Inc. All rights reserved.
-// Licensed under the MIT License. See LICENSE file for details.
 
 /**
  * WebSocket Improvise Module