mstro-app 0.5.0 → 0.5.1

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

@@ -1,5 +1,4 @@
 // Copyright (c) 2025-present Mstro, Inc. All rights reserved.
-// Licensed under the MIT License. See LICENSE file for details.
 
 import { resolve } from 'node:path';
 import { loadSkillPrompt } from '../plan/agent-loader.js';
@@ -8,7 +7,7 @@ import { handleGitCommitDiff, handleGitDiff, handleGitShowCommit } from './git-d
 import { handleGitDiscoverRepos, handleGitLog, handleGitSetDirectory } from './git-log-handlers.js';
 import { handleGitPRMessage } from './git-pr-handlers.js';
 import { handleGitCreateTag, handleGitListTags, handleGitPushTag } from './git-tag-handlers.js';
-import { executeGitCommand, parseGitStatus, sendGitError, spawnHaikuWithPrompt, stripCoauthorLines, truncateDiff } from './git-utils.js';
+import { classifyHaikuFailure, executeGitCommand, type HaikuResult, logCommitMessageEvent, parseGitStatus, sendGitError, spawnHaikuWithPrompt, stripCoauthorLines, truncateDiff } from './git-utils.js';
 import { handleGitWorktreeMessage } from './git-worktree-handlers.js';
 import type { HandlerContext } from './handler-context.js';
 import type { GitStatusResponse, WebSocketMessage, WSContext } from './types.js';
@@ -200,7 +199,50 @@ async function handleGitCommit(ctx: HandlerContext, ws: WSContext, msg: WebSocke
   }
 }
 
+/** Tabs currently generating an AI commit message. Prevents double-clicks
+ *  from spawning two `claude` subprocesses for the same tab and clobbering
+ *  the UI when the second response arrives. Module-scoped so it survives
+ *  HandlerContext lifecycle (one process per CLI machine). */
+const inflightCommitMessage = new Set<string>();
+
+const haikuSucceeded = (r: HaikuResult): boolean => r.exitCode === 0 && r.stdout.trim().length > 0;
+// rate_limit is terminal because a 500ms retry hits the same window the API
+// expects you to wait ~60s on; retrying just doubles user-perceived latency
+// before we surface the "wait a minute" message.
+const RETRY_TERMINAL_REASONS = new Set(['binary_missing', 'auth', 'rate_limit']);
+
+/** Run Haiku with one retry on transient failures. Logs each attempt. */
+async function runHaikuForCommitMessage(
+  prompt: string, systemPrompt: string, workingDir: string,
+  meta: { tabId: string; startedAt: number; diffBytes: number },
+): Promise<{ result: HaikuResult; attempts: number }> {
+  let result = await spawnHaikuWithPrompt(prompt, systemPrompt, workingDir);
+  if (haikuSucceeded(result)) return { result, attempts: 1 };
+
+  const first = classifyHaikuFailure(result);
+  if (RETRY_TERMINAL_REASONS.has(first.reason)) return { result, attempts: 1 };
+
+  logCommitMessageEvent(workingDir, {
+    tabId: meta.tabId, attempt: 1, success: false, reason: first.reason,
+    exitCode: result.exitCode, timedOut: result.timedOut,
+    stderrTail: (result.stderr || '').trim().slice(-500),
+    latencyMs: Date.now() - meta.startedAt, diffBytes: meta.diffBytes, willRetry: true,
+  });
+  await new Promise(r => setTimeout(r, 500));
+  result = await spawnHaikuWithPrompt(prompt, systemPrompt, workingDir);
+  return { result, attempts: 2 };
+}
+
 async function handleGitCommitWithAI(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, workingDir: string): Promise<void> {
+  if (inflightCommitMessage.has(tabId)) {
+    // Drop the duplicate request silently — the user already has one in
+    // flight. Surfacing "already generating" as a toast is noisier than
+    // just ignoring the click.
+    return;
+  }
+  inflightCommitMessage.add(tabId);
+  const startedAt = Date.now();
+
   try {
     const statusResult = await executeGitCommand(['status', '--porcelain=v1'], workingDir);
     const { staged } = parseGitStatus(statusResult.stdout);
@@ -216,25 +258,37 @@ async function handleGitCommitWithAI(ctx: HandlerContext, ws: WSContext, msg: We
     const recentCommits = logResult.stdout.trim() || 'No recent commits';
     const stagedFiles = staged.map(f => `${f.status} ${f.path}`).join('\n');
     const diff = truncateDiff(diffResult.stdout);
+    const diffBytes = diffResult.stdout.length;
 
     const prompt = loadSkillPrompt('commit-message', { recentCommits, stagedFiles, diff }, workingDir)
       ?? `You are generating a git commit message for the following staged changes.\n\nRECENT COMMIT MESSAGES (for style reference):\n${recentCommits}\n\nSTAGED FILES:\n${stagedFiles}\n\nDIFF OF STAGED CHANGES:\n${diff}\n\nGenerate a commit message: imperative mood, max 72 characters, focus on "why". Respond with ONLY the commit message.`;
 
-    const result = await spawnHaikuWithPrompt(
-      prompt,
-      'You are a commit message assistant. Respond with only the commit message, no preamble or explanation.',
-      workingDir,
-    );
+    const systemPrompt = 'You are a commit message assistant. Respond with only the commit message, no preamble or explanation.';
+
+    const { result, attempts } = await runHaikuForCommitMessage(prompt, systemPrompt, workingDir, { tabId, startedAt, diffBytes });
 
-    if (result.exitCode !== 0 || !result.stdout.trim()) {
-      console.error('[WebSocketImproviseHandler] Claude commit message error:', result.stderr || 'No output');
-      ctx.send(ws, { type: 'gitError', tabId, data: { error: 'Failed to generate commit message' } });
+    if (!haikuSucceeded(result)) {
+      const failure = classifyHaikuFailure(result);
+      logCommitMessageEvent(workingDir, {
+        tabId, attempt: attempts, success: false, reason: failure.reason,
+        exitCode: result.exitCode, timedOut: result.timedOut,
+        stderrTail: (result.stderr || '').trim().slice(-500),
+        latencyMs: Date.now() - startedAt, diffBytes,
+      });
+      console.error(`[git] commit-message failed (${failure.reason}):`, result.stderr || 'no stderr');
+      ctx.send(ws, { type: 'gitError', tabId, data: { error: failure.userMessage } });
       return;
     }
 
     const commitMessage = extractCommitMessage(result.stdout.trim());
     const autoCommit = !!msg.data?.autoCommit;
 
+    logCommitMessageEvent(workingDir, {
+      tabId, attempt: attempts, success: true,
+      latencyMs: Date.now() - startedAt, diffBytes, autoCommit,
+      messageLength: commitMessage.length,
+    });
+
     ctx.send(ws, { type: 'gitCommitMessage', tabId, data: { message: commitMessage, autoCommit } });
 
     if (autoCommit) {
@@ -250,6 +304,8 @@ async function handleGitCommitWithAI(ctx: HandlerContext, ws: WSContext, msg: We
     }
   } catch (error: unknown) {
     sendGitError(ctx, ws, tabId, error);
+  } finally {
+    inflightCommitMessage.delete(tabId);
   }
 }
 

package/server/services/websocket/git-head-watcher.ts

@@ -1,5 +1,4 @@
 // Copyright (c) 2025-present Mstro, Inc. All rights reserved.
-// Licensed under the MIT License. See LICENSE file for details.
 
 import { existsSync, type FSWatcher, watch } from 'node:fs';
 import { join } from 'node:path';

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

@@ -1,5 +1,4 @@
 // Copyright (c) 2025-present Mstro, Inc. All rights reserved.
-// Licensed under the MIT License. See LICENSE file for details.
 
 import { existsSync, readdirSync } from 'node:fs';
 import { join, resolve } from 'node:path';

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

@@ -1,5 +1,4 @@
 // Copyright (c) 2025-present Mstro, Inc. All rights reserved.
-// Licensed under the MIT License. See LICENSE file for details.
 
 import { loadSkillPrompt } from '../plan/agent-loader.js';
 import { getPrBaseBranch, setPrBaseBranch } from '../settings.js';

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

@@ -1,5 +1,4 @@
 // Copyright (c) 2025-present Mstro, Inc. All rights reserved.
-// Licensed under the MIT License. See LICENSE file for details.
 
 import { executeGitCommand, sendGitError } from './git-utils.js';
 import type { HandlerContext } from './handler-context.js';

package/server/services/websocket/git-utils.ts

@@ -1,8 +1,8 @@
 // Copyright (c) 2025-present Mstro, Inc. All rights reserved.
-// Licensed under the MIT License. See LICENSE file for details.
 
 import { spawn } from 'node:child_process';
-import { existsSync, mkdirSync, unlinkSync, writeFileSync } from 'node:fs';
+import { randomUUID } from 'node:crypto';
+import { appendFileSync, existsSync, mkdirSync, unlinkSync, writeFileSync } from 'node:fs';
 import { join } from 'node:path';
 import type { HandlerContext } from './handler-context.js';
 import type { GitFileStatus, WSContext } from './types.js';
@@ -171,17 +171,31 @@ export function truncateDiff(diff: string, maxLength = 8000): string {
   return `${diff.slice(0, headSize)}\n\n... [diff truncated] ...\n\n${diff.slice(-tailSize)}`;
 }
 
-/** Spawn Claude Haiku with a prompt file and return captured output. */
+/** Result of a Haiku spawn. `timedOut` is set when the process was killed by our timer. */
+export interface HaikuResult {
+  stdout: string;
+  stderr: string;
+  exitCode: number;
+  timedOut: boolean;
+}
+
+/** Spawn Claude Haiku with a prompt file and return captured output.
+ *
+ * Default timeout is 90s — measured p99 of Haiku responses is ~31s for tiny
+ * bouncer prompts; commit-message and PR-description prompts ship a full diff
+ * and run materially longer. 90s gives a 3x margin while still failing fast
+ * enough that the UI doesn't spin forever. */
 export function spawnHaikuWithPrompt(
   prompt: string,
   systemPrompt: string,
   workingDir: string,
-  timeoutMs = 30000,
-): Promise<{ stdout: string; stderr: string; exitCode: number }> {
+  timeoutMs = 90000,
+): Promise<HaikuResult> {
   return new Promise((resolve) => {
     const tempDir = join(workingDir, '.mstro', 'tmp');
     if (!existsSync(tempDir)) mkdirSync(tempDir, { recursive: true });
-    const promptFile = join(tempDir, `haiku-${Date.now()}.txt`);
+    // randomUUID — Date.now() collided when two requests fired in the same ms.
+    const promptFile = join(tempDir, `haiku-${randomUUID()}.txt`);
     writeFileSync(promptFile, prompt);
 
     const args = ['--print', '--model', 'haiku', '--system-prompt', systemPrompt, promptFile];
@@ -189,25 +203,71 @@ export function spawnHaikuWithPrompt(
 
     let stdout = '';
     let stderr = '';
+    let timedOut = false;
     proc.stdout?.on('data', (d: Buffer) => { stdout += d.toString(); });
     proc.stderr?.on('data', (d: Buffer) => { stderr += d.toString(); });
 
-    const timer = setTimeout(() => proc.kill(), timeoutMs);
+    const timer = setTimeout(() => { timedOut = true; proc.kill(); }, timeoutMs);
 
     proc.on('close', (code) => {
       clearTimeout(timer);
       try { unlinkSync(promptFile); } catch { /* ignore cleanup errors */ }
-      resolve({ stdout, stderr, exitCode: code ?? 1 });
+      resolve({ stdout, stderr, exitCode: code ?? 1, timedOut });
     });
 
     proc.on('error', (err: Error) => {
       clearTimeout(timer);
       try { unlinkSync(promptFile); } catch { /* ignore cleanup errors */ }
-      resolve({ stdout: '', stderr: err.message, exitCode: 1 });
+      // ENOENT here means the `claude` binary wasn't found on PATH.
+      const enoent = (err as NodeJS.ErrnoException).code === 'ENOENT';
+      resolve({ stdout: '', stderr: enoent ? `claude: command not found (${err.message})` : err.message, exitCode: 1, timedOut: false });
     });
   });
 }
 
+/** Map a Haiku failure to an actionable user-facing message + a stable reason code. */
+export function classifyHaikuFailure(result: HaikuResult): { reason: string; userMessage: string } {
+  const stderr = (result.stderr || '').toLowerCase();
+  const stdout = (result.stdout || '').toLowerCase();
+  const combined = `${stderr} ${stdout}`;
+
+  if (result.timedOut) {
+    return { reason: 'timeout', userMessage: 'Claude took too long to respond. Try again, or stage fewer files at once.' };
+  }
+  if (combined.includes('command not found') || combined.includes('enoent')) {
+    return { reason: 'binary_missing', userMessage: 'Claude CLI is not installed or not on PATH. Install it with `npm i -g @anthropic-ai/claude-code` and re-launch mstro.' };
+  }
+  if (/not\s+(logged|signed)\s*in|not\s+authenticated|please\s+(log|sign)\s*in|api\s*key|unauthor/.test(combined)) {
+    return { reason: 'auth', userMessage: 'Claude CLI is not authenticated. Run `claude /login` in a terminal, then try again.' };
+  }
+  if (/rate\s*limit|usage\s*limit|429|too\s+many\s+requests/.test(combined)) {
+    return { reason: 'rate_limit', userMessage: 'Claude rate limit reached. Wait a minute and try again.' };
+  }
+  if (/overloaded|529|service\s+unavailable|503/.test(combined)) {
+    return { reason: 'overloaded', userMessage: "Claude is temporarily overloaded. Try again in a few seconds." };
+  }
+  if (result.exitCode === 0 && !result.stdout.trim()) {
+    return { reason: 'empty_output', userMessage: 'Claude returned an empty response. Try again.' };
+  }
+
+  const tail = (result.stderr || result.stdout || '').trim().slice(-200);
+  return { reason: 'unknown', userMessage: `Failed to generate commit message${tail ? `: ${tail}` : ''}` };
+}
+
+/** Append a structured event to .mstro/logs/commit-message.jsonl for observability. */
+export function logCommitMessageEvent(workingDir: string, entry: Record<string, unknown>): void {
+  try {
+    const logDir = join(workingDir, '.mstro', 'logs');
+    if (!existsSync(logDir)) mkdirSync(logDir, { recursive: true });
+    const logFile = join(logDir, 'commit-message.jsonl');
+    const line = `${JSON.stringify({ timestamp: new Date().toISOString(), ...entry })}\n`;
+    appendFileSync(logFile, line, 'utf-8');
+  } catch (err) {
+    // Logging must never break the user-facing flow.
+    console.error('[git] failed to write commit-message log:', err);
+  }
+}
+
 /**
  * Strip injected coauthor/attribution lines from a commit message.
  */

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

@@ -1,5 +1,4 @@
 // 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';
@@ -7,8 +6,21 @@ 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);
@@ -44,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]?.();
 }
@@ -339,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 || {};
@@ -374,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) } });
@@ -482,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;
@@ -494,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);
+
+    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) } });
+  }
+}
 
-    if (removedWorktreePath) {
-      cleanupWorktreeReferences(ctx, workingDir, removedWorktreePath);
+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,5 +1,4 @@
 // 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';

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
@@ -235,7 +234,7 @@ 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);
     }
   }
@@ -259,14 +258,14 @@ export class WebSocketImproviseHandler implements HandlerContext {
 
   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);