mstro-app 0.3.5 → 0.3.6

package/server/services/websocket/file-upload-handler.ts

@@ -0,0 +1,259 @@
+// Copyright (c) 2025-present Mstro, Inc. All rights reserved.
+// Licensed under the MIT License. See LICENSE file for details.
+
+/**
+ * Chunked File Upload Handler
+ *
+ * Receives files in chunks over WebSocket from remote web clients,
+ * writes them to .mstro/tmp/attachments/{tabId}/, and sends progress acks back.
+ */
+
+import type { WriteStream } from 'node:fs';
+import { createWriteStream, existsSync, mkdirSync, rmSync, statSync } from 'node:fs';
+import { join } from 'node:path';
+import type { WebSocketResponse, WSContext } from './types.js';
+
+interface UploadState {
+  uploadId: string;
+  fileName: string;
+  fileSize: number;
+  mimeType: string;
+  isImage: boolean;
+  totalChunks: number;
+  receivedChunks: number;
+  filePath: string;
+  stream: WriteStream;
+  lastActivity: number;
+}
+
+/** Completed upload that's ready to be referenced in an execute message */
+export interface CompletedUpload {
+  uploadId: string;
+  fileName: string;
+  filePath: string;
+  isImage: boolean;
+  mimeType: string;
+  fileSize: number;
+}
+
+const UPLOAD_TIMEOUT_MS = 120_000; // 2 minutes idle timeout
+const MAX_FILE_SIZE = 50 * 1024 * 1024; // 50MB
+
+export class FileUploadHandler {
+  private activeUploads = new Map<string, UploadState>();
+  private completedUploads = new Map<string, CompletedUpload[]>(); // tabId -> completed uploads
+  private cleanupInterval: ReturnType<typeof setInterval>;
+
+  constructor(private workingDir: string) {
+    // Periodically clean up stale uploads
+    this.cleanupInterval = setInterval(() => this.cleanupStaleUploads(), 30_000);
+  }
+
+  /** Get completed uploads for a tab and clear them */
+  getAndClearCompletedUploads(tabId: string): CompletedUpload[] {
+    const uploads = this.completedUploads.get(tabId) || [];
+    this.completedUploads.delete(tabId);
+    return uploads;
+  }
+
+  /** Get completed uploads for a tab without clearing */
+  getCompletedUploads(tabId: string): CompletedUpload[] {
+    return this.completedUploads.get(tabId) || [];
+  }
+
+  handleUploadStart(
+    ws: WSContext,
+    send: (ws: WSContext, response: WebSocketResponse) => void,
+    tabId: string,
+    data: { uploadId: string; fileName: string; fileSize: number; mimeType: string; isImage: boolean; totalChunks: number }
+  ): void {
+    const { uploadId, fileName, fileSize, mimeType, isImage, totalChunks } = data;
+
+    // Validate file size
+    if (fileSize > MAX_FILE_SIZE) {
+      send(ws, {
+        type: 'fileUploadError' as WebSocketResponse['type'],
+        tabId,
+        data: { uploadId, error: `File too large: ${(fileSize / 1024 / 1024).toFixed(1)}MB exceeds ${MAX_FILE_SIZE / 1024 / 1024}MB limit` }
+      });
+      return;
+    }
+
+    // Create attachment directory
+    const attachDir = join(this.workingDir, '.mstro', 'tmp', 'attachments', tabId);
+    if (!existsSync(attachDir)) {
+      mkdirSync(attachDir, { recursive: true });
+    }
+
+    // Handle duplicate file names
+    let targetFileName = fileName;
+    let counter = 1;
+    while (existsSync(join(attachDir, targetFileName))) {
+      const ext = fileName.lastIndexOf('.') !== -1 ? fileName.slice(fileName.lastIndexOf('.')) : '';
+      const base = fileName.lastIndexOf('.') !== -1 ? fileName.slice(0, fileName.lastIndexOf('.')) : fileName;
+      targetFileName = `${base}-${counter}${ext}`;
+      counter++;
+    }
+
+    const filePath = join(attachDir, targetFileName);
+    const stream = createWriteStream(filePath);
+
+    const uploadState: UploadState = {
+      uploadId,
+      fileName: targetFileName,
+      fileSize,
+      mimeType,
+      isImage,
+      totalChunks,
+      receivedChunks: 0,
+      filePath,
+      stream,
+      lastActivity: Date.now(),
+    };
+
+    this.activeUploads.set(uploadId, uploadState);
+
+    // Send ack for start
+    send(ws, {
+      type: 'fileUploadAck' as WebSocketResponse['type'],
+      tabId,
+      data: { uploadId, chunkIndex: -1, status: 'ok' }
+    });
+  }
+
+  handleUploadChunk(
+    ws: WSContext,
+    send: (ws: WSContext, response: WebSocketResponse) => void,
+    tabId: string,
+    data: { uploadId: string; chunkIndex: number; content: string }
+  ): void {
+    const { uploadId, chunkIndex, content } = data;
+    const upload = this.activeUploads.get(uploadId);
+
+    if (!upload) {
+      send(ws, {
+        type: 'fileUploadError' as WebSocketResponse['type'],
+        tabId,
+        data: { uploadId, error: 'Upload not found or expired' }
+      });
+      return;
+    }
+
+    try {
+      const buffer = Buffer.from(content, 'base64');
+      upload.stream.write(buffer);
+      upload.receivedChunks++;
+      upload.lastActivity = Date.now();
+
+      send(ws, {
+        type: 'fileUploadAck' as WebSocketResponse['type'],
+        tabId,
+        data: { uploadId, chunkIndex, status: 'ok' }
+      });
+    } catch (err) {
+      const errorMsg = err instanceof Error ? err.message : String(err);
+      send(ws, {
+        type: 'fileUploadError' as WebSocketResponse['type'],
+        tabId,
+        data: { uploadId, error: `Chunk write failed: ${errorMsg}` }
+      });
+      this.cancelUpload(uploadId);
+    }
+  }
+
+  handleUploadComplete(
+    ws: WSContext,
+    send: (ws: WSContext, response: WebSocketResponse) => void,
+    tabId: string,
+    data: { uploadId: string }
+  ): void {
+    const { uploadId } = data;
+    const upload = this.activeUploads.get(uploadId);
+
+    if (!upload) {
+      send(ws, {
+        type: 'fileUploadError' as WebSocketResponse['type'],
+        tabId,
+        data: { uploadId, error: 'Upload not found or expired' }
+      });
+      return;
+    }
+
+    upload.stream.end(() => {
+      // Verify file was written
+      try {
+        const stat = statSync(upload.filePath);
+        const completed: CompletedUpload = {
+          uploadId,
+          fileName: upload.fileName,
+          filePath: upload.filePath,
+          isImage: upload.isImage,
+          mimeType: upload.mimeType,
+          fileSize: stat.size,
+        };
+
+        // Store completed upload for this tab
+        const tabUploads = this.completedUploads.get(tabId) || [];
+        tabUploads.push(completed);
+        this.completedUploads.set(tabId, tabUploads);
+
+        this.activeUploads.delete(uploadId);
+
+        send(ws, {
+          type: 'fileUploadReady' as WebSocketResponse['type'],
+          tabId,
+          data: { uploadId, filePath: upload.filePath, fileName: upload.fileName }
+        });
+      } catch (err) {
+        const errorMsg = err instanceof Error ? err.message : String(err);
+        send(ws, {
+          type: 'fileUploadError' as WebSocketResponse['type'],
+          tabId,
+          data: { uploadId, error: `File verification failed: ${errorMsg}` }
+        });
+        this.activeUploads.delete(uploadId);
+      }
+    });
+  }
+
+  handleUploadCancel(
+    _ws: WSContext,
+    _send: (ws: WSContext, response: WebSocketResponse) => void,
+    _tabId: string,
+    data: { uploadId: string }
+  ): void {
+    this.cancelUpload(data.uploadId);
+  }
+
+  private cancelUpload(uploadId: string): void {
+    const upload = this.activeUploads.get(uploadId);
+    if (!upload) return;
+
+    try {
+      upload.stream.destroy();
+      if (existsSync(upload.filePath)) {
+        rmSync(upload.filePath, { force: true });
+      }
+    } catch {
+      // Ignore cleanup errors
+    }
+    this.activeUploads.delete(uploadId);
+  }
+
+  private cleanupStaleUploads(): void {
+    const now = Date.now();
+    for (const [uploadId, upload] of this.activeUploads) {
+      if (now - upload.lastActivity > UPLOAD_TIMEOUT_MS) {
+        console.warn(`[FileUploadHandler] Upload ${uploadId} timed out, cleaning up`);
+        this.cancelUpload(uploadId);
+      }
+    }
+  }
+
+  destroy(): void {
+    clearInterval(this.cleanupInterval);
+    for (const uploadId of this.activeUploads.keys()) {
+      this.cancelUpload(uploadId);
+    }
+  }
+}

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

@@ -188,13 +188,13 @@ const GIT_PR_TYPES = new Set([
 
 // Worktree/merge message types that route to git-worktree-handlers
 const GIT_WORKTREE_TYPES = new Set([
-  'gitWorktreeList', 'gitWorktreeCreate', 'gitWorktreeRemove',
+  'gitWorktreeList', 'gitWorktreeCreate', 'gitWorktreeCreateAndAssign', 'gitWorktreeRemove',
   'tabWorktreeSwitch', 'gitWorktreePush', 'gitWorktreeCreatePR',
   'gitMergePreview', 'gitWorktreeMerge', 'gitMergeAbort', 'gitMergeComplete',
 ]);
 
 /** Route git messages to appropriate sub-handler */
-export function handleGitMessage(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, workingDir: string): void {
+export async function handleGitMessage(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, workingDir: string): Promise<void> {
   const gitDir = ctx.gitDirectories.get(tabId) || workingDir;
 
   if (GIT_PR_TYPES.has(msg.type)) {
@@ -202,7 +202,7 @@ export function handleGitMessage(ctx: HandlerContext, ws: WSContext, msg: WebSoc
     return;
   }
   if (GIT_WORKTREE_TYPES.has(msg.type)) {
-    handleGitWorktreeMessage(ctx, ws, msg, tabId, gitDir, workingDir);
+    await handleGitWorktreeMessage(ctx, ws, msg, tabId, gitDir, workingDir);
     return;
   }
 

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

@@ -6,10 +6,11 @@ import { executeGitCommand, handleGitStatus, spawnWithOutput } from './git-handl
 import type { HandlerContext } from './handler-context.js';
 import type { WebSocketMessage, WorktreeInfo, WSContext } from './types.js';
 
-export function handleGitWorktreeMessage(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, gitDir: string, workingDir: string): void {
-  const handlers: Record<string, () => void> = {
+export async function handleGitWorktreeMessage(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, gitDir: string, workingDir: string): Promise<void> {
+  const handlers: Record<string, () => Promise<void>> = {
     gitWorktreeList: () => handleGitWorktreeList(ctx, ws, tabId, gitDir),
     gitWorktreeCreate: () => handleGitWorktreeCreate(ctx, ws, msg, tabId, gitDir),
+    gitWorktreeCreateAndAssign: () => handleGitWorktreeCreateAndAssign(ctx, ws, msg, tabId, gitDir, workingDir),
     gitWorktreeRemove: () => handleGitWorktreeRemove(ctx, ws, msg, tabId, gitDir),
     tabWorktreeSwitch: () => handleTabWorktreeSwitch(ctx, ws, msg, tabId, workingDir),
     gitWorktreePush: () => handleGitWorktreePush(ctx, ws, msg, tabId, gitDir),
@@ -19,7 +20,7 @@ export function handleGitWorktreeMessage(ctx: HandlerContext, ws: WSContext, msg
     gitMergeAbort: () => handleGitMergeAbort(ctx, ws, tabId, gitDir),
     gitMergeComplete: () => handleGitMergeComplete(ctx, ws, msg, tabId, gitDir),
   };
-  handlers[msg.type]?.();
+  await handlers[msg.type]?.();
 }
 
 function applyWorktreePorcelainLine(line: string, worktrees: WorktreeInfo[], current: Partial<WorktreeInfo>): Partial<WorktreeInfo> {
@@ -90,6 +91,44 @@ async function handleGitWorktreeCreate(ctx: HandlerContext, ws: WSContext, msg:
   }
 }
 
+async function handleGitWorktreeCreateAndAssign(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, gitDir: string, workingDir: string): Promise<void> {
+  try {
+    const { branchName, baseBranch, path: worktreePath } = msg.data || {};
+    if (!branchName) {
+      ctx.send(ws, { type: 'gitError', tabId, data: { error: 'Branch name is required' } });
+      return;
+    }
+
+    const repoName = gitDir.split('/').pop() || 'repo';
+    const wtPath = worktreePath || join(dirname(gitDir), `${repoName}-worktrees`, branchName);
+
+    const args = ['worktree', 'add', wtPath, '-b', branchName, ...(baseBranch ? [baseBranch] : [])];
+    const result = await executeGitCommand(args, gitDir);
+    if (result.exitCode !== 0) {
+      ctx.send(ws, { type: 'gitError', tabId, data: { error: result.stderr || 'Failed to create worktree' } });
+      return;
+    }
+
+    const headResult = await executeGitCommand(['rev-parse', '--short', 'HEAD'], wtPath);
+
+    // Assign to tab
+    ctx.gitDirectories.set(tabId, wtPath);
+    ctx.gitBranches.set(tabId, branchName);
+    const registry = ctx.getRegistry(workingDir);
+    registry.updateTabWorktree(tabId, wtPath, branchName);
+
+    ctx.send(ws, {
+      type: 'gitWorktreeCreatedAndAssigned',
+      tabId,
+      data: { tabId, path: wtPath, branch: branchName, head: headResult.stdout.trim() },
+    });
+
+    handleGitStatus(ctx, ws, tabId, wtPath);
+  } catch (error: unknown) {
+    ctx.send(ws, { type: 'gitError', tabId, data: { error: error instanceof Error ? error.message : String(error) } });
+  }
+}
+
 async function handleGitWorktreeRemove(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, workingDir: string): Promise<void> {
   try {
     const { path: wtPath, force, deleteBranch } = msg.data || {};
@@ -127,8 +166,11 @@ async function handleTabWorktreeSwitch(ctx: HandlerContext, ws: WSContext, msg:
   try {
     const { tabId: targetTabId, worktreePath } = msg.data || {};
     const resolvedTabId = targetTabId || tabId;
+    const registry = ctx.getRegistry(workingDir);
     if (!worktreePath) {
       ctx.gitDirectories.delete(resolvedTabId);
+      ctx.gitBranches.delete(resolvedTabId);
+      registry.updateTabWorktree(resolvedTabId, null, null);
       ctx.send(ws, { type: 'tabWorktreeSwitched', tabId: resolvedTabId, data: { tabId: resolvedTabId, worktreePath: workingDir, branch: '' } });
       handleGitStatus(ctx, ws, resolvedTabId, workingDir);
       return;
@@ -138,6 +180,8 @@ async function handleTabWorktreeSwitch(ctx: HandlerContext, ws: WSContext, msg:
 
     const branchResult = await executeGitCommand(['rev-parse', '--abbrev-ref', 'HEAD'], worktreePath);
     const branch = branchResult.stdout.trim();
+    ctx.gitBranches.set(resolvedTabId, branch);
+    registry.updateTabWorktree(resolvedTabId, worktreePath, branch);
 
     ctx.send(ws, { type: 'tabWorktreeSwitched', tabId: resolvedTabId, data: { tabId: resolvedTabId, worktreePath, branch } });
     handleGitStatus(ctx, ws, resolvedTabId, worktreePath);

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

@@ -4,6 +4,7 @@
 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 { SessionRegistry } from './session-registry.js';
 import type { WebSocketResponse, WSContext } from './types.js';
 
@@ -25,11 +26,13 @@ export interface HandlerContext {
   connections: Map<WSContext, Map<string, string>>;
   allConnections: Set<WSContext>;
   gitDirectories: Map<string, string>;
+  gitBranches: Map<string, string>;
   activeSearches: Map<string, ChildProcess>;
   terminalSubscribers: Map<string, Set<WSContext>>;
   terminalListenerCleanups: Map<string, () => void>;
   autocompleteService: AutocompleteService;
   usageReporter: UsageReporter | null;
+  fileUploadHandler: FileUploadHandler | null;
 
   // Registry access
   getRegistry(workingDir: string): SessionRegistry;

package/server/services/websocket/handler.ts

@@ -16,6 +16,7 @@ import type { ImprovisationSessionManager } from '../../cli/improvisation-sessio
 import { captureException } from '../sentry.js';
 import { AutocompleteService } from './autocomplete.js';
 import { handleFileExplorerMessage, handleFileMessage } from './file-explorer-handlers.js';
+import { FileUploadHandler } from './file-upload-handler.js';
 import { handleGitMessage } from './git-handlers.js';
 import type { HandlerContext, UsageReporter } from './handler-context.js';
 import { handleHistoryMessage, handleSessionMessage, initializeTab, resumeHistoricalSession } from './session-handlers.js';
@@ -34,11 +35,13 @@ export class WebSocketImproviseHandler implements HandlerContext {
   private frecencyPath: string;
   usageReporter: UsageReporter | null = null;
   gitDirectories: Map<string, string> = new Map();
+  gitBranches: Map<string, string> = new Map();
   private sessionRegistry: SessionRegistry | null = null;
   allConnections: Set<WSContext> = new Set();
   activeSearches: Map<string, ChildProcess> = new Map();
   terminalListenerCleanups: Map<string, () => void> = new Map();
   terminalSubscribers: Map<string, Set<WSContext>> = new Map();
+  fileUploadHandler: FileUploadHandler | null = null;
 
   constructor() {
     this.frecencyPath = join(homedir(), '.mstro', 'autocomplete-frecency.json');
@@ -192,6 +195,7 @@ export class WebSocketImproviseHandler implements HandlerContext {
       case 'gitPushTag':
       case 'gitWorktreeList':
       case 'gitWorktreeCreate':
+      case 'gitWorktreeCreateAndAssign':
       case 'gitWorktreeRemove':
       case 'tabWorktreeSwitch':
       case 'gitWorktreePush':
@@ -221,11 +225,40 @@ export class WebSocketImproviseHandler implements HandlerContext {
         return handleGetSettings(this, ws);
       case 'updateSettings':
         return handleUpdateSettings(this, ws, msg);
+      // File upload messages (chunked remote uploads)
+      case 'fileUploadStart':
+      case 'fileUploadChunk':
+      case 'fileUploadComplete':
+      case 'fileUploadCancel':
+        return this.handleFileUploadMessage(ws, msg, tabId, workingDir);
       default:
         throw new Error(`Unknown message type: ${msg.type}`);
     }
   }
 
+  private handleFileUploadMessage(ws: WSContext, msg: WebSocketMessage, tabId: string, workingDir: string): void {
+    if (!this.fileUploadHandler) {
+      this.fileUploadHandler = new FileUploadHandler(workingDir);
+    }
+    const handler = this.fileUploadHandler;
+    const send = this.send.bind(this);
+
+    switch (msg.type) {
+      case 'fileUploadStart':
+        handler.handleUploadStart(ws, send, tabId, msg.data);
+        break;
+      case 'fileUploadChunk':
+        handler.handleUploadChunk(ws, send, tabId, msg.data);
+        break;
+      case 'fileUploadComplete':
+        handler.handleUploadComplete(ws, send, tabId, msg.data);
+        break;
+      case 'fileUploadCancel':
+        handler.handleUploadCancel(ws, send, tabId, msg.data);
+        break;
+    }
+  }
+
   handleClose(ws: WSContext): void {
     this.connections.delete(ws);
     this.allConnections.delete(ws);

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

@@ -3,7 +3,7 @@
 
 import { existsSync, readdirSync, readFileSync, unlinkSync } from 'node:fs';
 import { join } from 'node:path';
-import { ImprovisationSessionManager } from '../../cli/improvisation-session-manager.js';
+import { type FileAttachment, ImprovisationSessionManager } from '../../cli/improvisation-session-manager.js';
 import { getModel } from '../settings.js';
 import type { HandlerContext } from './handler-context.js';
 import type { SessionRegistry } from './session-registry.js';
@@ -128,6 +128,29 @@ export function setupSessionListeners(ctx: HandlerContext, session: Improvisatio
   });
 }
 
+/** Merge pre-uploaded files (from chunked upload) with any inline attachments */
+function mergePreUploadedAttachments(ctx: HandlerContext, tabId: string, inlineAttachments?: FileAttachment[]): FileAttachment[] | undefined {
+  if (!ctx.fileUploadHandler) return inlineAttachments;
+  const preUploaded = ctx.fileUploadHandler.getAndClearCompletedUploads(tabId);
+  if (preUploaded.length === 0) return inlineAttachments;
+
+  const merged: (FileAttachment & { _preUploaded?: boolean })[] = [...(inlineAttachments || [])];
+  for (const upload of preUploaded) {
+    const alreadyIncluded = merged.some(a => a.fileName === upload.fileName);
+    if (!alreadyIncluded) {
+      merged.push({
+        fileName: upload.fileName,
+        filePath: upload.filePath,
+        content: '',
+        isImage: upload.isImage,
+        mimeType: upload.mimeType,
+        _preUploaded: true,
+      });
+    }
+  }
+  return merged;
+}
+
 export function handleSessionMessage(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, permission?: 'control' | 'view'): void {
   switch (msg.type) {
     case 'execute': {
@@ -135,7 +158,8 @@ export function handleSessionMessage(ctx: HandlerContext, ws: WSContext, msg: We
       const session = requireSession(ctx, ws, tabId);
       const sandboxed = permission === 'control' || permission === 'view';
       const worktreeDir = ctx.gitDirectories.get(tabId);
-      session.executePrompt(msg.data.prompt, msg.data.attachments, { sandboxed, workingDir: worktreeDir });
+      const attachments = mergePreUploadedAttachments(ctx, tabId, msg.data.attachments);
+      session.executePrompt(msg.data.prompt, attachments, { sandboxed, workingDir: worktreeDir });
       break;
     }
     case 'cancel': {
@@ -206,6 +230,47 @@ export function handleHistoryMessage(ctx: HandlerContext, ws: WSContext, msg: We
   }
 }
 
+function tryResumeFromDisk(
+  ctx: HandlerContext,
+  ws: WSContext,
+  tabId: string,
+  workingDir: string,
+  registrySessionId: string,
+  tabMap: Map<string, string> | undefined,
+  registry: SessionRegistry
+): boolean {
+  try {
+    const diskSession = ImprovisationSessionManager.resumeFromHistory(workingDir, registrySessionId);
+    setupSessionListeners(ctx, diskSession, ws, tabId);
+    const diskSessionId = diskSession.getSessionInfo().sessionId;
+    ctx.sessions.set(diskSessionId, diskSession);
+    if (tabMap) tabMap.set(tabId, diskSessionId);
+    registry.touchTab(tabId);
+
+    // Restore worktree state from registry
+    const regTab = registry.getTab(tabId);
+    if (regTab?.worktreePath && !ctx.gitDirectories.has(tabId)) {
+      ctx.gitDirectories.set(tabId, regTab.worktreePath);
+      if (regTab.worktreeBranch) ctx.gitBranches.set(tabId, regTab.worktreeBranch);
+    }
+    const worktreePath = ctx.gitDirectories.get(tabId);
+    const worktreeBranch = ctx.gitBranches.get(tabId);
+
+    ctx.send(ws, {
+      type: 'tabInitialized',
+      tabId,
+      data: {
+        ...diskSession.getSessionInfo(),
+        outputHistory: buildOutputHistory(diskSession),
+        ...(worktreePath ? { worktreePath, worktreeBranch } : {}),
+      }
+    });
+    return true;
+  } catch {
+    return false;
+  }
+}
+
 export async function initializeTab(ctx: HandlerContext, ws: WSContext, tabId: string, workingDir: string, tabName?: string): Promise<void> {
   const tabMap = ctx.connections.get(ws);
   const registry = ctx.getRegistry(workingDir);
@@ -229,25 +294,8 @@ export async function initializeTab(ctx: HandlerContext, ws: WSContext, tabId: s
       return;
     }
 
-    try {
-      const diskSession = ImprovisationSessionManager.resumeFromHistory(workingDir, registrySessionId);
-      setupSessionListeners(ctx, diskSession, ws, tabId);
-      const diskSessionId = diskSession.getSessionInfo().sessionId;
-      ctx.sessions.set(diskSessionId, diskSession);
-      if (tabMap) tabMap.set(tabId, diskSessionId);
-      registry.touchTab(tabId);
-
-      ctx.send(ws, {
-        type: 'tabInitialized',
-        tabId,
-        data: {
-          ...diskSession.getSessionInfo(),
-          outputHistory: buildOutputHistory(diskSession),
-        }
-      });
+    if (tryResumeFromDisk(ctx, ws, tabId, workingDir, registrySessionId, tabMap, registry)) {
       return;
-    } catch {
-      // Disk session not found — fall through to create new
     }
   }
 
@@ -352,12 +400,22 @@ function reattachSession(
   if (tabMap) tabMap.set(tabId, sessionId);
   registry.touchTab(tabId);
 
+  // Restore worktree state from registry if not already in memory
+  const regTab = registry.getTab(tabId);
+  if (regTab?.worktreePath && !ctx.gitDirectories.has(tabId)) {
+    ctx.gitDirectories.set(tabId, regTab.worktreePath);
+    if (regTab.worktreeBranch) ctx.gitBranches.set(tabId, regTab.worktreeBranch);
+  }
+
   const outputHistory = buildOutputHistory(session);
 
   const executionEvents = session.isExecuting
     ? session.getExecutionEventLog()
     : undefined;
 
+  const worktreePath = ctx.gitDirectories.get(tabId);
+  const worktreeBranch = ctx.gitBranches.get(tabId);
+
   ctx.send(ws, {
     type: 'tabInitialized',
     tabId,
@@ -367,6 +425,7 @@ function reattachSession(
       isExecuting: session.isExecuting,
       executionEvents,
       ...(session.isExecuting && session.executionStartTimestamp ? { executionStartTimestamp: session.executionStartTimestamp } : {}),
+      ...(worktreePath ? { worktreePath, worktreeBranch } : {}),
     }
   });
 }

package/server/services/websocket/session-registry.ts

@@ -21,6 +21,8 @@ export interface RegisteredTab {
   lastActivityAt: string
   order: number
   hasUnviewedCompletion: boolean
+  worktreePath?: string
+  worktreeBranch?: string
 }
 
 interface RegistryData {
@@ -167,6 +169,22 @@ export class SessionRegistry {
     }
   }
 
+  /**
+   * Update worktree assignment for a tab. Pass null to clear.
+   */
+  updateTabWorktree(tabId: string, worktreePath: string | null, worktreeBranch: string | null): void {
+    if (this.data.tabs[tabId]) {
+      if (worktreePath) {
+        this.data.tabs[tabId].worktreePath = worktreePath
+        this.data.tabs[tabId].worktreeBranch = worktreeBranch || undefined
+      } else {
+        delete this.data.tabs[tabId].worktreePath
+        delete this.data.tabs[tabId].worktreeBranch
+      }
+      this.save()
+    }
+  }
+
   /**
    * Reorder tabs. Accepts an ordered array of tabIds and reassigns order values.
    */