mstro-app 0.3.5 → 0.3.7

package/server/cli/improvisation-session-manager.ts

@@ -132,6 +132,12 @@ export class ImprovisationSessionManager extends EventEmitter {
   private executionEventLog: Array<{ type: string; data: unknown; timestamp: number }> = [];
   /** Set by cancel() to signal the retry loop to exit */
   private _cancelled: boolean = false;
+  /** True when cancel() has already emitted movementComplete (prevents double-emit) */
+  private _cancelCompleteEmitted: boolean = false;
+  /** Current execution's user prompt (for cancel to build movement record) */
+  private _currentUserPrompt: string = '';
+  /** Current execution's sequence number (for cancel to build movement record) */
+  private _currentSequenceNumber: number = 0;
 
   /**
    * Resume from a historical session.
@@ -283,6 +289,13 @@ export class ImprovisationSessionManager extends EventEmitter {
 
     const paths: string[] = [];
     for (const attachment of attachments) {
+      // Pre-uploaded files are already on disk from chunked upload
+      if ((attachment as FileAttachment & { _preUploaded?: boolean })._preUploaded) {
+        if (existsSync(attachment.filePath)) {
+          paths.push(attachment.filePath);
+        }
+        continue;
+      }
       const filePath = join(attachDir, attachment.fileName);
       try {
         // All paste content arrives as base64 — decode to binary
@@ -321,10 +334,13 @@ export class ImprovisationSessionManager extends EventEmitter {
     const _execStart = Date.now();
     this._isExecuting = true;
     this._cancelled = false;
+    this._cancelCompleteEmitted = false;
     this._executionStartTimestamp = _execStart;
     this.executionEventLog = [];
 
     const sequenceNumber = this.history.movements.length + 1;
+    this._currentUserPrompt = userPrompt;
+    this._currentSequenceNumber = sequenceNumber;
     this.emit('onMovementStart', sequenceNumber, userPrompt);
     trackEvent(AnalyticsEvents.IMPROVISE_PROMPT_RECEIVED, {
       prompt_length: userPrompt.length,
@@ -417,6 +433,13 @@ export class ImprovisationSessionManager extends EventEmitter {
     this.executionEventLog = [];
     this.currentRunner = null;
 
+    // If cancel() already emitted movementComplete, just clean up state —
+    // don't double-emit or double-persist.
+    if (this._cancelCompleteEmitted) {
+      const existing = this.history.movements.find(m => m.sequenceNumber === sequenceNumber);
+      if (existing) return existing;
+    }
+
     const cancelledMovement: MovementRecord = {
       id: `prompt-${sequenceNumber}`,
       sequenceNumber,
@@ -485,16 +508,50 @@ export class ImprovisationSessionManager extends EventEmitter {
     return result;
   }
 
+  /** MIME types that the Claude API can accept as image content blocks */
+  private static readonly SUPPORTED_IMAGE_MIMES = new Set([
+    'image/jpeg', 'image/png', 'image/gif', 'image/webp',
+  ]);
+
+  /** Hydrate pre-uploaded images from disk and downgrade unsupported formats */
+  private hydrateAndFilterAttachments(attachments: FileAttachment[]): void {
+    for (const attachment of attachments) {
+      // Pre-uploaded images need their content read from disk
+      const preUploaded = (attachment as FileAttachment & { _preUploaded?: boolean })._preUploaded;
+      if (preUploaded && attachment.isImage && !attachment.content && existsSync(attachment.filePath)) {
+        try {
+          attachment.content = readFileSync(attachment.filePath).toString('base64');
+        } catch (err) {
+          console.error(`Failed to read pre-uploaded image ${attachment.filePath}:`, err);
+          attachment.isImage = false;
+        }
+      }
+
+      // Downgrade unsupported image formats (SVG, BMP, TIFF, ICO, etc.) to text attachments
+      if (attachment.isImage) {
+        const mime = (attachment.mimeType || '').toLowerCase();
+        if (mime && !ImprovisationSessionManager.SUPPORTED_IMAGE_MIMES.has(mime)) {
+          attachment.isImage = false;
+        }
+      }
+    }
+  }
+
   /** Prepare prompt with attachments and limit image count */
   private preparePromptAndAttachments(
     userPrompt: string,
     attachments: FileAttachment[] | undefined,
   ): { prompt: string; imageAttachments: FileAttachment[] | undefined } {
+    if (attachments) {
+      this.hydrateAndFilterAttachments(attachments);
+    }
+
     const diskPaths = attachments ? this.persistAttachments(attachments) : [];
     const prompt = this.buildPromptWithAttachments(userPrompt, attachments, diskPaths);
 
     const MAX_IMAGE_ATTACHMENTS = 20;
-    const allImages = attachments?.filter(a => a.isImage);
+    // Only include images that have valid content
+    const allImages = attachments?.filter(a => a.isImage && a.content);
     let imageAttachments = allImages;
     if (allImages && allImages.length > MAX_IMAGE_ATTACHMENTS) {
       imageAttachments = allImages.slice(-MAX_IMAGE_ATTACHMENTS);
@@ -1469,16 +1526,47 @@ export class ImprovisationSessionManager extends EventEmitter {
   }
 
   /**
-   * Cancel current execution
+   * Cancel current execution — immediately emits movementComplete so the web
+   * gets instant feedback, then cleans up the process tree in the background.
    */
   cancel(): void {
     this._cancelled = true;
+
     if (this.currentRunner) {
       this.currentRunner.cleanup();
       this.currentRunner = null;
-      this.queueOutput('\n⚠ Execution cancelled\n');
-      this.flushOutputQueue();
     }
+
+    // Emit movementComplete immediately so the web UI updates without waiting
+    // for the process tree to fully die (SIGTERM → SIGKILL can take up to 5s).
+    if (this._isExecuting && !this._cancelCompleteEmitted) {
+      this._cancelCompleteEmitted = true;
+      const execStart = this._executionStartTimestamp || Date.now();
+      this._isExecuting = false;
+      this._executionStartTimestamp = undefined;
+
+      const cancelledMovement: MovementRecord = {
+        id: `prompt-${this._currentSequenceNumber}`,
+        sequenceNumber: this._currentSequenceNumber,
+        userPrompt: this._currentUserPrompt,
+        timestamp: new Date().toISOString(),
+        tokensUsed: 0,
+        summary: '',
+        filesModified: [],
+        errorOutput: 'Execution cancelled by user',
+        durationMs: Date.now() - execStart,
+      };
+      this.persistMovement(cancelledMovement);
+
+      const fallbackResult = {
+        completed: false, needsHandoff: false, totalTokens: 0, sessionId: '',
+        output: '', exitCode: 1, signalName: 'SIGTERM',
+      } as HeadlessRunResult;
+      this.emitMovementComplete(cancelledMovement, fallbackResult, execStart, this._currentSequenceNumber);
+    }
+
+    this.queueOutput('\n⚠ Execution cancelled\n');
+    this.flushOutputQueue();
   }
 
   /**

package/server/services/platform.ts

@@ -117,6 +117,9 @@ interface ConnectionCallbacks {
 /**
  * Platform WebSocket connection with token-based authentication
  */
+/** Number of missed pongs before treating connection as dead */
+const MAX_MISSED_PONGS = 2
+
 export class PlatformConnection {
   private ws: WebSocket | null = null
   private reconnectTimeout: ReturnType<typeof setTimeout> | null = null
@@ -130,6 +133,7 @@ export class PlatformConnection {
   private isConnected = false
   private tokenRefreshInterval: ReturnType<typeof setInterval> | null = null
   private heartbeatInterval: ReturnType<typeof setInterval> | null = null
+  private missedPongs = 0
 
   constructor(
     workingDirectory: string,
@@ -184,19 +188,33 @@ export class PlatformConnection {
   }
 
   /**
-   * Start heartbeat to keep connection alive and refresh server-side TTL
+   * Start heartbeat to keep connection alive and refresh server-side TTL.
+   * Tracks missed pongs — if the server doesn't respond to MAX_MISSED_PONGS
+   * consecutive pings, the connection is considered dead and force-closed
+   * to trigger reconnection.
    */
   private startHeartbeat(): void {
+    this.missedPongs = 0
     // Send ping every 2 minutes (server TTL is 5 minutes)
-    this.heartbeatInterval = setInterval(() => {
-      if (this.ws && this.isConnected) {
-        try {
-          this.ws.send(JSON.stringify({ type: 'ping' }))
-        } catch {
-          // Ignore send errors - will reconnect if disconnected
-        }
-      }
-    }, 2 * 60 * 1000)
+    this.heartbeatInterval = setInterval(() => this.heartbeatTick(), 2 * 60 * 1000)
+  }
+
+  private heartbeatTick(): void {
+    if (!this.ws || !this.isConnected) return
+
+    if (this.missedPongs >= MAX_MISSED_PONGS) {
+      console.log(`[Platform] ${this.missedPongs} pongs missed — forcing reconnect`)
+      this.missedPongs = 0
+      this.stopHeartbeat()
+      try { this.ws.close() } catch { /* ignore */ }
+      return
+    }
+    this.missedPongs++
+    try {
+      this.ws.send(JSON.stringify({ type: 'ping' }))
+    } catch {
+      // Send failed — onclose will handle reconnect
+    }
   }
 
   /**
@@ -363,7 +381,7 @@ export class PlatformConnection {
         break
 
       case 'pong':
-        // Heartbeat response, ignore
+        this.missedPongs = 0
         break
 
       default:

package/server/services/terminal/pty-manager.ts

@@ -102,6 +102,47 @@ export function getPtyInstallInstructions(): string {
 // Import type separately for type-checking (doesn't require the module to load)
 type IPty = import('node-pty').IPty;
 
+/**
+ * Fixed-size buffer that retains the most recent PTY output for replay on reconnect.
+ * Stores raw string chunks and evicts oldest data when the total exceeds maxBytes.
+ */
+class ScrollbackBuffer {
+  private chunks: string[] = [];
+  private totalLength = 0;
+  private maxBytes: number;
+
+  constructor(maxBytes: number) {
+    this.maxBytes = maxBytes;
+  }
+
+  append(data: string): void {
+    this.chunks.push(data);
+    this.totalLength += data.length;
+    // Evict oldest chunks until under budget
+    while (this.totalLength > this.maxBytes && this.chunks.length > 1) {
+      const removed = this.chunks.shift()!;
+      this.totalLength -= removed.length;
+    }
+    // If a single chunk exceeds max, truncate from the front
+    if (this.totalLength > this.maxBytes && this.chunks.length === 1) {
+      const excess = this.totalLength - this.maxBytes;
+      this.chunks[0] = this.chunks[0].slice(excess);
+      this.totalLength = this.chunks[0].length;
+    }
+  }
+
+  getContents(): string {
+    return this.chunks.join('');
+  }
+
+  clear(): void {
+    this.chunks = [];
+    this.totalLength = 0;
+  }
+}
+
+const SCROLLBACK_MAX_BYTES = 256 * 1024; // 256KB
+
 export interface PTYSession {
   id: string;
   pty: IPty;
@@ -117,6 +158,8 @@ export interface PTYSession {
   // Output coalescing: buffer small chunks into fewer WS messages
   _outputBuffer: string;
   _outputTimer: ReturnType<typeof setTimeout> | null;
+  // Scrollback ring buffer for replay on reconnect
+  scrollback: ScrollbackBuffer;
 }
 
 /**
@@ -201,7 +244,7 @@ export class PTYManager extends EventEmitter {
     rows: number = 24,
     requestedShell?: string,
     options?: { sandboxed?: boolean }
-  ): { shell: string; cwd: string; isReconnect: boolean } {
+  ): { shell: string; cwd: string; isReconnect: boolean; platform: string } {
     // Check if node-pty is available
     if (!pty) {
       throw new Error(`PTY_NOT_AVAILABLE:${getPtyInstallInstructions()}`);
@@ -221,6 +264,7 @@ export class PTYManager extends EventEmitter {
         shell: existingSession.shell,
         cwd: existingSession.cwd,
         isReconnect: true,
+        platform: platform(),
       };
     }
 
@@ -259,6 +303,7 @@ export class PTYManager extends EventEmitter {
         rows,
         _outputBuffer: '',
         _outputTimer: null,
+        scrollback: new ScrollbackBuffer(SCROLLBACK_MAX_BYTES),
       };
       this.terminals.set(terminalId, session);
 
@@ -267,7 +312,9 @@ export class PTYManager extends EventEmitter {
       // wraps echoed chars in multi-part ANSI sequences (RPROMPT, syntax highlighting).
       // A longer window on macOS ensures these multi-part sequences arrive as one chunk,
       // which the browser's predictive echo can match correctly.
-      const OUTPUT_COALESCE_MS = platform() === 'darwin' ? 24 : 8;
+      // 32ms on macOS captures full zsh redraw cycles (RPROMPT + syntax highlighting)
+      // that 24ms often split across coalesce boundaries.
+      const OUTPUT_COALESCE_MS = platform() === 'darwin' ? 32 : 8;
       // High-water mark: flush immediately when buffer exceeds this size
       // to prevent unbounded memory growth during high-output commands (e.g. `yes`)
       const OUTPUT_HIGH_WATER = 64 * 1024; // 64KB
@@ -288,6 +335,7 @@ export class PTYManager extends EventEmitter {
       };
 
       ptyProcess.onData((data: string) => {
+        session.scrollback.append(data);
         session.lastActivityAt = Date.now();
         session._outputBuffer += data;
         // Flush immediately if buffer exceeds high-water mark
@@ -310,7 +358,7 @@ export class PTYManager extends EventEmitter {
         this.terminals.delete(terminalId);
       });
 
-      return { shell: session.shell, cwd, isReconnect: false };
+      return { shell: session.shell, cwd, isReconnect: false, platform: platform() };
     } catch (error: unknown) {
       console.error(`[PTYManager] Failed to create terminal ${terminalId}:`, error);
       this.emit('error', terminalId, error instanceof Error ? error.message : 'Failed to create terminal');
@@ -387,6 +435,15 @@ export class PTYManager extends EventEmitter {
     }
   }
 
+  /**
+   * Get scrollback buffer contents for replay on reconnect
+   */
+  getScrollback(terminalId: string): string | null {
+    const session = this.terminals.get(terminalId);
+    if (!session) return null;
+    return session.scrollback.getContents();
+  }
+
   /**
    * Get terminal session info
    */

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;
   }