claude-code

src/query/config.ts

@@ -0,0 +1,47 @@
+import { getSessionId } from '../bootstrap/state.js'
+import { checkStatsigFeatureGate_CACHED_MAY_BE_STALE } from '../services/analytics/growthbook.js'
+import type { SessionId } from '../types/ids.js'
+import { isEnvTruthy } from '../utils/envUtils.js'
+
+// -- config
+
+// Immutable values snapshotted once at query() entry. Separating these from
+// the per-iteration State struct and the mutable ToolUseContext makes future
+// step() extraction tractable — a pure reducer can take (state, event, config)
+// where config is plain data.
+//
+// Intentionally excludes feature() gates — those are tree-shaking boundaries
+// and must stay inline at the guarded blocks for dead-code elimination.
+export type QueryConfig = {
+  sessionId: SessionId
+
+  // Runtime gates (env/statsig). NOT feature() gates — see above.
+  gates: {
+    // Statsig — CACHED_MAY_BE_STALE already admits staleness, so snapshotting
+    // once per query() call stays within the existing contract.
+    streamingToolExecution: boolean
+    emitToolUseSummaries: boolean
+    isAnt: boolean
+    fastModeEnabled: boolean
+  }
+}
+
+export function buildQueryConfig(): QueryConfig {
+  return {
+    sessionId: getSessionId(),
+    gates: {
+      streamingToolExecution: checkStatsigFeatureGate_CACHED_MAY_BE_STALE(
+        'tengu_streaming_tool_execution2',
+      ),
+      emitToolUseSummaries: isEnvTruthy(
+        process.env.CLAUDE_CODE_EMIT_TOOL_USE_SUMMARIES,
+      ),
+      isAnt: process.env.USER_TYPE === 'ant',
+      // Inlined from fastMode.ts to avoid pulling its heavy module graph
+      // (axios, settings, auth, model, oauth, config) into test shards that
+      // didn't previously load it — changes init order and breaks unrelated tests.
+      fastModeEnabled: !isEnvTruthy(process.env.CLAUDE_CODE_DISABLE_FAST_MODE),
+    },
+  }
+}
+

src/query/deps.ts

@@ -0,0 +1,41 @@
+import { randomUUID } from 'crypto'
+import { queryModelWithStreaming } from '../services/api/claude.js'
+import { autoCompactIfNeeded } from '../services/compact/autoCompact.js'
+import { microcompactMessages } from '../services/compact/microCompact.js'
+
+// -- deps
+
+// I/O dependencies for query(). Passing a `deps` override into QueryParams
+// lets tests inject fakes directly instead of spyOn-per-module — the most
+// common mocks (callModel, autocompact) are each spied in 6-8 test files
+// today with module-import-and-spy boilerplate.
+//
+// Using `typeof fn` keeps signatures in sync with the real implementations
+// automatically. This file imports the real functions for both typing and
+// the production factory — tests that import this file for typing are
+// already importing query.ts (which imports everything), so there's no
+// new module-graph cost.
+//
+// Scope is intentionally narrow (4 deps) to prove the pattern. Followup
+// PRs can add runTools, handleStopHooks, logEvent, queue ops, etc.
+export type QueryDeps = {
+  // -- model
+  callModel: typeof queryModelWithStreaming
+
+  // -- compaction
+  microcompact: typeof microcompactMessages
+  autocompact: typeof autoCompactIfNeeded
+
+  // -- platform
+  uuid: () => string
+}
+
+export function productionDeps(): QueryDeps {
+  return {
+    callModel: queryModelWithStreaming,
+    microcompact: microcompactMessages,
+    autocompact: autoCompactIfNeeded,
+    uuid: randomUUID,
+  }
+}
+

src/query/stopHooks.ts

@@ -0,0 +1,474 @@
+import { feature } from 'bun:bundle'
+import { getShortcutDisplay } from '../keybindings/shortcutFormat.js'
+import { isExtractModeActive } from '../memdir/paths.js'
+import {
+  type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+  logEvent,
+} from '../services/analytics/index.js'
+import type { ToolUseContext } from '../Tool.js'
+import type { HookProgress } from '../types/hooks.js'
+import type {
+  AssistantMessage,
+  Message,
+  RequestStartEvent,
+  StopHookInfo,
+  StreamEvent,
+  TombstoneMessage,
+  ToolUseSummaryMessage,
+} from '../types/message.js'
+import { createAttachmentMessage } from '../utils/attachments.js'
+import { logForDebugging } from '../utils/debug.js'
+import { errorMessage } from '../utils/errors.js'
+import type { REPLHookContext } from '../utils/hooks/postSamplingHooks.js'
+import {
+  executeStopHooks,
+  executeTaskCompletedHooks,
+  executeTeammateIdleHooks,
+  getStopHookMessage,
+  getTaskCompletedHookMessage,
+  getTeammateIdleHookMessage,
+} from '../utils/hooks.js'
+import {
+  createStopHookSummaryMessage,
+  createSystemMessage,
+  createUserInterruptionMessage,
+  createUserMessage,
+} from '../utils/messages.js'
+import type { SystemPrompt } from '../utils/systemPromptType.js'
+import { getTaskListId, listTasks } from '../utils/tasks.js'
+import { getAgentName, getTeamName, isTeammate } from '../utils/teammate.js'
+
+/* eslint-disable @typescript-eslint/no-require-imports */
+const extractMemoriesModule = feature('EXTRACT_MEMORIES')
+  ? (require('../services/extractMemories/extractMemories.js') as typeof import('../services/extractMemories/extractMemories.js'))
+  : null
+const jobClassifierModule = feature('TEMPLATES')
+  ? (require('../jobs/classifier.js') as typeof import('../jobs/classifier.js'))
+  : null
+
+/* eslint-enable @typescript-eslint/no-require-imports */
+
+import type { QuerySource } from '../constants/querySource.js'
+import { executeAutoDream } from '../services/autoDream/autoDream.js'
+import { executePromptSuggestion } from '../services/PromptSuggestion/promptSuggestion.js'
+import { isBareMode, isEnvDefinedFalsy } from '../utils/envUtils.js'
+import {
+  createCacheSafeParams,
+  saveCacheSafeParams,
+} from '../utils/forkedAgent.js'
+
+type StopHookResult = {
+  blockingErrors: Message[]
+  preventContinuation: boolean
+}
+
+export async function* handleStopHooks(
+  messagesForQuery: Message[],
+  assistantMessages: AssistantMessage[],
+  systemPrompt: SystemPrompt,
+  userContext: { [k: string]: string },
+  systemContext: { [k: string]: string },
+  toolUseContext: ToolUseContext,
+  querySource: QuerySource,
+  stopHookActive?: boolean,
+): AsyncGenerator<
+  | StreamEvent
+  | RequestStartEvent
+  | Message
+  | TombstoneMessage
+  | ToolUseSummaryMessage,
+  StopHookResult
+> {
+  const hookStartTime = Date.now()
+
+  const stopHookContext: REPLHookContext = {
+    messages: [...messagesForQuery, ...assistantMessages],
+    systemPrompt,
+    userContext,
+    systemContext,
+    toolUseContext,
+    querySource,
+  }
+  // Only save params for main session queries — subagents must not overwrite.
+  // Outside the prompt-suggestion gate: the REPL /btw command and the
+  // side_question SDK control_request both read this snapshot, and neither
+  // depends on prompt suggestions being enabled.
+  if (querySource === 'repl_main_thread' || querySource === 'sdk') {
+    saveCacheSafeParams(createCacheSafeParams(stopHookContext))
+  }
+
+  // Template job classification: when running as a dispatched job, classify
+  // state after each turn. Gate on repl_main_thread so background forks
+  // (extract-memories, auto-dream) don't pollute the timeline with their own
+  // assistant messages. Await the classifier so state.json is written before
+  // the turn returns — otherwise `claude list` shows stale state for the gap.
+  // Env key hardcoded (vs importing JOB_ENV_KEY from jobs/state) to match the
+  // require()-gated jobs/ import pattern above; spawn.test.ts asserts the
+  // string matches.
+  if (
+    feature('TEMPLATES') &&
+    process.env.CLAUDE_JOB_DIR &&
+    querySource.startsWith('repl_main_thread') &&
+    !toolUseContext.agentId
+  ) {
+    // Full turn history — assistantMessages resets each queryLoop iteration,
+    // so tool calls from earlier iterations (Agent spawn, then summary) need
+    // messagesForQuery to be visible in the tool-call summary.
+    const turnAssistantMessages = stopHookContext.messages.filter(
+      (m): m is AssistantMessage => m.type === 'assistant',
+    )
+    const p = jobClassifierModule!
+      .classifyAndWriteState(process.env.CLAUDE_JOB_DIR, turnAssistantMessages)
+      .catch(err => {
+        logForDebugging(`[job] classifier error: ${errorMessage(err)}`, {
+          level: 'error',
+        })
+      })
+    await Promise.race([
+      p,
+      // eslint-disable-next-line no-restricted-syntax -- sleep() has no .unref(); timer must not block exit
+      new Promise<void>(r => setTimeout(r, 60_000).unref()),
+    ])
+  }
+  // --bare / SIMPLE: skip background bookkeeping (prompt suggestion,
+  // memory extraction, auto-dream). Scripted -p calls don't want auto-memory
+  // or forked agents contending for resources during shutdown.
+  if (!isBareMode()) {
+    // Inline env check for dead code elimination in external builds
+    if (!isEnvDefinedFalsy(process.env.CLAUDE_CODE_ENABLE_PROMPT_SUGGESTION)) {
+      void executePromptSuggestion(stopHookContext)
+    }
+    if (
+      feature('EXTRACT_MEMORIES') &&
+      !toolUseContext.agentId &&
+      isExtractModeActive()
+    ) {
+      // Fire-and-forget in both interactive and non-interactive. For -p/SDK,
+      // print.ts drains the in-flight promise after flushing the response
+      // but before gracefulShutdownSync (see drainPendingExtraction).
+      void extractMemoriesModule!.executeExtractMemories(
+        stopHookContext,
+        toolUseContext.appendSystemMessage,
+      )
+    }
+    if (!toolUseContext.agentId) {
+      void executeAutoDream(stopHookContext, toolUseContext.appendSystemMessage)
+    }
+  }
+
+  // chicago MCP: auto-unhide + lock release at turn end.
+  // Main thread only — the CU lock is a process-wide module-level variable,
+  // so a subagent's stopHooks releasing it leaves the main thread's cleanup
+  // seeing isLockHeldLocally()===false → no exit notification, and unhides
+  // mid-turn. Subagents don't start CU sessions so this is a pure skip.
+  if (feature('CHICAGO_MCP') && !toolUseContext.agentId) {
+    try {
+      const { cleanupComputerUseAfterTurn } = await import(
+        '../utils/computerUse/cleanup.js'
+      )
+      await cleanupComputerUseAfterTurn(toolUseContext)
+    } catch {
+      // Failures are silent — this is dogfooding cleanup, not critical path
+    }
+  }
+
+  try {
+    const blockingErrors = []
+    const appState = toolUseContext.getAppState()
+    const permissionMode = appState.toolPermissionContext.mode
+
+    const generator = executeStopHooks(
+      permissionMode,
+      toolUseContext.abortController.signal,
+      undefined,
+      stopHookActive ?? false,
+      toolUseContext.agentId,
+      toolUseContext,
+      [...messagesForQuery, ...assistantMessages],
+      toolUseContext.agentType,
+    )
+
+    // Consume all progress messages and get blocking errors
+    let stopHookToolUseID = ''
+    let hookCount = 0
+    let preventedContinuation = false
+    let stopReason = ''
+    let hasOutput = false
+    const hookErrors: string[] = []
+    const hookInfos: StopHookInfo[] = []
+
+    for await (const result of generator) {
+      if (result.message) {
+        yield result.message
+        // Track toolUseID from progress messages and count hooks
+        if (result.message.type === 'progress' && result.message.toolUseID) {
+          stopHookToolUseID = result.message.toolUseID
+          hookCount++
+          // Extract hook command and prompt text from progress data
+          const progressData = result.message.data as HookProgress
+          if (progressData.command) {
+            hookInfos.push({
+              command: progressData.command,
+              promptText: progressData.promptText,
+            })
+          }
+        }
+        // Track errors and output from attachments
+        if (result.message.type === 'attachment') {
+          const attachment = result.message.attachment
+          if (
+            'hookEvent' in attachment &&
+            (attachment.hookEvent === 'Stop' ||
+              attachment.hookEvent === 'SubagentStop')
+          ) {
+            if (attachment.type === 'hook_non_blocking_error') {
+              hookErrors.push(
+                attachment.stderr || `Exit code ${attachment.exitCode}`,
+              )
+              // Non-blocking errors always have output
+              hasOutput = true
+            } else if (attachment.type === 'hook_error_during_execution') {
+              hookErrors.push(attachment.content)
+              hasOutput = true
+            } else if (attachment.type === 'hook_success') {
+              // Check if successful hook produced any stdout/stderr
+              if (
+                (attachment.stdout && attachment.stdout.trim()) ||
+                (attachment.stderr && attachment.stderr.trim())
+              ) {
+                hasOutput = true
+              }
+            }
+            // Extract per-hook duration for timing visibility.
+            // Hooks run in parallel; match by command + first unassigned entry.
+            if ('durationMs' in attachment && 'command' in attachment) {
+              const info = hookInfos.find(
+                i =>
+                  i.command === attachment.command &&
+                  i.durationMs === undefined,
+              )
+              if (info) {
+                info.durationMs = attachment.durationMs
+              }
+            }
+          }
+        }
+      }
+      if (result.blockingError) {
+        const userMessage = createUserMessage({
+          content: getStopHookMessage(result.blockingError),
+          isMeta: true, // Hide from UI (shown in summary message instead)
+        })
+        blockingErrors.push(userMessage)
+        yield userMessage
+        hasOutput = true
+        // Add to hookErrors so it appears in the summary
+        hookErrors.push(result.blockingError.blockingError)
+      }
+      // Check if hook wants to prevent continuation
+      if (result.preventContinuation) {
+        preventedContinuation = true
+        stopReason = result.stopReason || 'Stop hook prevented continuation'
+        // Create attachment to track the stopped continuation (for structured data)
+        yield createAttachmentMessage({
+          type: 'hook_stopped_continuation',
+          message: stopReason,
+          hookName: 'Stop',
+          toolUseID: stopHookToolUseID,
+          hookEvent: 'Stop',
+        })
+      }
+
+      // Check if we were aborted during hook execution
+      if (toolUseContext.abortController.signal.aborted) {
+        logEvent('tengu_pre_stop_hooks_cancelled', {
+          queryChainId: toolUseContext.queryTracking
+            ?.chainId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+
+          queryDepth: toolUseContext.queryTracking?.depth,
+        })
+        yield createUserInterruptionMessage({
+          toolUse: false,
+        })
+        return { blockingErrors: [], preventContinuation: true }
+      }
+    }
+
+    // Create summary system message if hooks ran
+    if (hookCount > 0) {
+      yield createStopHookSummaryMessage(
+        hookCount,
+        hookInfos,
+        hookErrors,
+        preventedContinuation,
+        stopReason,
+        hasOutput,
+        'suggestion',
+        stopHookToolUseID,
+      )
+
+      // Send notification about errors (shown in verbose/transcript mode via ctrl+o)
+      if (hookErrors.length > 0) {
+        const expandShortcut = getShortcutDisplay(
+          'app:toggleTranscript',
+          'Global',
+          'ctrl+o',
+        )
+        toolUseContext.addNotification?.({
+          key: 'stop-hook-error',
+          text: `Stop hook error occurred \u00b7 ${expandShortcut} to see`,
+          priority: 'immediate',
+        })
+      }
+    }
+
+    if (preventedContinuation) {
+      return { blockingErrors: [], preventContinuation: true }
+    }
+
+    // Collect blocking errors from stop hooks
+    if (blockingErrors.length > 0) {
+      return { blockingErrors, preventContinuation: false }
+    }
+
+    // After Stop hooks pass, run TeammateIdle and TaskCompleted hooks if this is a teammate
+    if (isTeammate()) {
+      const teammateName = getAgentName() ?? ''
+      const teamName = getTeamName() ?? ''
+      const teammateBlockingErrors: Message[] = []
+      let teammatePreventedContinuation = false
+      let teammateStopReason: string | undefined
+      // Each hook executor generates its own toolUseID — capture from progress
+      // messages (same pattern as stopHookToolUseID at L142), not the Stop ID.
+      let teammateHookToolUseID = ''
+
+      // Run TaskCompleted hooks for any in-progress tasks owned by this teammate
+      const taskListId = getTaskListId()
+      const tasks = await listTasks(taskListId)
+      const inProgressTasks = tasks.filter(
+        t => t.status === 'in_progress' && t.owner === teammateName,
+      )
+
+      for (const task of inProgressTasks) {
+        const taskCompletedGenerator = executeTaskCompletedHooks(
+          task.id,
+          task.subject,
+          task.description,
+          teammateName,
+          teamName,
+          permissionMode,
+          toolUseContext.abortController.signal,
+          undefined,
+          toolUseContext,
+        )
+
+        for await (const result of taskCompletedGenerator) {
+          if (result.message) {
+            if (
+              result.message.type === 'progress' &&
+              result.message.toolUseID
+            ) {
+              teammateHookToolUseID = result.message.toolUseID
+            }
+            yield result.message
+          }
+          if (result.blockingError) {
+            const userMessage = createUserMessage({
+              content: getTaskCompletedHookMessage(result.blockingError),
+              isMeta: true,
+            })
+            teammateBlockingErrors.push(userMessage)
+            yield userMessage
+          }
+          // Match Stop hook behavior: allow preventContinuation/stopReason
+          if (result.preventContinuation) {
+            teammatePreventedContinuation = true
+            teammateStopReason =
+              result.stopReason || 'TaskCompleted hook prevented continuation'
+            yield createAttachmentMessage({
+              type: 'hook_stopped_continuation',
+              message: teammateStopReason,
+              hookName: 'TaskCompleted',
+              toolUseID: teammateHookToolUseID,
+              hookEvent: 'TaskCompleted',
+            })
+          }
+          if (toolUseContext.abortController.signal.aborted) {
+            return { blockingErrors: [], preventContinuation: true }
+          }
+        }
+      }
+
+      // Run TeammateIdle hooks
+      const teammateIdleGenerator = executeTeammateIdleHooks(
+        teammateName,
+        teamName,
+        permissionMode,
+        toolUseContext.abortController.signal,
+      )
+
+      for await (const result of teammateIdleGenerator) {
+        if (result.message) {
+          if (result.message.type === 'progress' && result.message.toolUseID) {
+            teammateHookToolUseID = result.message.toolUseID
+          }
+          yield result.message
+        }
+        if (result.blockingError) {
+          const userMessage = createUserMessage({
+            content: getTeammateIdleHookMessage(result.blockingError),
+            isMeta: true,
+          })
+          teammateBlockingErrors.push(userMessage)
+          yield userMessage
+        }
+        // Match Stop hook behavior: allow preventContinuation/stopReason
+        if (result.preventContinuation) {
+          teammatePreventedContinuation = true
+          teammateStopReason =
+            result.stopReason || 'TeammateIdle hook prevented continuation'
+          yield createAttachmentMessage({
+            type: 'hook_stopped_continuation',
+            message: teammateStopReason,
+            hookName: 'TeammateIdle',
+            toolUseID: teammateHookToolUseID,
+            hookEvent: 'TeammateIdle',
+          })
+        }
+        if (toolUseContext.abortController.signal.aborted) {
+          return { blockingErrors: [], preventContinuation: true }
+        }
+      }
+
+      if (teammatePreventedContinuation) {
+        return { blockingErrors: [], preventContinuation: true }
+      }
+
+      if (teammateBlockingErrors.length > 0) {
+        return {
+          blockingErrors: teammateBlockingErrors,
+          preventContinuation: false,
+        }
+      }
+    }
+
+    return { blockingErrors: [], preventContinuation: false }
+  } catch (error) {
+    const durationMs = Date.now() - hookStartTime
+    logEvent('tengu_stop_hook_error', {
+      duration: durationMs,
+
+      queryChainId: toolUseContext.queryTracking
+        ?.chainId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
+      queryDepth: toolUseContext.queryTracking?.depth,
+    })
+    // Yield a system message that is not visible to the model for the user
+    // to debug their hook.
+    yield createSystemMessage(
+      `Stop hook failed: ${errorMessage(error)}`,
+      'warning',
+    )
+    return { blockingErrors: [], preventContinuation: false }
+  }
+}
+

src/query/tokenBudget.ts

@@ -0,0 +1,94 @@
+import { getBudgetContinuationMessage } from '../utils/tokenBudget.js'
+
+const COMPLETION_THRESHOLD = 0.9
+const DIMINISHING_THRESHOLD = 500
+
+export type BudgetTracker = {
+  continuationCount: number
+  lastDeltaTokens: number
+  lastGlobalTurnTokens: number
+  startedAt: number
+}
+
+export function createBudgetTracker(): BudgetTracker {
+  return {
+    continuationCount: 0,
+    lastDeltaTokens: 0,
+    lastGlobalTurnTokens: 0,
+    startedAt: Date.now(),
+  }
+}
+
+type ContinueDecision = {
+  action: 'continue'
+  nudgeMessage: string
+  continuationCount: number
+  pct: number
+  turnTokens: number
+  budget: number
+}
+
+type StopDecision = {
+  action: 'stop'
+  completionEvent: {
+    continuationCount: number
+    pct: number
+    turnTokens: number
+    budget: number
+    diminishingReturns: boolean
+    durationMs: number
+  } | null
+}
+
+export type TokenBudgetDecision = ContinueDecision | StopDecision
+
+export function checkTokenBudget(
+  tracker: BudgetTracker,
+  agentId: string | undefined,
+  budget: number | null,
+  globalTurnTokens: number,
+): TokenBudgetDecision {
+  if (agentId || budget === null || budget <= 0) {
+    return { action: 'stop', completionEvent: null }
+  }
+
+  const turnTokens = globalTurnTokens
+  const pct = Math.round((turnTokens / budget) * 100)
+  const deltaSinceLastCheck = globalTurnTokens - tracker.lastGlobalTurnTokens
+
+  const isDiminishing =
+    tracker.continuationCount >= 3 &&
+    deltaSinceLastCheck < DIMINISHING_THRESHOLD &&
+    tracker.lastDeltaTokens < DIMINISHING_THRESHOLD
+
+  if (!isDiminishing && turnTokens < budget * COMPLETION_THRESHOLD) {
+    tracker.continuationCount++
+    tracker.lastDeltaTokens = deltaSinceLastCheck
+    tracker.lastGlobalTurnTokens = globalTurnTokens
+    return {
+      action: 'continue',
+      nudgeMessage: getBudgetContinuationMessage(pct, turnTokens, budget),
+      continuationCount: tracker.continuationCount,
+      pct,
+      turnTokens,
+      budget,
+    }
+  }
+
+  if (isDiminishing || tracker.continuationCount > 0) {
+    return {
+      action: 'stop',
+      completionEvent: {
+        continuationCount: tracker.continuationCount,
+        pct,
+        turnTokens,
+        budget,
+        diminishingReturns: isDiminishing,
+        durationMs: Date.now() - tracker.startedAt,
+      },
+    }
+  }
+
+  return { action: 'stop', completionEvent: null }
+}
+

src/query/transitions.ts

@@ -0,0 +1,37 @@
+/**
+ * Transition types for the query loop.
+ *
+ * Terminal: why the loop exited (returned).
+ * Continue: why the loop continued to the next iteration (not returned).
+ */
+
+/** Terminal transition — the query loop returned. */
+export type Terminal = {
+  reason:
+    | 'completed'
+    | 'blocking_limit'
+    | 'image_error'
+    | 'model_error'
+    | 'aborted_streaming'
+    | 'aborted_tools'
+    | 'prompt_too_long'
+    | 'stop_hook_prevented'
+    | 'hook_stopped'
+    | 'max_turns'
+    | (string & {})
+  error?: unknown
+}
+
+/** Continue transition — the loop will iterate again. */
+export type Continue = {
+  reason:
+    | 'tool_use'
+    | 'reactive_compact_retry'
+    | 'max_output_tokens_recovery'
+    | 'max_output_tokens_escalate'
+    | 'collapse_drain_retry'
+    | 'stop_hook_blocking'
+    | 'token_budget_continuation'
+    | 'queued_command'
+    | (string & {})
+}