mstro-app 0.2.0 → 0.3.0

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

@@ -0,0 +1,403 @@
+// Copyright (c) 2025-present Mstro, Inc. All rights reserved.
+// Licensed under the MIT License. See LICENSE file for details.
+
+import { dirname, join } from 'node:path';
+import { executeGitCommand, handleGitStatus, spawnWithOutput } from './git-handlers.js';
+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> = {
+    gitWorktreeList: () => handleGitWorktreeList(ctx, ws, tabId, gitDir),
+    gitWorktreeCreate: () => handleGitWorktreeCreate(ctx, ws, msg, tabId, gitDir),
+    gitWorktreeRemove: () => handleGitWorktreeRemove(ctx, ws, msg, tabId, gitDir),
+    tabWorktreeSwitch: () => handleTabWorktreeSwitch(ctx, ws, msg, tabId, workingDir),
+    gitWorktreePush: () => handleGitWorktreePush(ctx, ws, msg, tabId, gitDir),
+    gitWorktreeCreatePR: () => handleGitWorktreeCreatePR(ctx, ws, msg, tabId, gitDir),
+    gitMergePreview: () => handleGitMergePreview(ctx, ws, msg, tabId, gitDir),
+    gitWorktreeMerge: () => handleGitWorktreeMerge(ctx, ws, msg, tabId, gitDir),
+    gitMergeAbort: () => handleGitMergeAbort(ctx, ws, tabId, gitDir),
+    gitMergeComplete: () => handleGitMergeComplete(ctx, ws, msg, tabId, gitDir),
+  };
+  handlers[msg.type]?.();
+}
+
+function applyWorktreePorcelainLine(line: string, worktrees: WorktreeInfo[], current: Partial<WorktreeInfo>): Partial<WorktreeInfo> {
+  if (line.startsWith('worktree ')) {
+    if (current.path) worktrees.push(current as WorktreeInfo);
+    return { path: line.slice(9).trim(), isMain: false, isBare: false };
+  }
+  if (line.startsWith('HEAD ')) current.head = line.slice(5).trim();
+  else if (line.startsWith('branch ')) current.branch = line.slice(7).trim().replace('refs/heads/', '');
+  else if (line === 'bare') current.isBare = true;
+  else if (line === 'prunable') current.prunable = true;
+  return current;
+}
+
+function parseWorktreePorcelain(stdout: string): WorktreeInfo[] {
+  const worktrees: WorktreeInfo[] = [];
+  let current: Partial<WorktreeInfo> = {};
+
+  for (const line of stdout.split('\n')) {
+    current = applyWorktreePorcelainLine(line, worktrees, current);
+  }
+  if (current.path) worktrees.push(current as WorktreeInfo);
+  if (worktrees.length > 0) worktrees[0].isMain = true;
+  return worktrees;
+}
+
+async function handleGitWorktreeList(ctx: HandlerContext, ws: WSContext, tabId: string, workingDir: string): Promise<void> {
+  try {
+    const result = await executeGitCommand(['worktree', 'list', '--porcelain'], workingDir);
+    if (result.exitCode !== 0) {
+      ctx.send(ws, { type: 'gitError', tabId, data: { error: result.stderr || 'Failed to list worktrees' } });
+      return;
+    }
+    const worktrees = parseWorktreePorcelain(result.stdout);
+    ctx.send(ws, { type: 'gitWorktreeListResult', tabId, data: { worktrees } });
+  } catch (error: any) {
+    ctx.send(ws, { type: 'gitError', tabId, data: { error: error.message } });
+  }
+}
+
+async function handleGitWorktreeCreate(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: 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 = workingDir.split('/').pop() || 'repo';
+    const wtPath = worktreePath || join(dirname(workingDir), `${repoName}-worktrees`, branchName);
+
+    const args = ['worktree', 'add', wtPath, '-b', branchName, ...(baseBranch ? [baseBranch] : [])];
+    const result = await executeGitCommand(args, workingDir);
+    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);
+
+    ctx.send(ws, {
+      type: 'gitWorktreeCreated',
+      tabId,
+      data: { path: wtPath, branch: branchName, head: headResult.stdout.trim() },
+    });
+  } catch (error: any) {
+    ctx.send(ws, { type: 'gitError', tabId, data: { error: error.message } });
+  }
+}
+
+async function handleGitWorktreeRemove(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, workingDir: string): Promise<void> {
+  try {
+    const { path: wtPath, force, deleteBranch } = msg.data || {};
+    if (!wtPath) {
+      ctx.send(ws, { type: 'gitError', tabId, data: { error: 'Worktree path is required' } });
+      return;
+    }
+
+    let branchToDelete: string | undefined;
+    if (deleteBranch) {
+      const branchResult = await executeGitCommand(['rev-parse', '--abbrev-ref', 'HEAD'], wtPath);
+      branchToDelete = branchResult.stdout.trim();
+    }
+
+    const args = force ? ['worktree', 'remove', '--force', wtPath] : ['worktree', 'remove', wtPath];
+    const result = await executeGitCommand(args, workingDir);
+    if (result.exitCode !== 0) {
+      ctx.send(ws, { type: 'gitError', tabId, data: { error: result.stderr || 'Failed to remove worktree' } });
+      return;
+    }
+
+    if (branchToDelete && deleteBranch) {
+      await executeGitCommand(['branch', '-d', branchToDelete], workingDir);
+    }
+
+    await executeGitCommand(['worktree', 'prune'], workingDir);
+
+    ctx.send(ws, { type: 'gitWorktreeRemoved', tabId, data: { path: wtPath } });
+  } catch (error: any) {
+    ctx.send(ws, { type: 'gitError', tabId, data: { error: error.message } });
+  }
+}
+
+async function handleTabWorktreeSwitch(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, workingDir: string): Promise<void> {
+  try {
+    const { tabId: targetTabId, worktreePath } = msg.data || {};
+    const resolvedTabId = targetTabId || tabId;
+    if (!worktreePath) {
+      ctx.gitDirectories.delete(resolvedTabId);
+      ctx.send(ws, { type: 'tabWorktreeSwitched', tabId: resolvedTabId, data: { tabId: resolvedTabId, worktreePath: workingDir, branch: '' } });
+      handleGitStatus(ctx, ws, resolvedTabId, workingDir);
+      return;
+    }
+
+    ctx.gitDirectories.set(resolvedTabId, worktreePath);
+
+    const branchResult = await executeGitCommand(['rev-parse', '--abbrev-ref', 'HEAD'], worktreePath);
+    const branch = branchResult.stdout.trim();
+
+    ctx.send(ws, { type: 'tabWorktreeSwitched', tabId: resolvedTabId, data: { tabId: resolvedTabId, worktreePath, branch } });
+    handleGitStatus(ctx, ws, resolvedTabId, worktreePath);
+  } catch (error: any) {
+    ctx.send(ws, { type: 'gitError', tabId, data: { error: error.message } });
+  }
+}
+
+async function pushWithUpstreamRetry(
+  worktreePath: string,
+  pushRemote: string,
+  pushBranch: string,
+  setUpstream: boolean,
+): Promise<{ exitCode: number; output: string; error: string }> {
+  const args = setUpstream
+    ? ['push', '--set-upstream', pushRemote, pushBranch]
+    : ['push', pushRemote, pushBranch];
+
+  const result = await executeGitCommand(args, worktreePath);
+  if (result.exitCode === 0) return { exitCode: 0, output: result.stderr || result.stdout, error: '' };
+
+  const needsUpstream = result.stderr.includes('no upstream') || result.stderr.includes('has no upstream');
+  if (!needsUpstream) return { exitCode: result.exitCode, output: '', error: result.stderr || 'Failed to push' };
+
+  const retry = await executeGitCommand(['push', '--set-upstream', pushRemote, pushBranch], worktreePath);
+  if (retry.exitCode !== 0) return { exitCode: retry.exitCode, output: '', error: retry.stderr || 'Failed to push' };
+  return { exitCode: 0, output: retry.stderr || retry.stdout, error: '' };
+}
+
+async function handleGitWorktreePush(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, _workingDir: string): Promise<void> {
+  try {
+    const { worktreePath, remote, branch, setUpstream } = msg.data || {};
+    if (!worktreePath) {
+      ctx.send(ws, { type: 'gitError', tabId, data: { error: 'Worktree path is required' } });
+      return;
+    }
+
+    const pushRemote = remote || 'origin';
+    const branchResult = await executeGitCommand(['rev-parse', '--abbrev-ref', 'HEAD'], worktreePath);
+    const pushBranch = branch || branchResult.stdout.trim();
+
+    const pushResult = await pushWithUpstreamRetry(worktreePath, pushRemote, pushBranch, !!setUpstream);
+    if (pushResult.exitCode !== 0) {
+      ctx.send(ws, { type: 'gitError', tabId, data: { error: pushResult.error } });
+      return;
+    }
+    ctx.send(ws, { type: 'gitWorktreePushed', tabId, data: { output: pushResult.output, upstream: `${pushRemote}/${pushBranch}` } });
+  } catch (error: any) {
+    ctx.send(ws, { type: 'gitError', tabId, data: { error: error.message } });
+  }
+}
+
+async function handleGitWorktreeCreatePR(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, _workingDir: string): Promise<void> {
+  try {
+    const { worktreePath, title, body, baseBranch, draft } = msg.data || {};
+    if (!worktreePath) {
+      ctx.send(ws, { type: 'gitError', tabId, data: { error: 'Worktree path is required' } });
+      return;
+    }
+
+    const branchResult = await executeGitCommand(['rev-parse', '--abbrev-ref', 'HEAD'], worktreePath);
+    const branchName = branchResult.stdout.trim();
+    const prTitle = title || branchName.replace(/[-_/]/g, ' ').replace(/^\w/, c => c.toUpperCase());
+
+    const args = ['pr', 'create', '--title', prTitle];
+    if (body) args.push('--body', body);
+    if (baseBranch) args.push('--base', baseBranch);
+    if (draft) args.push('--draft');
+
+    const result = await spawnWithOutput('gh', args, worktreePath);
+    if (result.exitCode !== 0) {
+      ctx.send(ws, { type: 'gitError', tabId, data: { error: result.stderr || 'Failed to create PR' } });
+      return;
+    }
+
+    const prUrl = result.stdout.trim();
+    const prNumberMatch = prUrl.match(/\/(\d+)$/);
+    const prNumber = prNumberMatch ? parseInt(prNumberMatch[1], 10) : 0;
+
+    ctx.send(ws, { type: 'gitWorktreePRCreated', tabId, data: { prUrl, prNumber } });
+  } catch (error: any) {
+    ctx.send(ws, { type: 'gitError', tabId, data: { error: error.message } });
+  }
+}
+
+async function handleGitMergePreview(ctx: HandlerContext, ws: WSContext, msg: WebSocketMessage, tabId: string, workingDir: string): Promise<void> {
+  try {
+    const { sourceBranch, targetBranch } = msg.data || {};
+    if (!sourceBranch || !targetBranch) {
+      ctx.send(ws, { type: 'gitError', tabId, data: { error: 'Source and target branches are required' } });
+      return;
+    }
+
+    let clean = true;
+    let conflicts: string[] = [];
+    const mergeTreeResult = await executeGitCommand(['merge-tree', '--write-tree', targetBranch, sourceBranch], workingDir);
+    if (mergeTreeResult.exitCode !== 0) {
+      clean = false;
+      conflicts = mergeTreeResult.stdout.split('\n')
+        .filter(line => line.includes('CONFLICT'))
+        .map(line => {
+          const match = line.match(/CONFLICT.*:\s+(.+)/);
+          return match?.[1]?.trim() || line;
+        });
+    }
+
+    const statResult = await executeGitCommand(['diff', `${targetBranch}...${sourceBranch}`, '--stat'], workingDir);
+    const stat = statResult.stdout.trim();
+
+    const logResult = await executeGitCommand(
+      ['log', `${targetBranch}..${sourceBranch}`, '--oneline', '--format=%h|%s'],
+      workingDir
+    );
+    const commits = logResult.stdout.trim().split('\n')
+      .filter(line => line.trim())
+      .map(line => {
+        const [hash, ...rest] = line.split('|');
+        return { hash: hash.trim(), message: rest.join('|').trim() };
+      });
+
+    ctx.send(ws, {
+      type: 'gitMergePreviewResult',
+      tabId,
+      data: { clean, conflicts, stat, commits, ahead: commits.length },
+    });
+  } catch (error: any) {
+    ctx.send(ws, { type: 'gitError', tabId, data: { error: error.message } });
+  }
+}
+
+async function resolveMainWorktreePath(workingDir: string): Promise<string> {
+  const wtListResult = await executeGitCommand(['worktree', 'list', '--porcelain'], workingDir);
+  const firstLine = wtListResult.stdout.split('\n').find(l => l.startsWith('worktree '));
+  return firstLine ? firstLine.slice(9).trim() : workingDir;
+}
+
+async function executeMergeStrategy(
+  strategy: string,
+  sourceBranch: string,
+  commitMessage: string | undefined,
+  mainPath: string,
+): Promise<{ exitCode: number; error?: string }> {
+  if (strategy === 'squash') {
+    const squashResult = await executeGitCommand(['merge', '--squash', sourceBranch], mainPath);
+    if (squashResult.exitCode !== 0) return { exitCode: squashResult.exitCode, error: squashResult.stderr };
+    const msg2 = commitMessage || `Squash merge branch '${sourceBranch}'`;
+    const commitResult = await executeGitCommand(['commit', '-m', msg2], mainPath);
+    if (commitResult.exitCode !== 0) return { exitCode: commitResult.exitCode, error: commitResult.stderr || 'Failed to commit squash merge' };
+    return { exitCode: 0 };
+  }
+  if (strategy === 'rebase') {
+    const result = await executeGitCommand(['merge', '--ff-only', sourceBranch], mainPath);
+    return { exitCode: result.exitCode, error: result.stderr };
+  }
+  const mergeArgs = commitMessage ? ['merge', sourceBranch, '-m', commitMessage] : ['merge', sourceBranch];
+  const result = await executeGitCommand(mergeArgs, mainPath);
+  return { exitCode: result.exitCode, error: result.stderr };
+}
+
+async function detectMergeConflicts(mainPath: string): Promise<string[]> {
+  const result = await executeGitCommand(['diff', '--name-only', '--diff-filter=U'], mainPath);
+  return result.stdout.trim().split('\n').filter(f => f.trim());
+}
+
+async function cleanupAfterMerge(
+  mainPath: string,
+  sourceBranch: string,
+  strategy: string,
+  deleteWorktree: boolean,
+  deleteBranch: boolean,
+): Promise<void> {
+  if (deleteWorktree) {
+    const wtList = await executeGitCommand(['worktree', 'list', '--porcelain'], mainPath);
+    const worktreePath = findWorktreePathForBranch(wtList.stdout, sourceBranch);
+    if (worktreePath && worktreePath !== mainPath) {
+      await executeGitCommand(['worktree', 'remove', worktreePath], mainPath);
+    }
+  }
+  if (deleteBranch) {
+    const deleteFlag = strategy === 'squash' ? '-D' : '-d';
+    await executeGitCommand(['branch', deleteFlag, sourceBranch], mainPath);
+  }
+  await executeGitCommand(['worktree', 'prune'], mainPath);
+}
+
+function findWorktreePathForBranch(porcelainOutput: string, branchName: string): string | null {
+  let currentWtPath = '';
+  for (const line of porcelainOutput.split('\n')) {
+    if (line.startsWith('worktree ')) currentWtPath = line.slice(9).trim();
+    if (line.startsWith('branch ') && line.includes(branchName)) return currentWtPath;
+  }
+  return null;
+}
+
+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 || {};
+    if (!sourceBranch || !targetBranch || !strategy) {
+      ctx.send(ws, { type: 'gitError', tabId, data: { error: 'Source branch, target branch, and strategy are required' } });
+      return;
+    }
+
+    const mainPath = await resolveMainWorktreePath(workingDir);
+
+    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` } });
+      return;
+    }
+
+    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' };
+      ctx.send(ws, { type: 'gitWorktreeMergeResult', tabId, data });
+      return;
+    }
+
+    const commitHashResult = await executeGitCommand(['rev-parse', '--short', 'HEAD'], mainPath);
+    await cleanupAfterMerge(mainPath, sourceBranch, strategy, !!deleteWorktree, !!deleteBranch);
+
+    ctx.send(ws, { type: 'gitWorktreeMergeResult', tabId, data: { success: true, mergeCommit: commitHashResult.stdout.trim() } });
+  } catch (error: any) {
+    ctx.send(ws, { type: 'gitError', tabId, data: { error: error.message } });
+  }
+}
+
+async function handleGitMergeAbort(ctx: HandlerContext, ws: WSContext, tabId: string, workingDir: string): Promise<void> {
+  try {
+    const mainPath = await resolveMainWorktreePath(workingDir);
+
+    const result = await executeGitCommand(['merge', '--abort'], mainPath);
+    if (result.exitCode !== 0) {
+      ctx.send(ws, { type: 'gitError', tabId, data: { error: result.stderr || 'Failed to abort merge' } });
+      return;
+    }
+
+    ctx.send(ws, { type: 'gitMergeAborted', tabId, data: { aborted: true } });
+  } catch (error: any) {
+    ctx.send(ws, { type: 'gitError', tabId, data: { error: error.message } });
+  }
+}
+
+async function handleGitMergeComplete(ctx: HandlerContext, ws: WSContext, _msg: WebSocketMessage, tabId: string, workingDir: string): Promise<void> {
+  try {
+    const mainPath = await resolveMainWorktreePath(workingDir);
+
+    const result = await executeGitCommand(['commit', '--no-edit'], mainPath);
+    if (result.exitCode !== 0) {
+      ctx.send(ws, { type: 'gitError', tabId, data: { error: result.stderr || 'Failed to complete merge' } });
+      return;
+    }
+
+    const hashResult = await executeGitCommand(['rev-parse', '--short', 'HEAD'], mainPath);
+    const mergeCommit = hashResult.stdout.trim();
+
+    ctx.send(ws, { type: 'gitMergeCompleted', tabId, data: { success: true, mergeCommit } });
+  } catch (error: any) {
+    ctx.send(ws, { type: 'gitError', tabId, data: { error: error.message } });
+  }
+}

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

@@ -0,0 +1,44 @@
+// 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 { SessionRegistry } from './session-registry.js';
+import type { WebSocketResponse, WSContext } from './types.js';
+
+export interface UsageReport {
+  tokensUsed: number;
+  sessionId?: string;
+  movementId?: string;
+}
+
+export type UsageReporter = (report: UsageReport) => void;
+
+/**
+ * Shared context passed to all domain handler functions.
+ * The WebSocketImproviseHandler class satisfies this interface directly.
+ */
+export interface HandlerContext {
+  // Shared state
+  sessions: Map<string, ImprovisationSessionManager>;
+  connections: Map<WSContext, Map<string, string>>;
+  allConnections: Set<WSContext>;
+  gitDirectories: Map<string, string>;
+  activeSearches: Map<string, ChildProcess>;
+  terminalSubscribers: Map<string, Set<WSContext>>;
+  terminalListenerCleanups: Map<string, () => void>;
+  autocompleteService: AutocompleteService;
+  usageReporter: UsageReporter | null;
+
+  // Registry access
+  getRegistry(workingDir: string): SessionRegistry;
+
+  // Communication utilities
+  send(ws: WSContext, response: WebSocketResponse): void;
+  broadcastToOthers(sender: WSContext, response: WebSocketResponse): void;
+  broadcastToAll(response: WebSocketResponse): void;
+
+  // Frecency
+  recordFileSelection(filePath: string): void;
+}

package/server/services/websocket/handler.test.ts

@@ -15,6 +15,6 @@ describe('WebSocket handler code quality', () => {
   })
 
   it('imports mkdirSync from fs at the top level', () => {
-    expect(handlerSource).toContain("import { existsSync, mkdirSync, readdirSync, readFileSync, unlinkSync, writeFileSync } from 'node:fs'")
+    expect(handlerSource).toContain("import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'node:fs'")
   })
 })