mstro-app 0.1.58 → 0.3.0

package/server/cli/headless/tool-watchdog.ts

@@ -0,0 +1,398 @@
+// Copyright (c) 2025-present Mstro, Inc. All rights reserved.
+// Licensed under the MIT License. See LICENSE file for details.
+
+/**
+ * Tool Watchdog
+ *
+ * Per-tool adaptive timeout system using TCP RTO-style EMA tracking (RFC 6298).
+ * Monitors individual tool call durations and kills tools that exceed their
+ * adaptive timeout, preserving work via checkpoint-and-retry.
+ *
+ * Three-tier timeout strategy:
+ * 1. EMA tracking: timeout = estimatedDuration + 4 * deviation
+ * 2. Floor/ceiling bounds: never kill below floor, always kill at ceiling
+ * 3. Haiku tiebreaker: optional AI assessment before killing ambiguous cases
+ */
+
+import type {
+  ExecutionCheckpoint,
+  ToolDurationTracker,
+  ToolTimeoutProfile,
+  ToolUseAccumulator,
+} from './types.js';
+
+// RFC 6298 smoothing constants
+const ALPHA = 0.125;  // smoothing factor for duration EMA
+const BETA = 0.25;    // smoothing factor for deviation EMA
+const DEVIATION_MULTIPLIER = 4;  // timeout = est + 4*dev (same as TCP)
+
+/** Default timeout profiles per tool type */
+export const DEFAULT_TOOL_TIMEOUT_PROFILES: Record<string, ToolTimeoutProfile> = {
+  WebFetch: {
+    coldStartMs: 180_000,      // 3 min — accounts for slow sites + Haiku inference
+    floorMs: 120_000,          // 2 min absolute minimum
+    ceilingMs: 300_000,        // 5 min hard cap
+    useAdaptive: true,
+    useHaikuTiebreaker: true,
+  },
+  WebSearch: {
+    coldStartMs: 90_000,       // 1.5 min
+    floorMs: 60_000,           // 1 min minimum
+    ceilingMs: 180_000,        // 3 min hard cap
+    useAdaptive: true,
+    useHaikuTiebreaker: false,
+  },
+  Task: {
+    coldStartMs: 900_000,      // 15 min — subagents are inherently long-running
+    floorMs: 600_000,          // 10 min minimum (research agents routinely take 7-10 min)
+    ceilingMs: 2_700_000,      // 45 min hard cap
+    useAdaptive: true,         // learn from past Task durations via EMA
+    useHaikuTiebreaker: true,
+  },
+  Bash: {
+    coldStartMs: 300_000,      // 5 min
+    floorMs: 120_000,          // 2 min minimum
+    ceilingMs: 600_000,        // 10 min hard cap
+    useAdaptive: false,
+    useHaikuTiebreaker: true,
+  },
+  // Local filesystem tools — these go through Claude Code's streaming stdio protocol,
+  // NOT direct filesystem I/O. Large files/results can take 30-60s+ to stream.
+  // Read/Grep have bimodal distributions (tiny vs huge responses) that defeat EMA,
+  // so adaptive is disabled for them. Floors are generous to prevent premature kills.
+  Read: {
+    coldStartMs: 120_000,      // 2 min — large files stream slowly through stdio protocol
+    floorMs: 60_000,           // 1 min minimum — prevents EMA-driven premature kills
+    ceilingMs: 300_000,        // 5 min ceiling (very large files, slow mounts)
+    useAdaptive: false,        // bimodal: 1-line file vs 2000-line file defeats EMA
+    useHaikuTiebreaker: true,  // safety net: assess before killing the whole process
+  },
+  Grep: {
+    coldStartMs: 120_000,      // 2 min — broad searches return large result sets
+    floorMs: 60_000,           // 1 min minimum
+    ceilingMs: 300_000,        // 5 min ceiling
+    useAdaptive: false,        // bimodal: single-file vs codebase-wide search
+    useHaikuTiebreaker: true,  // safety net before killing
+  },
+  Glob: {
+    coldStartMs: 60_000,       // 1 min — pattern matching can be slow on large trees
+    floorMs: 30_000,           // 30s minimum
+    ceilingMs: 180_000,        // 3 min ceiling
+    useAdaptive: true,
+    useHaikuTiebreaker: true,
+  },
+  Edit: {
+    coldStartMs: 60_000,       // 1 min — edits go through streaming protocol too
+    floorMs: 30_000,           // 30s minimum
+    ceilingMs: 180_000,        // 3 min ceiling
+    useAdaptive: true,
+    useHaikuTiebreaker: true,
+  },
+  Write: {
+    coldStartMs: 60_000,       // 1 min
+    floorMs: 30_000,           // 30s minimum
+    ceilingMs: 180_000,        // 3 min ceiling
+    useAdaptive: true,
+    useHaikuTiebreaker: true,
+  },
+};
+
+const DEFAULT_TOOL_TIMEOUT_PROFILE: ToolTimeoutProfile = {
+  coldStartMs: 300_000,
+  floorMs: 120_000,
+  ceilingMs: 600_000,
+  useAdaptive: false,
+  useHaikuTiebreaker: true,
+};
+
+export interface ToolWatchdogOptions {
+  profiles?: Record<string, Partial<ToolTimeoutProfile>>;
+  verbose?: boolean;
+  /** Called before killing — if returns 'extend', reschedule with extensionMs */
+  onTiebreaker?: (toolName: string, toolInput: Record<string, unknown>, elapsedMs: number, tokenSilenceMs?: number) => Promise<{ action: 'extend' | 'kill'; extensionMs: number; reason: string }>;
+  /** Returns ms since last token activity. Called at tiebreaker time for fresh data. */
+  getTokenSilenceMs?: () => number | undefined;
+}
+
+interface ActiveWatch {
+  toolName: string;
+  toolInput: Record<string, unknown>;
+  startTime: number;
+  timer: ReturnType<typeof setTimeout>;
+  timeoutMs: number;
+  tiebreakerAttempted: boolean;
+}
+
+export class ToolWatchdog {
+  private trackers: Map<string, ToolDurationTracker> = new Map();
+  private profiles: Record<string, ToolTimeoutProfile>;
+  private activeWatches: Map<string, ActiveWatch> = new Map();
+  private verbose: boolean;
+  private onTiebreaker?: ToolWatchdogOptions['onTiebreaker'];
+  private getTokenSilenceMs?: () => number | undefined;
+
+  constructor(options: ToolWatchdogOptions = {}) {
+    this.verbose = options.verbose ?? false;
+    this.onTiebreaker = options.onTiebreaker;
+    this.getTokenSilenceMs = options.getTokenSilenceMs;
+
+    // Merge user profiles with defaults
+    this.profiles = { ...DEFAULT_TOOL_TIMEOUT_PROFILES };
+    if (options.profiles) {
+      for (const [name, partial] of Object.entries(options.profiles)) {
+        const base = this.profiles[name] || DEFAULT_TOOL_TIMEOUT_PROFILE;
+        this.profiles[name] = { ...base, ...partial };
+      }
+    }
+  }
+
+  /** Record a tool completion — updates the EMA tracker for its type */
+  recordCompletion(toolName: string, durationMs: number): void {
+    const profile = this.getProfile(toolName);
+    if (!profile.useAdaptive) return;
+
+    const tracker = this.trackers.get(toolName);
+    if (!tracker) {
+      this.trackers.set(toolName, {
+        estimatedDuration: durationMs,
+        deviation: durationMs / 2,
+        sampleCount: 1,
+      });
+      if (this.verbose) {
+        console.log(`[WATCHDOG] ${toolName}: first sample ${durationMs}ms, initial timeout ${this.getTimeout(toolName)}ms`);
+      }
+      return;
+    }
+
+    // RFC 6298 update
+    tracker.deviation = (1 - BETA) * tracker.deviation + BETA * Math.abs(durationMs - tracker.estimatedDuration);
+    tracker.estimatedDuration = (1 - ALPHA) * tracker.estimatedDuration + ALPHA * durationMs;
+    tracker.sampleCount++;
+
+    if (this.verbose) {
+      console.log(`[WATCHDOG] ${toolName}: sample #${tracker.sampleCount} ${durationMs}ms, est=${Math.round(tracker.estimatedDuration)}ms, dev=${Math.round(tracker.deviation)}ms, timeout=${this.getTimeout(toolName)}ms`);
+    }
+  }
+
+  /** Compute the current timeout for a tool type */
+  getTimeout(toolName: string): number {
+    const profile = this.getProfile(toolName);
+    if (!profile.useAdaptive) return profile.coldStartMs;
+
+    const tracker = this.trackers.get(toolName);
+    if (!tracker || tracker.sampleCount < 1) return profile.coldStartMs;
+
+    const adaptive = tracker.estimatedDuration + DEVIATION_MULTIPLIER * tracker.deviation;
+    return Math.max(profile.floorMs, Math.min(profile.ceilingMs, adaptive));
+  }
+
+  /** Get the profile for a tool (with fallback to default) */
+  getProfile(toolName: string): ToolTimeoutProfile {
+    return this.profiles[toolName] || DEFAULT_TOOL_TIMEOUT_PROFILE;
+  }
+
+  /** Start watching a tool call */
+  startWatch(toolId: string, toolName: string, toolInput: Record<string, unknown>, onTimeout: () => void): void {
+    // Clear any existing watch for this ID
+    this.clearWatch(toolId);
+
+    const timeoutMs = this.getTimeout(toolName);
+    const profile = this.getProfile(toolName);
+
+    if (this.verbose) {
+      console.log(`[WATCHDOG] Starting watch: ${toolName} (${toolId}), timeout=${Math.round(timeoutMs / 1000)}s`);
+    }
+
+    const timer = setTimeout(async () => {
+      const extended = await this.handleTimeoutWithTiebreaker(toolId, toolName, toolInput, profile, onTimeout);
+      if (!extended) {
+        // Don't delete the watch here — buildCheckpoint() needs it.
+        // handleToolTimeout() calls clearAll() after building the checkpoint.
+        onTimeout();
+      }
+    }, timeoutMs);
+
+    this.activeWatches.set(toolId, {
+      toolName,
+      toolInput,
+      startTime: Date.now(),
+      timer,
+      timeoutMs,
+      tiebreakerAttempted: false,
+    });
+  }
+
+  /** Handle timeout expiry: attempt tiebreaker if configured, return true if extended */
+  private async handleTimeoutWithTiebreaker(
+    toolId: string,
+    toolName: string,
+    toolInput: Record<string, unknown>,
+    profile: ToolTimeoutProfile,
+    onTimeout: () => void,
+  ): Promise<boolean> {
+    const watch = this.activeWatches.get(toolId);
+    if (!watch) return true;
+
+    const elapsedMs = Date.now() - watch.startTime;
+
+    if (!profile.useHaikuTiebreaker || !this.onTiebreaker || watch.tiebreakerAttempted) {
+      if (this.verbose) {
+        console.log(`[WATCHDOG] ${toolName} (${toolId}) timed out after ${Math.round(elapsedMs / 1000)}s, killing`);
+      }
+      return false;
+    }
+
+    return this.runTiebreaker(watch, toolId, toolName, toolInput, elapsedMs, onTimeout);
+  }
+
+  /** Execute the Haiku tiebreaker and reschedule if extended */
+  private async runTiebreaker(
+    watch: ActiveWatch,
+    toolId: string,
+    toolName: string,
+    toolInput: Record<string, unknown>,
+    elapsedMs: number,
+    onTimeout: () => void,
+  ): Promise<boolean> {
+    watch.tiebreakerAttempted = true;
+
+    if (this.verbose) {
+      console.log(`[WATCHDOG] ${toolName} (${toolId}) hit timeout after ${Math.round(elapsedMs / 1000)}s, running tiebreaker...`);
+    }
+
+    try {
+      const tokenSilenceMs = this.getTokenSilenceMs?.();
+      const verdict = await this.onTiebreaker!(toolName, toolInput, elapsedMs, tokenSilenceMs);
+
+      if (verdict.action === 'extend') {
+        if (this.verbose) {
+          console.log(`[WATCHDOG] Tiebreaker: extend ${toolName} by ${Math.round(verdict.extensionMs / 1000)}s — ${verdict.reason}`);
+        }
+        this.scheduleExtensionTimeout(watch, toolId, toolName, verdict.extensionMs, onTimeout);
+        watch.timeoutMs = elapsedMs + verdict.extensionMs;
+        return true;
+      }
+
+      if (this.verbose) {
+        console.log(`[WATCHDOG] Tiebreaker: kill ${toolName} — ${verdict.reason}`);
+      }
+    } catch (err) {
+      if (this.verbose) {
+        console.log(`[WATCHDOG] Tiebreaker failed: ${err}, proceeding with kill`);
+      }
+    }
+
+    return false;
+  }
+
+  /** Schedule a post-extension timeout that kills without another tiebreaker */
+  private scheduleExtensionTimeout(
+    watch: ActiveWatch,
+    toolId: string,
+    toolName: string,
+    extensionMs: number,
+    onTimeout: () => void,
+  ): void {
+    watch.timer = setTimeout(() => {
+      const w = this.activeWatches.get(toolId);
+      if (!w) return;
+      if (this.verbose) {
+        console.log(`[WATCHDOG] ${toolName} (${toolId}) still running after extension, killing`);
+      }
+      // Don't delete the watch — buildCheckpoint() needs it.
+      // handleToolTimeout() calls clearAll() after building the checkpoint.
+      onTimeout();
+    }, extensionMs);
+  }
+
+  /** Stop watching a tool (it completed normally) */
+  clearWatch(toolId: string): void {
+    const watch = this.activeWatches.get(toolId);
+    if (watch) {
+      clearTimeout(watch.timer);
+      this.activeWatches.delete(toolId);
+    }
+  }
+
+  /** Clear all active watches (process ending) */
+  clearAll(): void {
+    for (const [_id, watch] of this.activeWatches) {
+      clearTimeout(watch.timer);
+    }
+    this.activeWatches.clear();
+  }
+
+  /** Get the active watch for a tool ID (for checkpoint building) */
+  getActiveWatch(toolId: string): ActiveWatch | undefined {
+    return this.activeWatches.get(toolId);
+  }
+
+  /** Get all active watches */
+  getActiveWatches(): Map<string, ActiveWatch> {
+    return this.activeWatches;
+  }
+
+  /** Build an ExecutionCheckpoint from the current state */
+  buildCheckpoint(
+    originalPrompt: string,
+    assistantText: string,
+    thinkingText: string,
+    accumulatedToolUse: ToolUseAccumulator[],
+    hungToolId: string,
+    claudeSessionId: string | undefined,
+    processStartTime: number,
+  ): ExecutionCheckpoint | null {
+    const hungWatch = this.activeWatches.get(hungToolId);
+    if (!hungWatch) return null;
+
+    // Find the matching tool entry
+    const hungToolEntry = accumulatedToolUse.find(t => t.toolId === hungToolId);
+
+    // Build completed tools list (exclude the hung one)
+    const completedTools = accumulatedToolUse
+      .filter(t => t.toolId !== hungToolId && t.result !== undefined)
+      .map(t => ({
+        toolName: t.toolName,
+        toolId: t.toolId,
+        input: t.toolInput,
+        result: t.result || '',
+        isError: t.isError || false,
+        durationMs: t.duration || 0,
+      }));
+
+    // Build in-progress tools list (started but no result, excluding the hung one)
+    const inProgressTools = accumulatedToolUse
+      .filter(t => t.toolId !== hungToolId && t.result === undefined)
+      .map(t => ({
+        toolName: t.toolName,
+        toolId: t.toolId,
+        input: t.toolInput,
+      }));
+
+    // Extract URL from tool input if WebFetch/WebSearch
+    let url: string | undefined;
+    const toolInput = hungToolEntry?.toolInput || hungWatch.toolInput;
+    if (toolInput.url) {
+      url = String(toolInput.url);
+    } else if (toolInput.query) {
+      url = String(toolInput.query);
+    }
+
+    return {
+      originalPrompt,
+      assistantText,
+      thinkingText,
+      completedTools,
+      inProgressTools,
+      hungTool: {
+        toolName: hungWatch.toolName,
+        toolId: hungToolId,
+        input: toolInput,
+        timeoutMs: hungWatch.timeoutMs,
+        url,
+      },
+      claudeSessionId,
+      elapsedMs: Date.now() - processStartTime,
+    };
+  }
+}

package/server/cli/headless/types.ts

@@ -33,6 +33,57 @@ export interface ImageAttachment {
   mimeType?: string;      // MIME type (e.g., "image/png")
 }
 
+/** Per-tool-type timeout configuration with adaptive tracking */
+export interface ToolTimeoutProfile {
+  /** Initial timeout when no prior samples exist (ms) */
+  coldStartMs: number;
+  /** Minimum timeout — never kill before this (ms) */
+  floorMs: number;
+  /** Maximum timeout — always kill after this (ms) */
+  ceilingMs: number;
+  /** Track EMA of past durations and adapt timeout dynamically */
+  useAdaptive: boolean;
+  /** Spawn a Haiku call to assess before killing */
+  useHaikuTiebreaker: boolean;
+}
+
+/** Snapshot of execution state at the moment a tool times out */
+export interface ExecutionCheckpoint {
+  originalPrompt: string;
+  assistantText: string;
+  thinkingText: string;
+  completedTools: Array<{
+    toolName: string;
+    toolId: string;
+    input: Record<string, unknown>;
+    result: string;
+    isError: boolean;
+    durationMs: number;
+  }>;
+  /** Tools that were still running (not the hung tool) when the process was killed */
+  inProgressTools: Array<{
+    toolName: string;
+    toolId: string;
+    input: Record<string, unknown>;
+  }>;
+  hungTool: {
+    toolName: string;
+    toolId: string;
+    input: Record<string, unknown>;
+    timeoutMs: number;
+    url?: string;
+  };
+  claudeSessionId?: string;
+  elapsedMs: number;
+}
+
+/** EMA tracker for a single tool type's completion times */
+export interface ToolDurationTracker {
+  estimatedDuration: number;
+  deviation: number;
+  sampleCount: number;
+}
+
 export interface HeadlessConfig {
   workingDir: string;
   tokenBudgetThreshold: number;
@@ -46,6 +97,8 @@ export interface HeadlessConfig {
   outputCallback?: (text: string) => void;
   thinkingCallback?: (text: string) => void;
   toolUseCallback?: (event: ToolUseEvent) => void;
+  /** Called with cumulative API token counts as they arrive from the stream */
+  tokenUsageCallback?: (usage: { inputTokens: number; outputTokens: number }) => void;
   directPrompt?: string;
   promptContext?: PromptContext;
   continueSession?: boolean;
@@ -58,6 +111,16 @@ export interface HeadlessConfig {
   stallHardCapMs?: number;     // Absolute wall-clock kill cap (default: 3600000 = 60 min)
   /** Claude model for main execution (e.g., 'opus', 'sonnet'). 'default' = no --model flag. */
   model?: string;
+  /** Per-tool timeout profiles (merge with defaults) */
+  toolTimeoutProfiles?: Record<string, Partial<ToolTimeoutProfile>>;
+  /** Enable per-tool adaptive timeout watchdog (default: true) */
+  enableToolWatchdog?: boolean;
+  /** Max auto-retries on tool timeout (default: 2) */
+  maxAutoRetries?: number;
+  /** Called when a tool times out with checkpoint data */
+  onToolTimeout?: (checkpoint: ExecutionCheckpoint) => void;
+  /** When true, spawn Claude with sanitized env (strips secrets, HOME=workingDir) */
+  sandboxed?: boolean;
 }
 
 export interface SessionState {
@@ -76,6 +139,8 @@ export interface SessionResult {
   totalTokens: number;
   sessionId: string;
   error?: string;
+  /** Signal name if Claude process was killed (e.g., 'SIGTERM', 'SIGKILL') */
+  signalName?: string;
   conflicts?: Array<{
     filePath: string;
     modifiedBy: string[];
@@ -92,6 +157,14 @@ export interface SessionResult {
     duration?: number;
   }>;
   claudeSessionId?: string;
+  /** Number of Claude Code native tool timeouts detected during this execution */
+  nativeTimeoutCount?: number;
+  /** Assistant text buffered after native timeouts — not yet shown to user.
+   *  Flush to output if context is OK, discard if context was lost and recovery starts. */
+  postTimeoutOutput?: string;
+  /** Assistant text buffered during resume assessment — held back until thinking/tool activity
+   *  confirms Claude has context. Undefined when not in resume mode or buffer was flushed. */
+  resumeBufferedOutput?: string;
 }
 
 export interface ToolUseAccumulator {
@@ -104,23 +177,42 @@ export interface ToolUseAccumulator {
   duration?: number;
 }
 
+/** Map of toolId -> toolName for currently pending (started but not yet returned) tools */
+export type PendingToolMap = Map<string, string>;
+
 export interface ExecutionResult {
   output: string;
   error?: string;
   exitCode: number;
+  /** Signal name if process was killed (e.g., 'SIGTERM', 'SIGKILL') */
+  signalName?: string;
   assistantResponse?: string;
   thinkingOutput?: string;
   toolUseHistory?: ToolUseAccumulator[];
   claudeSessionId?: string;
+  /** Number of Claude Code native tool timeouts detected during this execution */
+  nativeTimeoutCount?: number;
+  /** Assistant text buffered after native timeouts — not yet sent to outputCallback.
+   *  The session manager should flush this to the client if context is OK, or discard if recovering. */
+  postTimeoutOutput?: string;
+  /** Assistant text buffered during resume assessment — held back until thinking/tool activity
+   *  confirms Claude has context. Undefined when not in resume mode or buffer was flushed. */
+  resumeBufferedOutput?: string;
+  /** Actual API token usage from Claude Code stream events (summed across all turns) */
+  apiTokenUsage?: { inputTokens: number; outputTokens: number };
 }
 
 /** Resolved config with all defaults applied */
-export type ResolvedHeadlessConfig = Omit<Required<HeadlessConfig>, 'outputCallback' | 'thinkingCallback' | 'toolUseCallback' | 'continueSession' | 'claudeSessionId' | 'imageAttachments' | 'model'> & {
+export type ResolvedHeadlessConfig = Omit<Required<HeadlessConfig>, 'outputCallback' | 'thinkingCallback' | 'toolUseCallback' | 'tokenUsageCallback' | 'continueSession' | 'claudeSessionId' | 'imageAttachments' | 'model' | 'toolTimeoutProfiles' | 'onToolTimeout' | 'sandboxed'> & {
   outputCallback?: (text: string) => void;
   thinkingCallback?: (text: string) => void;
   toolUseCallback?: (event: ToolUseEvent) => void;
+  tokenUsageCallback?: (usage: { inputTokens: number; outputTokens: number }) => void;
   continueSession?: boolean;
   claudeSessionId?: string;
   imageAttachments?: ImageAttachment[];
   model?: string;
+  toolTimeoutProfiles?: Record<string, Partial<ToolTimeoutProfile>>;
+  onToolTimeout?: (checkpoint: ExecutionCheckpoint) => void;
+  sandboxed?: boolean;
 };