mstro-app 0.4.2 → 0.4.4

package/server/cli/headless/claude-invoker-stream.ts

@@ -0,0 +1,353 @@
+// Copyright (c) 2025-present Mstro, Inc. All rights reserved.
+// Licensed under the MIT License. See LICENSE file for details.
+
+import { hlog } from './headless-logger.js';
+import type { NativeTimeoutDetector } from './native-timeout-detector.js';
+import { classifyError } from './stall-assessor.js';
+import type { ResolvedHeadlessConfig, ToolUseAccumulator } from './types.js';
+
+// biome-ignore lint/suspicious/noExplicitAny: external CLI stream JSON with heterogeneous shapes
+type StreamJson = any;
+
+export interface StreamHandlerContext {
+  config: ResolvedHeadlessConfig;
+  accumulatedAssistantResponse: string;
+  accumulatedThinking: string;
+  accumulatedToolUse: ToolUseAccumulator[];
+  toolInputBuffers: Map<number, { name: string; id: string; inputJson: string; startTime: number }>;
+  nativeTimeoutDetector: NativeTimeoutDetector;
+  /** When true, assistant text is buffered instead of forwarded to outputCallback.
+   *  Active during resume mode until thinking/tool activity confirms Claude has context. */
+  resumeAssessmentActive: boolean;
+  /** Buffered assistant text during resume assessment */
+  resumeAssessmentBuffer: string;
+  /** Cumulative API token usage from message_start/message_delta events */
+  apiTokenUsage: { inputTokens: number; outputTokens: number };
+  /** Tracks cumulative output_tokens within the current step (message_delta is cumulative per-step) */
+  currentStepOutputTokens: number;
+  /** Timestamp of the last token usage change (tokens still flowing = process alive) */
+  lastTokenActivityTime: number;
+  /** Claude Code result event stop_reason (e.g., 'end_turn', 'max_tokens') */
+  stopReason?: string;
+}
+
+/** Log messages when verbose mode is enabled */
+export function verboseLog(verbose: boolean | undefined, ...msgs: string[]): void {
+  if (verbose) {
+    for (const msg of msgs) hlog(msg);
+  }
+}
+
+// ========== Stream Event Handlers ==========
+
+function handleSessionCapture(
+  parsed: StreamJson,
+  captured: { claudeSessionId?: string }
+): void {
+  if (parsed.type === 'system' && parsed.subtype === 'init' && parsed.session_id) {
+    captured.claudeSessionId = parsed.session_id;
+  }
+  if (parsed.type === 'result' && parsed.session_id && !captured.claudeSessionId) {
+    captured.claudeSessionId = parsed.session_id;
+  }
+}
+
+function handleThinkingDelta(event: StreamJson, ctx: StreamHandlerContext): string {
+  if (
+    event.type !== 'content_block_delta' ||
+    event.delta?.type !== 'thinking_delta' ||
+    !event.delta?.thinking
+  ) {
+    return ctx.accumulatedThinking;
+  }
+
+  if (ctx.resumeAssessmentActive) {
+    ctx.resumeAssessmentActive = false;
+    if (ctx.resumeAssessmentBuffer) {
+      ctx.config.outputCallback?.(ctx.resumeAssessmentBuffer);
+      ctx.resumeAssessmentBuffer = '';
+    }
+  }
+
+  const thinking = event.delta.thinking;
+  const updated = ctx.accumulatedThinking + thinking;
+
+  if (ctx.config.thinkingCallback) {
+    ctx.config.thinkingCallback(thinking);
+  } else if (ctx.config.outputCallback) {
+    ctx.config.outputCallback(thinking);
+  } else {
+    process.stdout.write(thinking);
+  }
+
+  return updated;
+}
+
+function handleTextDelta(event: StreamJson, ctx: StreamHandlerContext): string {
+  if (
+    event.type !== 'content_block_delta' ||
+    event.delta?.type !== 'text_delta' ||
+    !event.delta?.text
+  ) {
+    return ctx.accumulatedAssistantResponse;
+  }
+
+  const text = event.delta.text;
+  const updated = ctx.accumulatedAssistantResponse + text;
+
+  const { passthrough, timeouts } = ctx.nativeTimeoutDetector.processChunk(text);
+
+  for (const timeout of timeouts) {
+    ctx.config.outputCallback?.(
+      `\n[[MSTRO_NATIVE_TIMEOUT]] ${timeout.toolName} timed out \u2014 ${timeout.action} with ${timeout.preservedCount} results preserved\n`
+    );
+  }
+
+  if (ctx.resumeAssessmentActive) {
+    if (passthrough) {
+      ctx.resumeAssessmentBuffer += passthrough;
+    }
+    return updated;
+  }
+
+  if (passthrough && ctx.config.outputCallback) {
+    ctx.config.outputCallback(passthrough);
+  }
+
+  return updated;
+}
+
+/** Accumulate input tokens from a message_start event */
+function handleMessageStartTokens(event: StreamJson, ctx: StreamHandlerContext): boolean {
+  if (event.type !== 'message_start' || !event.message?.usage) return false;
+  const usage = event.message.usage;
+  ctx.currentStepOutputTokens = 0;
+  let changed = false;
+  if (typeof usage.input_tokens === 'number') {
+    ctx.apiTokenUsage.inputTokens += usage.input_tokens;
+    changed = true;
+  }
+  if (typeof usage.cache_creation_input_tokens === 'number') {
+    ctx.apiTokenUsage.inputTokens += usage.cache_creation_input_tokens;
+    changed = true;
+  }
+  if (typeof usage.cache_read_input_tokens === 'number') {
+    ctx.apiTokenUsage.inputTokens += usage.cache_read_input_tokens;
+    changed = true;
+  }
+  verboseLog(ctx.config.verbose,
+    `[TOKENS] message_start: input=${usage.input_tokens ?? 0} cache_create=${usage.cache_creation_input_tokens ?? 0} cache_read=${usage.cache_read_input_tokens ?? 0} → total_input=${ctx.apiTokenUsage.inputTokens}`);
+  return changed;
+}
+
+/** Accumulate output tokens from a message_delta event (cumulative tracking) */
+function handleMessageDeltaTokens(event: StreamJson, ctx: StreamHandlerContext): boolean {
+  if (event.type !== 'message_delta' || !event.usage) return false;
+  if (typeof event.usage.output_tokens !== 'number') return false;
+  const increment = event.usage.output_tokens - ctx.currentStepOutputTokens;
+  verboseLog(ctx.config.verbose,
+    `[TOKENS] message_delta: output=${event.usage.output_tokens} (step_prev=${ctx.currentStepOutputTokens} increment=${increment}) → total_output=${ctx.apiTokenUsage.outputTokens + Math.max(increment, 0)}`);
+  if (increment <= 0) return false;
+  ctx.apiTokenUsage.outputTokens += increment;
+  ctx.currentStepOutputTokens = event.usage.output_tokens;
+  return true;
+}
+
+function handleTokenUsage(event: StreamJson, ctx: StreamHandlerContext): void {
+  const changed = handleMessageStartTokens(event, ctx) || handleMessageDeltaTokens(event, ctx);
+  if (changed) {
+    ctx.lastTokenActivityTime = Date.now();
+    ctx.config.tokenUsageCallback?.({ ...ctx.apiTokenUsage });
+  }
+}
+
+/** Extract definitive token usage from the result event */
+function handleResultTokenUsage(parsed: StreamJson, ctx: StreamHandlerContext): void {
+  if (!parsed.usage) return;
+  const u = parsed.usage;
+  const input = (typeof u.input_tokens === 'number' ? u.input_tokens : 0)
+    + (typeof u.cache_creation_input_tokens === 'number' ? u.cache_creation_input_tokens : 0)
+    + (typeof u.cache_read_input_tokens === 'number' ? u.cache_read_input_tokens : 0);
+  const output = typeof u.output_tokens === 'number' ? u.output_tokens : 0;
+
+  if (input > 0 || output > 0) {
+    verboseLog(ctx.config.verbose,
+      `[TOKENS] Result event usage: input=${input} output=${output} ` +
+      `(stream accumulated: input=${ctx.apiTokenUsage.inputTokens} output=${ctx.apiTokenUsage.outputTokens})`);
+    ctx.apiTokenUsage = { inputTokens: input, outputTokens: output };
+    ctx.lastTokenActivityTime = Date.now();
+    ctx.config.tokenUsageCallback?.({ ...ctx.apiTokenUsage });
+  }
+}
+
+function handleToolResult(parsed: StreamJson, ctx: StreamHandlerContext): void {
+  if (parsed.type !== 'user' || !parsed.message?.content) {
+    return;
+  }
+
+  for (const content of parsed.message.content) {
+    if (content.type !== 'tool_result') {
+      continue;
+    }
+
+    const toolId = content.tool_use_id;
+    const result = content.content;
+    const isError = content.is_error || false;
+    const resultStr = typeof result === 'string' ? result : JSON.stringify(result);
+
+    const toolEntry = ctx.accumulatedToolUse.find(t => t.toolId === toolId);
+    if (toolEntry) {
+      toolEntry.result = resultStr;
+      toolEntry.isError = isError;
+      toolEntry.duration = Date.now() - toolEntry.startTime;
+    }
+
+    if (ctx.config.toolUseCallback) {
+      ctx.config.toolUseCallback({ type: 'tool_result', toolId, result: resultStr, isError });
+    }
+  }
+}
+
+// ========== Stream Processing ==========
+
+export function processStreamEvent(parsed: StreamJson, ctx: StreamHandlerContext): void {
+  if (parsed.type === 'error') {
+    const errorMessage = parsed.error?.message || parsed.message || JSON.stringify(parsed);
+    ctx.config.outputCallback?.(`\n[[MSTRO_ERROR:CLAUDE_ERROR]] ${errorMessage}\n`);
+    return;
+  }
+
+  if (parsed.type === 'result') {
+    handleResultTokenUsage(parsed, ctx);
+    if (parsed.stop_reason) {
+      ctx.stopReason = parsed.stop_reason;
+    }
+    if (parsed.is_error) {
+      const errorMessage = parsed.error || parsed.result || 'Unknown error in result';
+      ctx.config.outputCallback?.(`\n[[MSTRO_ERROR:CLAUDE_RESULT_ERROR]] ${errorMessage}\n`);
+      return;
+    }
+  }
+
+  if (parsed.type === 'stream_event' && parsed.event) {
+    const event = parsed.event;
+    ctx.accumulatedThinking = handleThinkingDelta(event, ctx);
+    ctx.accumulatedAssistantResponse = handleTextDelta(event, ctx);
+    handleToolStreamEvents(event, ctx);
+    handleTokenUsage(event, ctx);
+  }
+  handleToolResult(parsed, ctx);
+}
+
+/** Handle tool_use content_block_start */
+function handleToolStart(event: StreamJson, ctx: StreamHandlerContext): void {
+  if (event.type !== 'content_block_start' || event.content_block?.type !== 'tool_use') return;
+
+  if (ctx.resumeAssessmentActive) {
+    ctx.resumeAssessmentActive = false;
+    if (ctx.resumeAssessmentBuffer) {
+      ctx.config.outputCallback?.(ctx.resumeAssessmentBuffer);
+      ctx.resumeAssessmentBuffer = '';
+    }
+  }
+
+  const toolName = event.content_block.name;
+  const toolId = event.content_block.id;
+  const index = event.index;
+
+  ctx.toolInputBuffers.set(index, { name: toolName, id: toolId, inputJson: '', startTime: Date.now() });
+  ctx.config.toolUseCallback?.({ type: 'tool_start', toolName, toolId, index });
+}
+
+/** Handle input_json_delta for tool input streaming */
+function handleToolInputDelta(event: StreamJson, ctx: StreamHandlerContext): void {
+  if (event.type !== 'content_block_delta' || event.delta?.type !== 'input_json_delta') return;
+
+  const index = event.index;
+  const partialJson = event.delta.partial_json;
+  const toolBuffer = ctx.toolInputBuffers.get(index);
+  if (toolBuffer) toolBuffer.inputJson += partialJson;
+  ctx.config.toolUseCallback?.({ type: 'tool_input_delta', partialJson, index });
+}
+
+/** Handle content_block_stop — finalize tool input and emit tool_complete */
+function handleToolComplete(event: StreamJson, ctx: StreamHandlerContext): void {
+  if (event.type !== 'content_block_stop') return;
+
+  const index = event.index;
+  const toolBuffer = ctx.toolInputBuffers.get(index);
+  if (!toolBuffer) return;
+
+  let completeInput: Record<string, unknown> = {};
+  try { completeInput = JSON.parse(toolBuffer.inputJson); } catch { /* incomplete JSON */ }
+
+  ctx.accumulatedToolUse.push({
+    toolName: toolBuffer.name, toolId: toolBuffer.id,
+    toolInput: completeInput, startTime: toolBuffer.startTime
+  });
+  ctx.toolInputBuffers.delete(index);
+
+  ctx.config.toolUseCallback?.({
+    type: 'tool_complete', toolName: toolBuffer.name, toolId: toolBuffer.id,
+    index, completeInput
+  });
+}
+
+/** Handle tool-specific stream events (start, input delta, complete) */
+function handleToolStreamEvents(event: StreamJson, ctx: StreamHandlerContext): void {
+  handleToolStart(event, ctx);
+  handleToolInputDelta(event, ctx);
+  handleToolComplete(event, ctx);
+}
+
+export function processStreamLines(
+  buffer: string,
+  sessionCapture: { claudeSessionId?: string },
+  ctx: StreamHandlerContext
+): string {
+  const lines = buffer.split('\n');
+  const remainder = lines.pop() || '';
+
+  for (const line of lines) {
+    if (!line.trim()) continue;
+    try {
+      const parsed = JSON.parse(line);
+      handleSessionCapture(parsed, sessionCapture);
+      processStreamEvent(parsed, ctx);
+    } catch { /* Ignore parse errors */ }
+  }
+
+  return remainder;
+}
+
+/** Flush native timeout detector buffers and return post-timeout output if any */
+export function flushNativeTimeoutBuffers(ctx: StreamHandlerContext): string | undefined {
+  const remaining = ctx.nativeTimeoutDetector.flush();
+  const buffered = ctx.nativeTimeoutDetector.bufferedPostTimeoutOutput;
+  const postTimeout = (buffered + remaining) || undefined;
+
+  if (!postTimeout && remaining) {
+    ctx.config.outputCallback?.(remaining);
+  }
+
+  return postTimeout;
+}
+
+/** Classify unmatched stderr via Haiku when process exits with error */
+export async function classifyUnmatchedStderr(
+  stderr: string,
+  errorAlreadySurfaced: boolean,
+  code: number | null,
+  config: ResolvedHeadlessConfig,
+): Promise<void> {
+  if (!stderr || errorAlreadySurfaced || code === 0) return;
+
+  try {
+    const classified = await classifyError(stderr, config.claudeCommand, config.verbose);
+    if (classified) {
+      config.outputCallback?.(`\n[[MSTRO_ERROR:${classified.errorCode}]] ${classified.message}\n`);
+    }
+  } catch {
+    // Haiku classification failed — proceed without it
+  }
+}

package/server/cli/headless/claude-invoker-tools.ts

@@ -0,0 +1,187 @@
+// 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 { StallState } from './claude-invoker-stall.js';
+import type { StreamHandlerContext } from './claude-invoker-stream.js';
+import { verboseLog } from './claude-invoker-stream.js';
+import { killProcessGroup } from './runner.js';
+import { assessToolTimeout } from './stall-assessor.js';
+import { ToolWatchdog } from './tool-watchdog.js';
+import type { ResolvedHeadlessConfig, ToolUseEvent } from './types.js';
+
+// ========== Tool Tracking ==========
+
+/** Summarize a tool's input for stall assessment context */
+function summarizeToolInput(input: Record<string, unknown>): string | undefined {
+  try {
+    if (input.description) return String(input.description).slice(0, 200);
+    if (input.prompt) return String(input.prompt).slice(0, 200);
+    if (input.command) return String(input.command).slice(0, 200);
+    if (input.pattern) return `pattern: ${String(input.pattern).slice(0, 100)}`;
+    return JSON.stringify(input).slice(0, 200);
+  } catch {
+    return undefined;
+  }
+}
+
+/** Shared mutable state for tool event handlers */
+interface ToolTrackingState {
+  pendingTools: Map<string, string>;
+  counters: { lastToolInputSummary: string | undefined; totalToolCalls: number };
+  toolIdToName: Map<string, string>;
+  toolIdToInput: Map<string, Record<string, unknown>>;
+  watchdog: ToolWatchdog | null;
+  stallState: StallState;
+  ctx: StreamHandlerContext;
+  onTimeout: (hungToolId: string) => void;
+}
+
+export interface ToolTrackingResult {
+  pendingTools: Map<string, string>;
+  watchdog: ToolWatchdog | null;
+  toolWatchdogActive: boolean;
+  counters: { lastToolInputSummary: string | undefined; totalToolCalls: number };
+  /** Must be called after stallCheckInterval is created, to wire up the kill handler */
+  setKillContext: (claudeProcess: ChildProcess, stallCheckInterval: ReturnType<typeof setInterval>) => void;
+}
+
+function onToolStart(event: ToolUseEvent, s: ToolTrackingState): void {
+  const id = event.toolId!;
+  s.pendingTools.set(id, event.toolName!);
+  s.counters.totalToolCalls++;
+  s.toolIdToName.set(id, event.toolName!);
+  if (s.watchdog) {
+    s.watchdog.startWatch(id, event.toolName!, {}, () => { s.onTimeout(id); });
+  }
+}
+
+function onToolComplete(event: ToolUseEvent, s: ToolTrackingState): void {
+  const id = event.toolId!;
+  const input = event.completeInput ?? {};
+  s.counters.lastToolInputSummary = summarizeToolInput(input);
+  s.toolIdToInput.set(id, input);
+  if (!s.watchdog) return;
+  const toolName = s.toolIdToName.get(id);
+  if (toolName) {
+    s.watchdog.startWatch(id, toolName, input, () => { s.onTimeout(id); });
+  }
+}
+
+function onToolResult(event: ToolUseEvent, s: ToolTrackingState): void {
+  const id = event.toolId!;
+  s.pendingTools.delete(id);
+  s.stallState.stallWarningEmitted = false;
+  s.stallState.lastActivityTime = Date.now();
+  const toolEntry = s.ctx.accumulatedToolUse.find(t => t.toolId === id);
+  if (!s.watchdog || !toolEntry) return;
+  const toolName = s.toolIdToName.get(id);
+  if (toolName && toolEntry.duration) {
+    s.watchdog.recordCompletion(toolName, toolEntry.duration);
+  }
+  s.watchdog.clearWatch(id);
+}
+
+/** Resolve a display URL from tool input for timeout messages */
+function resolveToolUrl(toolInput: Record<string, unknown>): string | undefined {
+  if (toolInput.url) return String(toolInput.url);
+  if (toolInput.query) return String(toolInput.query);
+  return undefined;
+}
+
+/** Handle a tool timeout by building a checkpoint and killing the process */
+function executeToolTimeout(
+  hungToolId: string,
+  watchdog: ToolWatchdog,
+  killCtx: { claudeProcess: ChildProcess; stallCheckInterval: ReturnType<typeof setInterval> },
+  s: ToolTrackingState,
+  config: ResolvedHeadlessConfig,
+  prompt: string,
+  sessionCapture: { claudeSessionId?: string },
+  perfStart: number,
+): void {
+  const checkpoint = watchdog.buildCheckpoint(
+    prompt, s.ctx.accumulatedAssistantResponse, s.ctx.accumulatedThinking,
+    s.ctx.accumulatedToolUse, hungToolId, sessionCapture.claudeSessionId, perfStart,
+  );
+
+  const toolName = s.toolIdToName.get(hungToolId) || 'unknown';
+  const toolInput = s.toolIdToInput.get(hungToolId) || {};
+  const timeoutMs = watchdog.getTimeout(toolName);
+  const url = resolveToolUrl(toolInput);
+
+  config.outputCallback?.(
+    `\n[[MSTRO_TOOL_TIMEOUT]] ${toolName} timed out after ${Math.round(timeoutMs / 1000)}s${url ? ` fetching: ${url.slice(0, 100)}` : ''}. ${s.ctx.accumulatedToolUse.filter(t => t.result !== undefined).length} completed results preserved.\n`
+  );
+
+  if (checkpoint) {
+    config.onToolTimeout?.(checkpoint);
+  }
+
+  verboseLog(config.verbose, `[WATCHDOG] Killing process due to ${toolName} timeout`);
+  watchdog.clearAll();
+  clearInterval(killCtx.stallCheckInterval);
+  if (killCtx.claudeProcess.pid) killProcessGroup(killCtx.claudeProcess.pid, 'SIGTERM');
+  const proc = killCtx.claudeProcess;
+  setTimeout(() => { if (!proc.killed && proc.pid) killProcessGroup(proc.pid, 'SIGKILL'); }, 5000);
+}
+
+/** Set up tool activity tracking and watchdog */
+export function setupToolTracking(
+  config: ResolvedHeadlessConfig,
+  stallState: StallState,
+  ctx: StreamHandlerContext,
+  sessionCapture: { claudeSessionId?: string },
+  prompt: string,
+  perfStart: number,
+): ToolTrackingResult {
+  const pendingTools = new Map<string, string>();
+  const counters = { lastToolInputSummary: undefined as string | undefined, totalToolCalls: 0 };
+
+  const toolWatchdogActive = config.enableToolWatchdog !== false;
+  const watchdog = toolWatchdogActive
+    ? new ToolWatchdog({
+        profiles: config.toolTimeoutProfiles,
+        verbose: config.verbose,
+        onTiebreaker: async (toolName, toolInput, elapsedMs, tokenSilenceMs) => {
+          return assessToolTimeout(toolName, toolInput, elapsedMs, config.claudeCommand, config.verbose, tokenSilenceMs);
+        },
+        getTokenSilenceMs: () => {
+          const last = ctx.lastTokenActivityTime;
+          return last > 0 ? Date.now() - last : undefined;
+        },
+      })
+    : null;
+
+  let killCtx: { claudeProcess: ChildProcess; stallCheckInterval: ReturnType<typeof setInterval> } | null = null;
+
+  const trackingState: ToolTrackingState = {
+    pendingTools, counters,
+    toolIdToName: new Map(), toolIdToInput: new Map(),
+    watchdog, stallState, ctx,
+    onTimeout: (hungToolId) => {
+      if (!watchdog || !killCtx) return;
+      executeToolTimeout(hungToolId, watchdog, killCtx, trackingState, config, prompt, sessionCapture, perfStart);
+    },
+  };
+
+  const origToolUseCallback = config.toolUseCallback;
+
+  config.toolUseCallback = (event) => {
+    if (event.type === 'tool_start' && event.toolName && event.toolId) {
+      onToolStart(event, trackingState);
+    } else if (event.type === 'tool_complete' && event.completeInput && event.toolId) {
+      onToolComplete(event, trackingState);
+    } else if (event.type === 'tool_result' && event.toolId) {
+      onToolResult(event, trackingState);
+    }
+    origToolUseCallback?.(event);
+  };
+
+  return {
+    pendingTools, watchdog, toolWatchdogActive, counters,
+    setKillContext: (claudeProcess, stallCheckInterval) => {
+      killCtx = { claudeProcess, stallCheckInterval };
+    },
+  };
+}