mstro-app 0.4.52 → 0.5.1

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

@@ -1,22 +1,41 @@
 // Copyright (c) 2025-present Mstro, Inc. All rights reserved.
-// Licensed under the MIT License. See LICENSE file for details.
 
 /**
  * Improvisation Session Manager v2
  *
  * Optimized for fast, direct prompt execution in Improvise mode.
  * For complex multi-part prompts with parallel/sequential movements, use Compose tab instead.
+ *
+ * Delegates to focused helpers:
+ * - improvisation-output-queue.ts  — buffered stdout flush loop
+ * - improvisation-history-store.ts — .mstro/history/*.json load/save
+ * - improvisation-movements.ts     — pure movement-record builders
+ * - improvisation-retry.ts         — retry decision tree + recovery strategies
  */
 
 import { EventEmitter } from 'node:events';
-import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'node:fs';
+import { existsSync, readFileSync } from 'node:fs';
 import { join } from 'node:path';
 import { AnalyticsEvents, trackEvent } from '../services/analytics.js';
 import { herror } from './headless/headless-logger.js';
 import { cleanupAttachments, preparePromptAndAttachments } from './improvisation-attachments.js';
+import {
+  ensureHistoryDir,
+  loadHistory,
+  resolveHistoryPaths,
+  saveHistory,
+} from './improvisation-history-store.js';
+import {
+  buildCancelledMovement,
+  buildErrorMovement,
+  buildSuccessMovement,
+  CANCELLED_FALLBACK_RESULT,
+  shouldAutoContinue,
+} from './improvisation-movements.js';
+import { OutputQueue } from './improvisation-output-queue.js';
 import type { RetryCallbacks, RetrySessionState } from './improvisation-retry.js';
-import {applyToolTimeoutRetry, 
-  createExecutionRunner,detectNativeTimeoutContextLoss, detectResumeContextLoss, 
+import {applyToolTimeoutRetry,
+  createExecutionRunner,detectNativeTimeoutContextLoss, detectResumeContextLoss,
   determineResumeStrategy,
   selectBestResult,
   shouldRetryContextLoss,
@@ -40,8 +59,7 @@ export class ImprovisationSessionManager extends EventEmitter {
     plan: unknown;
     resolve: (approved: boolean) => void;
   };
-  private outputQueue: Array<{ text: string; timestamp: number }> = [];
-  private queueTimer: NodeJS.Timeout | null = null;
+  private outputBuffer: OutputQueue;
   private isFirstPrompt: boolean = true;
   private claudeSessionId: string | undefined;
   private isResumedSession: boolean = false;
@@ -54,6 +72,7 @@ export class ImprovisationSessionManager extends EventEmitter {
   private _cancelCompleteEmitted: boolean = false;
   private _currentUserPrompt: string = '';
   private _currentSequenceNumber: number = 0;
+  private _hasPersistedToDisk: boolean = false;
 
   static resumeFromHistory(workingDir: string, historicalSessionId: string, overrides?: Partial<ImprovisationOptions>): ImprovisationSessionManager {
     const historyDir = join(workingDir, '.mstro', 'history');
@@ -79,6 +98,7 @@ export class ImprovisationSessionManager extends EventEmitter {
 
     manager.isResumedSession = true;
     manager.isFirstPrompt = true;
+    manager._hasPersistedToDisk = true;
     if (historyData.claudeSessionId) {
       manager.claudeSessionId = historyData.claudeSessionId;
     }
@@ -101,33 +121,29 @@ export class ImprovisationSessionManager extends EventEmitter {
     };
 
     this.sessionId = this.options.sessionId;
-    this.improviseDir = join(this.options.workingDir, '.mstro', 'history');
-    this.historyPath = join(this.improviseDir, `${this.sessionId.replace('improv-', '')}.json`);
+    const paths = resolveHistoryPaths(this.options.workingDir, this.sessionId);
+    this.improviseDir = paths.improviseDir;
+    this.historyPath = paths.historyPath;
+    ensureHistoryDir(this.improviseDir);
 
-    if (!existsSync(this.improviseDir)) {
-      mkdirSync(this.improviseDir, { recursive: true });
-    }
+    this.history = loadHistory(this.historyPath, this.sessionId);
+    // History is persisted lazily on the first `persistHistory` call (see
+    // `executePrompt`). Deferring the initial write keeps the Chat History
+    // view from showing "0 prompts" entries for tabs the user opens but
+    // never prompts.
 
-    this.history = this.loadHistory();
-    this.saveHistory(); // Persist immediately so the session file exists on disk from creation
-    this.startQueueProcessor();
+    this.outputBuffer = new OutputQueue(text => this.emit('onOutput', text));
+    this.outputBuffer.start();
   }
 
   // ========== Output Queue ==========
 
-  private startQueueProcessor(): void {
-    this.queueTimer = setInterval(() => { this.flushOutputQueue(); }, 50);
-  }
-
   private queueOutput(text: string): void {
-    this.outputQueue.push({ text, timestamp: Date.now() });
+    this.outputBuffer.queue_(text);
   }
 
   private flushOutputQueue(): void {
-    while (this.outputQueue.length > 0) {
-      const item = this.outputQueue.shift();
-      if (item) this.emit('onOutput', item.text);
-    }
+    this.outputBuffer.flush();
   }
 
   // ========== Main Execution ==========
@@ -174,7 +190,7 @@ export class ImprovisationSessionManager extends EventEmitter {
       ...(isAutoContinue && { isAutoContinue: true }),
     };
     this.history.movements.push(pendingMovement);
-    this.saveHistory();
+    this.persistHistory();
 
     try {
       this.executionEventLog.push({
@@ -212,7 +228,11 @@ export class ImprovisationSessionManager extends EventEmitter {
       this.captureSessionAndSurfaceErrors(result);
       this.isFirstPrompt = false;
 
-      const movement = this.buildMovementRecord(result, displayPrompt, sequenceNumber, _execStart, state.retryLog, isAutoContinue);
+      const movement = buildSuccessMovement(
+        result,
+        { sequenceNumber, userPrompt: displayPrompt, execStart: _execStart, isAutoContinue },
+        state.retryLog,
+      );
       this.handleConflicts(result);
       this.persistMovement(movement);
 
@@ -345,7 +365,14 @@ export class ImprovisationSessionManager extends EventEmitter {
     return false;
   }
 
-  // ========== Cancel Handling ==========
+  // ========== Cancel / Error Handling ==========
+
+  private resetExecutionState(): void {
+    this._isExecuting = false;
+    this._executionStartTimestamp = undefined;
+    this.executionEventLog = [];
+    this.currentRunner = null;
+  }
 
   private handleCancelledExecution(
     result: HeadlessRunResult | undefined,
@@ -353,39 +380,16 @@ export class ImprovisationSessionManager extends EventEmitter {
     sequenceNumber: number,
     execStart: number,
   ): MovementRecord {
-    this._isExecuting = false;
-    this._executionStartTimestamp = undefined;
-    this.executionEventLog = [];
-    this.currentRunner = null;
+    this.resetExecutionState();
 
     if (this._cancelCompleteEmitted) {
       const existing = this.history.movements.find(m => m.sequenceNumber === sequenceNumber);
       if (existing) return existing;
     }
 
-    const cancelledMovement: MovementRecord = {
-      id: `prompt-${sequenceNumber}`,
-      sequenceNumber,
-      userPrompt,
-      timestamp: new Date().toISOString(),
-      tokensUsed: result ? result.totalTokens : 0,
-      summary: '',
-      filesModified: [],
-      assistantResponse: result?.assistantResponse,
-      thinkingOutput: result?.thinkingOutput,
-      toolUseHistory: result?.toolUseHistory?.map(t => ({
-        toolName: t.toolName, toolId: t.toolId, toolInput: t.toolInput,
-        result: t.result,
-      })),
-      errorOutput: 'Execution cancelled by user',
-      durationMs: Date.now() - execStart,
-    };
+    const cancelledMovement = buildCancelledMovement(result, { sequenceNumber, userPrompt, execStart });
     this.persistMovement(cancelledMovement);
-    const fallbackResult = {
-      completed: false, needsHandoff: false, totalTokens: 0, sessionId: '',
-      output: '', exitCode: 1, signalName: 'SIGTERM',
-    } as HeadlessRunResult;
-    this.emitMovementComplete(cancelledMovement, result ?? fallbackResult, execStart, sequenceNumber);
+    this.emitMovementComplete(cancelledMovement, result ?? CANCELLED_FALLBACK_RESULT, execStart, sequenceNumber);
     return cancelledMovement;
   }
 
@@ -395,23 +399,10 @@ export class ImprovisationSessionManager extends EventEmitter {
     sequenceNumber: number,
     execStart: number,
   ): never {
-    this._isExecuting = false;
-    this._executionStartTimestamp = undefined;
-    this.executionEventLog = [];
-    this.currentRunner = null;
+    this.resetExecutionState();
 
     const errorMessage = error instanceof Error ? error.message : String(error);
-    const errorMovement: MovementRecord = {
-      id: `prompt-${sequenceNumber}`,
-      sequenceNumber,
-      userPrompt: displayPrompt,
-      timestamp: new Date().toISOString(),
-      tokensUsed: 0,
-      summary: '',
-      filesModified: [],
-      errorOutput: errorMessage,
-      durationMs: Date.now() - execStart,
-    };
+    const errorMovement = buildErrorMovement(errorMessage, { sequenceNumber, userPrompt: displayPrompt, execStart });
     this.persistMovement(errorMovement);
 
     this.emit('onMovementError', error);
@@ -440,35 +431,6 @@ export class ImprovisationSessionManager extends EventEmitter {
     }
   }
 
-  private buildMovementRecord(
-    result: HeadlessRunResult,
-    userPrompt: string,
-    sequenceNumber: number,
-    execStart: number,
-    retryLog?: import('./improvisation-types.js').RetryLogEntry[],
-    isAutoContinue?: boolean,
-  ): MovementRecord {
-    return {
-      id: `prompt-${sequenceNumber}`,
-      sequenceNumber,
-      userPrompt,
-      timestamp: new Date().toISOString(),
-      tokensUsed: result.totalTokens,
-      summary: '',
-      filesModified: [],
-      assistantResponse: result.assistantResponse,
-      thinkingOutput: result.thinkingOutput,
-      toolUseHistory: result.toolUseHistory?.map(t => ({
-        toolName: t.toolName, toolId: t.toolId, toolInput: t.toolInput,
-        result: t.result, isError: t.isError, duration: t.duration,
-      })),
-      errorOutput: result.error,
-      durationMs: Date.now() - execStart,
-      retryLog: retryLog && retryLog.length > 0 ? retryLog : undefined,
-      ...(isAutoContinue && { isAutoContinue: true }),
-    };
-  }
-
   private handleConflicts(result: HeadlessRunResult): void {
     if (!result.conflicts || result.conflicts.length === 0) return;
     this.queueOutput(`\n⚠ File conflicts detected: ${result.conflicts.length}`);
@@ -489,7 +451,7 @@ export class ImprovisationSessionManager extends EventEmitter {
       this.history.movements.push(movement);
       this.history.totalTokens += movement.tokensUsed;
     }
-    this.saveHistory();
+    this.persistHistory();
   }
 
   private emitMovementComplete(movement: MovementRecord, result: HeadlessRunResult, execStart: number, sequenceNumber: number): void {
@@ -515,26 +477,10 @@ export class ImprovisationSessionManager extends EventEmitter {
     const isStallKill = !this._cancelled && !!result.signalName;
     if (isStallKill && this._autoContinueCount < ImprovisationSessionManager.MAX_AUTO_CONTINUES) {
       this.scheduleAutoContinue('Process stalled');
-    } else if (this.shouldAutoContinue(result, userPrompt)) {
+    } else if (shouldAutoContinue(result, this._autoContinueCount, ImprovisationSessionManager.MAX_AUTO_CONTINUES, this._cancelled)) {
       this.scheduleAutoContinue();
     }
-  }
-
-  private shouldAutoContinue(result: HeadlessRunResult, _userPrompt: string): boolean {
-    if (this._autoContinueCount >= ImprovisationSessionManager.MAX_AUTO_CONTINUES) return false;
-    if (this._cancelled) return false;
-    if (!result.completed || result.signalName) return false;
-    if (result.stopReason !== 'end_turn') return false;
-
-    const thinkingLen = result.thinkingOutput?.length ?? 0;
-    const responseLen = result.assistantResponse?.length ?? 0;
-    const successfulToolCalls = result.toolUseHistory?.filter(t => t.result !== undefined && !t.isError).length ?? 0;
-
-    if (thinkingLen < 500 || responseLen > 1000) return false;
-    // When the agent executed tool calls and produced a non-trivial response,
-    // long thinking is expected — the work happened in the tools, not the text.
-    if (successfulToolCalls > 0 && responseLen > 200) return false;
-    return thinkingLen >= responseLen * 3;
+    void userPrompt;
   }
 
   private scheduleAutoContinue(reason?: string): void {
@@ -555,27 +501,12 @@ export class ImprovisationSessionManager extends EventEmitter {
 
   // ========== History I/O ==========
 
-  private loadHistory(): SessionHistory {
-    if (existsSync(this.historyPath)) {
-      try {
-        const data = readFileSync(this.historyPath, 'utf-8');
-        return JSON.parse(data);
-      } catch (error) {
-        herror('Failed to load history:', error);
-      }
+  private persistHistory(): void {
+    saveHistory(this.historyPath, this.history);
+    if (!this._hasPersistedToDisk) {
+      this._hasPersistedToDisk = true;
+      this.emit('onHistoryPersisted');
     }
-    return {
-      sessionId: this.sessionId,
-      startedAt: new Date().toISOString(),
-      lastActivityAt: new Date().toISOString(),
-      totalTokens: 0,
-      movements: [],
-    };
-  }
-
-  private saveHistory(): void {
-    this.history.lastActivityAt = new Date().toISOString();
-    writeFileSync(this.historyPath, JSON.stringify(this.history, null, 2));
   }
 
   getHistory(): SessionHistory {
@@ -592,7 +523,7 @@ export class ImprovisationSessionManager extends EventEmitter {
       this.currentRunner = null;
     }
 
-    this.destroyQueueTimer();
+    this.outputBuffer.destroy();
 
     if (this._isExecuting && !this._cancelCompleteEmitted) {
       this._cancelCompleteEmitted = true;
@@ -600,38 +531,20 @@ export class ImprovisationSessionManager extends EventEmitter {
       this._isExecuting = false;
       this._executionStartTimestamp = undefined;
 
-      const cancelledMovement: MovementRecord = {
-        id: `prompt-${this._currentSequenceNumber}`,
+      const cancelledMovement = buildCancelledMovement(undefined, {
         sequenceNumber: this._currentSequenceNumber,
         userPrompt: this._currentUserPrompt,
-        timestamp: new Date().toISOString(),
-        tokensUsed: 0,
-        summary: '',
-        filesModified: [],
-        errorOutput: 'Execution cancelled by user',
-        durationMs: Date.now() - execStart,
-      };
+        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.emitMovementComplete(cancelledMovement, CANCELLED_FALLBACK_RESULT, execStart, this._currentSequenceNumber);
     }
 
     this.flushOutputQueue();
   }
 
-  private destroyQueueTimer(): void {
-    if (this.queueTimer) {
-      clearInterval(this.queueTimer);
-      this.queueTimer = null;
-    }
-  }
-
   destroy(): void {
-    this.destroyQueueTimer();
+    this.outputBuffer.destroy();
     this.flushOutputQueue();
   }
 
@@ -642,7 +555,7 @@ export class ImprovisationSessionManager extends EventEmitter {
     this.isFirstPrompt = true;
     this.claudeSessionId = undefined;
     cleanupAttachments(this.options.workingDir, this.sessionId);
-    this.saveHistory();
+    this.persistHistory();
     this.emit('onSessionUpdate', this.getHistory());
   }
 
@@ -684,7 +597,7 @@ export class ImprovisationSessionManager extends EventEmitter {
   }
 
   startNewSession(overrides?: Partial<ImprovisationOptions>): ImprovisationSessionManager {
-    this.saveHistory();
+    this.persistHistory();
     return new ImprovisationSessionManager({
       ...this.options,
       sessionId: `improv-${Date.now()}`,

package/server/cli/improvisation-types.ts

@@ -1,5 +1,4 @@
 // Copyright (c) 2025-present Mstro, Inc. All rights reserved.
-// Licensed under the MIT License. See LICENSE file for details.
 
 /**
  * Types and interfaces for the Improvisation Session Manager.

package/server/cli/retry/retry-best-result.ts

@@ -0,0 +1,69 @@
+// Copyright (c) 2025-present Mstro, Inc. All rights reserved.
+
+/**
+ * Picks the best result across retry attempts. Prefers Haiku's judgment
+ * when available; falls back to a numeric score when Haiku is unreachable.
+ */
+
+import { hlog } from '../headless/headless-logger.js';
+import { assessBestResult } from '../headless/stall-assessor.js';
+import type { HeadlessRunResult, RetryLoopState } from '../improvisation-types.js';
+import { scoreRunResult } from '../improvisation-types.js';
+
+/** Select the best result across retries using Haiku assessment */
+export async function selectBestResult(
+  state: RetryLoopState,
+  result: HeadlessRunResult,
+  userPrompt: string,
+  verbose: boolean,
+): Promise<HeadlessRunResult> {
+  if (!state.bestResult || state.bestResult === result || state.retryNumber === 0) {
+    return result;
+  }
+
+  const claudeCmd = process.env.CLAUDE_COMMAND || 'claude';
+  const bestToolCount = state.bestResult.toolUseHistory?.filter(t => t.result !== undefined && !t.isError).length ?? 0;
+  const currentToolCount = result.toolUseHistory?.filter(t => t.result !== undefined && !t.isError).length ?? 0;
+
+  try {
+    const verdict = await assessBestResult({
+      originalPrompt: userPrompt,
+      resultA: {
+        successfulToolCalls: bestToolCount,
+        responseLength: state.bestResult.assistantResponse?.length ?? 0,
+        hasThinking: !!state.bestResult.thinkingOutput,
+        responseTail: (state.bestResult.assistantResponse ?? '').slice(-500),
+      },
+      resultB: {
+        successfulToolCalls: currentToolCount,
+        responseLength: result.assistantResponse?.length ?? 0,
+        hasThinking: !!result.thinkingOutput,
+        responseTail: (result.assistantResponse ?? '').slice(-500),
+      },
+    }, claudeCmd, verbose);
+
+    if (verdict.winner === 'A') {
+      if (verbose) hlog(`[BEST-RESULT] Haiku picked earlier attempt: ${verdict.reason}`);
+      return mergeResultSessionId(state.bestResult, result.claudeSessionId);
+    }
+    if (verbose) hlog(`[BEST-RESULT] Haiku picked final attempt: ${verdict.reason}`);
+    return result;
+  } catch {
+    return fallbackBestResult(state.bestResult, result, verbose);
+  }
+}
+
+function mergeResultSessionId(result: HeadlessRunResult, sessionId: string | undefined): HeadlessRunResult {
+  if (sessionId) return { ...result, claudeSessionId: sessionId };
+  return result;
+}
+
+function fallbackBestResult(bestResult: HeadlessRunResult, result: HeadlessRunResult, verbose: boolean): HeadlessRunResult {
+  if (scoreRunResult(bestResult) > scoreRunResult(result)) {
+    if (verbose) {
+      hlog(`[BEST-RESULT] Haiku unavailable, numeric fallback: earlier attempt (score ${scoreRunResult(bestResult)} vs ${scoreRunResult(result)})`);
+    }
+    return mergeResultSessionId(bestResult, result.claudeSessionId);
+  }
+  return result;
+}

package/server/cli/retry/retry-context-loss.ts

@@ -0,0 +1,86 @@
+// Copyright (c) 2025-present Mstro, Inc. All rights reserved.
+
+/**
+ * Context-loss detection: figures out when a run's output indicates the
+ * Claude session dropped its memory, either on `--resume` (Path 1) or after
+ * native tool timeouts scrambled the conversation (Path 2).
+ */
+
+import { hlog } from '../headless/headless-logger.js';
+import { assessContextLoss, type ContextLossContext } from '../headless/stall-assessor.js';
+import type { HeadlessRunResult, RetryLoopState } from '../improvisation-types.js';
+
+const WRITE_TOOL_NAMES = new Set(['Edit', 'Write', 'MultiEdit', 'NotebookEdit']);
+
+/** Detect resume context loss (Path 1): session expired on --resume */
+export function detectResumeContextLoss(
+  result: HeadlessRunResult,
+  state: RetryLoopState,
+  useResume: boolean,
+  maxRetries: number,
+  nativeTimeouts: number,
+  verbose: boolean,
+): void {
+  if (!useResume || state.checkpointRef.value || state.retryNumber >= maxRetries || nativeTimeouts > 0) {
+    return;
+  }
+  if (!result.assistantResponse || result.assistantResponse.trim().length === 0) {
+    state.contextLost = true;
+    if (verbose) hlog('[CONTEXT-RECOVERY] Resume context loss: null/empty response');
+  } else if (result.resumeBufferedOutput !== undefined) {
+    state.contextLost = true;
+    if (verbose) hlog('[CONTEXT-RECOVERY] Resume context loss: buffer never flushed (no thinking/tools)');
+  } else if (
+    (!result.toolUseHistory || result.toolUseHistory.length === 0) &&
+    !result.thinkingOutput &&
+    result.assistantResponse.length < 500
+  ) {
+    state.contextLost = true;
+    if (verbose) hlog('[CONTEXT-RECOVERY] Resume context loss: no tools, no thinking, short response');
+  }
+}
+
+/** Detect native timeout context loss (Path 2): tool timeouts caused confusion */
+export async function detectNativeTimeoutContextLoss(
+  result: HeadlessRunResult,
+  state: RetryLoopState,
+  maxRetries: number,
+  nativeTimeouts: number,
+  verbose: boolean,
+): Promise<void> {
+  if (state.contextLost) return;
+
+  const { effectiveTimeouts } = computeEffectiveTimeouts(result, nativeTimeouts);
+  if (effectiveTimeouts === 0 || !result.assistantResponse || state.checkpointRef.value || state.retryNumber >= maxRetries) {
+    return;
+  }
+
+  const contextLossCtx: ContextLossContext = {
+    assistantResponse: result.assistantResponse,
+    effectiveTimeouts,
+    nativeTimeoutCount: nativeTimeouts,
+    successfulToolCalls: result.toolUseHistory?.filter(t => t.result !== undefined && !t.isError).length ?? 0,
+    thinkingOutputLength: result.thinkingOutput?.length ?? 0,
+    hasSuccessfulWrite: result.toolUseHistory?.some(
+      t => WRITE_TOOL_NAMES.has(t.toolName) && t.result !== undefined && !t.isError
+    ) ?? false,
+  };
+
+  const claudeCmd = process.env.CLAUDE_COMMAND || 'claude';
+  const verdict = await assessContextLoss(contextLossCtx, claudeCmd, verbose);
+  state.contextLost = verdict.contextLost;
+  if (verbose) {
+    hlog(`[CONTEXT-RECOVERY] Haiku verdict: ${state.contextLost ? 'LOST' : 'OK'} — ${verdict.reason}`);
+  }
+}
+
+function computeEffectiveTimeouts(result: HeadlessRunResult, nativeTimeouts: number): { effectiveTimeouts: number } {
+  const succeededIds = new Set<string>();
+  const allIds = new Set<string>();
+  for (const t of result.toolUseHistory ?? []) {
+    allIds.add(t.toolId);
+    if (t.result !== undefined) succeededIds.add(t.toolId);
+  }
+  const toolsWithoutResult = [...allIds].filter(id => !succeededIds.has(id)).length;
+  return { effectiveTimeouts: Math.max(nativeTimeouts, toolsWithoutResult) };
+}

package/server/cli/retry/retry-premature-completion.ts

@@ -0,0 +1,112 @@
+// Copyright (c) 2025-present Mstro, Inc. All rights reserved.
+
+/**
+ * Detects when a claimed-complete run is actually unfinished (hit
+ * max_tokens, abandoned mid-task, or Haiku says the end_turn response is
+ * a stop short of the goal) and triggers a continuation retry.
+ */
+
+import { AnalyticsEvents, trackEvent } from '../../services/analytics.js';
+import { hlog } from '../headless/headless-logger.js';
+import { extractFinalTextBlock, isResponseAbandoned } from '../headless/retry-strategies.js';
+import { assessPrematureCompletion } from '../headless/stall-assessor.js';
+import type { HeadlessRunResult, RetryLoopState } from '../improvisation-types.js';
+import type { RetryCallbacks, RetrySessionState } from './retry-types.js';
+
+/** Guard checks for premature completion */
+function isPrematureCompletionCandidate(
+  result: HeadlessRunResult,
+  state: RetryLoopState,
+  maxRetries: number,
+): boolean {
+  if (!result.completed || result.signalName || state.retryNumber >= maxRetries) return false;
+  if (state.checkpointRef.value || state.contextLost) return false;
+  if (!result.claudeSessionId || !result.stopReason) return false;
+  return result.stopReason === 'max_tokens' || result.stopReason === 'end_turn';
+}
+
+/** Use Haiku to assess whether an end_turn response is genuinely complete */
+async function assessEndTurnCompletion(result: HeadlessRunResult, verbose: boolean): Promise<boolean> {
+  if (!result.assistantResponse) return false;
+
+  const successfulToolCalls = result.toolUseHistory?.filter(t => t.result !== undefined && !t.isError).length ?? 0;
+  const claudeCmd = process.env.CLAUDE_COMMAND || 'claude';
+  const verdict = await assessPrematureCompletion({
+    responseTail: extractFinalTextBlock(result.assistantResponse, 800),
+    successfulToolCalls,
+    hasThinking: !!result.thinkingOutput,
+    responseLength: result.assistantResponse.length,
+  }, claudeCmd, verbose);
+
+  if (verbose) {
+    hlog(`[PREMATURE-COMPLETION] Haiku verdict: ${verdict.isIncomplete ? 'INCOMPLETE' : 'COMPLETE'} — ${verdict.reason}`);
+  }
+  return verdict.isIncomplete;
+}
+
+/** Apply premature completion retry */
+function applyPrematureCompletionRetry(
+  result: HeadlessRunResult,
+  state: RetryLoopState,
+  session: RetrySessionState,
+  maxRetries: number,
+  stopReason: string,
+  isMaxTokens: boolean,
+  callbacks: RetryCallbacks,
+): void {
+  state.retryNumber++;
+  const reason = isMaxTokens ? 'Output limit reached' : 'Task appears unfinished (AI assessment)';
+
+  state.retryLog.push({
+    retryNumber: state.retryNumber,
+    path: 'PrematureCompletion',
+    reason,
+    timestamp: Date.now(),
+  });
+
+  callbacks.emit('onAutoRetry', {
+    retryNumber: state.retryNumber,
+    maxRetries,
+    toolName: `PrematureCompletion(${stopReason})`,
+    completedCount: result.toolUseHistory?.length ?? 0,
+  });
+
+  trackEvent(AnalyticsEvents.IMPROVISE_AUTO_RETRY, {
+    retry_number: state.retryNumber,
+    hung_tool: `premature_completion:${stopReason}`,
+    completed_tools: result.toolUseHistory?.length ?? 0,
+    resume_attempted: true,
+  });
+
+  callbacks.queueOutput(
+    `\n\n${reason} — resuming session (retry ${state.retryNumber}/${maxRetries}).\n\n`
+  );
+  callbacks.flushOutputQueue();
+
+  state.contextRecoverySessionId = result.claudeSessionId;
+  session.claudeSessionId = result.claudeSessionId;
+  state.currentPrompt = 'continue';
+}
+
+/** Detect and retry premature completion. Returns true if loop should continue. */
+export async function shouldRetryPrematureCompletion(
+  result: HeadlessRunResult,
+  state: RetryLoopState,
+  session: RetrySessionState,
+  maxRetries: number,
+  callbacks: RetryCallbacks,
+): Promise<boolean> {
+  if (!isPrematureCompletionCandidate(result, state, maxRetries)) {
+    return false;
+  }
+
+  const stopReason = result.stopReason!;
+  const isMaxTokens = stopReason === 'max_tokens';
+  const abandoned = isResponseAbandoned(result);
+  const isIncomplete = isMaxTokens || abandoned || await assessEndTurnCompletion(result, session.options.verbose);
+
+  if (!isIncomplete) return false;
+
+  applyPrematureCompletionRetry(result, state, session, maxRetries, stopReason, isMaxTokens, callbacks);
+  return true;
+}