1  import { feature } from 'bun:bundle'
  2  import { getShortcutDisplay } from '../keybindings/shortcutFormat.js'
  3  import { isExtractModeActive } from '../memdir/paths.js'
  4  import {
  5    type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
  6    logEvent,
  7  } from '../services/analytics/index.js'
  8  import type { ToolUseContext } from '../Tool.js'
  9  import type { HookProgress } from '../types/hooks.js'
 10  import type {
 11    AssistantMessage,
 12    Message,
 13    RequestStartEvent,
 14    StopHookInfo,
 15    StreamEvent,
 16    TombstoneMessage,
 17    ToolUseSummaryMessage,
 18  } from '../types/message.js'
 19  import { createAttachmentMessage } from '../utils/attachments.js'
 20  import { logForDebugging } from '../utils/debug.js'
 21  import { errorMessage } from '../utils/errors.js'
 22  import type { REPLHookContext } from '../utils/hooks/postSamplingHooks.js'
 23  import {
 24    executeStopHooks,
 25    executeTaskCompletedHooks,
 26    executeTeammateIdleHooks,
 27    getStopHookMessage,
 28    getTaskCompletedHookMessage,
 29    getTeammateIdleHookMessage,
 30  } from '../utils/hooks.js'
 31  import {
 32    createStopHookSummaryMessage,
 33    createSystemMessage,
 34    createUserInterruptionMessage,
 35    createUserMessage,
 36  } from '../utils/messages.js'
 37  import type { SystemPrompt } from '../utils/systemPromptType.js'
 38  import { getTaskListId, listTasks } from '../utils/tasks.js'
 39  import { getAgentName, getTeamName, isTeammate } from '../utils/teammate.js'
 40  
 41  /* eslint-disable @typescript-eslint/no-require-imports */
 42  const extractMemoriesModule = feature('EXTRACT_MEMORIES')
 43    ? (require('../services/extractMemories/extractMemories.js') as typeof import('../services/extractMemories/extractMemories.js'))
 44    : null
 45  const jobClassifierModule = feature('TEMPLATES')
 46    ? (require('../jobs/classifier.js') as typeof import('../jobs/classifier.js'))
 47    : null
 48  
 49  /* eslint-enable @typescript-eslint/no-require-imports */
 50  
 51  import type { QuerySource } from '../constants/querySource.js'
 52  import { executeAutoDream } from '../services/autoDream/autoDream.js'
 53  import { executePromptSuggestion } from '../services/PromptSuggestion/promptSuggestion.js'
 54  import { isBareMode, isEnvDefinedFalsy } from '../utils/envUtils.js'
 55  import {
 56    createCacheSafeParams,
 57    saveCacheSafeParams,
 58  } from '../utils/forkedAgent.js'
 59  
 60  type StopHookResult = {
 61    blockingErrors: Message[]
 62    preventContinuation: boolean
 63  }
 64  
 65  export async function* handleStopHooks(
 66    messagesForQuery: Message[],
 67    assistantMessages: AssistantMessage[],
 68    systemPrompt: SystemPrompt,
 69    userContext: { [k: string]: string },
 70    systemContext: { [k: string]: string },
 71    toolUseContext: ToolUseContext,
 72    querySource: QuerySource,
 73    stopHookActive?: boolean,
 74  ): AsyncGenerator<
 75    | StreamEvent
 76    | RequestStartEvent
 77    | Message
 78    | TombstoneMessage
 79    | ToolUseSummaryMessage,
 80    StopHookResult
 81  > {
 82    const hookStartTime = Date.now()
 83  
 84    const stopHookContext: REPLHookContext = {
 85      messages: [...messagesForQuery, ...assistantMessages],
 86      systemPrompt,
 87      userContext,
 88      systemContext,
 89      toolUseContext,
 90      querySource,
 91    }
 92    // Only save params for main session queries — subagents must not overwrite.
 93    // Outside the prompt-suggestion gate: the REPL /btw command and the
 94    // side_question SDK control_request both read this snapshot, and neither
 95    // depends on prompt suggestions being enabled.
 96    if (querySource === 'repl_main_thread' || querySource === 'sdk') {
 97      saveCacheSafeParams(createCacheSafeParams(stopHookContext))
 98    }
 99  
100    // Template job classification: when running as a dispatched job, classify
101    // state after each turn. Gate on repl_main_thread so background forks
102    // (extract-memories, auto-dream) don't pollute the timeline with their own
103    // assistant messages. Await the classifier so state.json is written before
104    // the turn returns — otherwise `claude list` shows stale state for the gap.
105    // Env key hardcoded (vs importing JOB_ENV_KEY from jobs/state) to match the
106    // require()-gated jobs/ import pattern above; spawn.test.ts asserts the
107    // string matches.
108    if (
109      feature('TEMPLATES') &&
110      process.env.CLAUDE_JOB_DIR &&
111      querySource.startsWith('repl_main_thread') &&
112      !toolUseContext.agentId
113    ) {
114      // Full turn history — assistantMessages resets each queryLoop iteration,
115      // so tool calls from earlier iterations (Agent spawn, then summary) need
116      // messagesForQuery to be visible in the tool-call summary.
117      const turnAssistantMessages = stopHookContext.messages.filter(
118        (m): m is AssistantMessage => m.type === 'assistant',
119      )
120      const p = jobClassifierModule!
121        .classifyAndWriteState(process.env.CLAUDE_JOB_DIR, turnAssistantMessages)
122        .catch(err => {
123          logForDebugging(`[job] classifier error: ${errorMessage(err)}`, {
124            level: 'error',
125          })
126        })
127      await Promise.race([
128        p,
129        // eslint-disable-next-line no-restricted-syntax -- sleep() has no .unref(); timer must not block exit
130        new Promise<void>(r => setTimeout(r, 60_000).unref()),
131      ])
132    }
133    // --bare / SIMPLE: skip background bookkeeping (prompt suggestion,
134    // memory extraction, auto-dream). Scripted -p calls don't want auto-memory
135    // or forked agents contending for resources during shutdown.
136    if (!isBareMode()) {
137      // Inline env check for dead code elimination in external builds
138      if (!isEnvDefinedFalsy(process.env.CLAUDE_CODE_ENABLE_PROMPT_SUGGESTION)) {
139        void executePromptSuggestion(stopHookContext)
140      }
141      if (
142        feature('EXTRACT_MEMORIES') &&
143        !toolUseContext.agentId &&
144        isExtractModeActive()
145      ) {
146        // Fire-and-forget in both interactive and non-interactive. For -p/SDK,
147        // print.ts drains the in-flight promise after flushing the response
148        // but before gracefulShutdownSync (see drainPendingExtraction).
149        void extractMemoriesModule!.executeExtractMemories(
150          stopHookContext,
151          toolUseContext.appendSystemMessage,
152        )
153      }
154      if (!toolUseContext.agentId) {
155        void executeAutoDream(stopHookContext, toolUseContext.appendSystemMessage)
156      }
157    }
158  
159    // chicago MCP: auto-unhide + lock release at turn end.
160    // Main thread only — the CU lock is a process-wide module-level variable,
161    // so a subagent's stopHooks releasing it leaves the main thread's cleanup
162    // seeing isLockHeldLocally()===false → no exit notification, and unhides
163    // mid-turn. Subagents don't start CU sessions so this is a pure skip.
164    if (feature('CHICAGO_MCP') && !toolUseContext.agentId) {
165      try {
166        const { cleanupComputerUseAfterTurn } = await import(
167          '../utils/computerUse/cleanup.js'
168        )
169        await cleanupComputerUseAfterTurn(toolUseContext)
170      } catch {
171        // Failures are silent — this is dogfooding cleanup, not critical path
172      }
173    }
174  
175    try {
176      const blockingErrors = []
177      const appState = toolUseContext.getAppState()
178      const permissionMode = appState.toolPermissionContext.mode
179  
180      const generator = executeStopHooks(
181        permissionMode,
182        toolUseContext.abortController.signal,
183        undefined,
184        stopHookActive ?? false,
185        toolUseContext.agentId,
186        toolUseContext,
187        [...messagesForQuery, ...assistantMessages],
188        toolUseContext.agentType,
189      )
190  
191      // Consume all progress messages and get blocking errors
192      let stopHookToolUseID = ''
193      let hookCount = 0
194      let preventedContinuation = false
195      let stopReason = ''
196      let hasOutput = false
197      const hookErrors: string[] = []
198      const hookInfos: StopHookInfo[] = []
199  
200      for await (const result of generator) {
201        if (result.message) {
202          yield result.message
203          // Track toolUseID from progress messages and count hooks
204          if (result.message.type === 'progress' && result.message.toolUseID) {
205            stopHookToolUseID = result.message.toolUseID
206            hookCount++
207            // Extract hook command and prompt text from progress data
208            const progressData = result.message.data as HookProgress
209            if (progressData.command) {
210              hookInfos.push({
211                command: progressData.command,
212                promptText: progressData.promptText,
213              })
214            }
215          }
216          // Track errors and output from attachments
217          if (result.message.type === 'attachment') {
218            const attachment = result.message.attachment
219            if (
220              'hookEvent' in attachment &&
221              (attachment.hookEvent === 'Stop' ||
222                attachment.hookEvent === 'SubagentStop')
223            ) {
224              if (attachment.type === 'hook_non_blocking_error') {
225                hookErrors.push(
226                  attachment.stderr || `Exit code ${attachment.exitCode}`,
227                )
228                // Non-blocking errors always have output
229                hasOutput = true
230              } else if (attachment.type === 'hook_error_during_execution') {
231                hookErrors.push(attachment.content)
232                hasOutput = true
233              } else if (attachment.type === 'hook_success') {
234                // Check if successful hook produced any stdout/stderr
235                if (
236                  (attachment.stdout && attachment.stdout.trim()) ||
237                  (attachment.stderr && attachment.stderr.trim())
238                ) {
239                  hasOutput = true
240                }
241              }
242              // Extract per-hook duration for timing visibility.
243              // Hooks run in parallel; match by command + first unassigned entry.
244              if ('durationMs' in attachment && 'command' in attachment) {
245                const info = hookInfos.find(
246                  i =>
247                    i.command === attachment.command &&
248                    i.durationMs === undefined,
249                )
250                if (info) {
251                  info.durationMs = attachment.durationMs
252                }
253              }
254            }
255          }
256        }
257        if (result.blockingError) {
258          const userMessage = createUserMessage({
259            content: getStopHookMessage(result.blockingError),
260            isMeta: true, // Hide from UI (shown in summary message instead)
261          })
262          blockingErrors.push(userMessage)
263          yield userMessage
264          hasOutput = true
265          // Add to hookErrors so it appears in the summary
266          hookErrors.push(result.blockingError.blockingError)
267        }
268        // Check if hook wants to prevent continuation
269        if (result.preventContinuation) {
270          preventedContinuation = true
271          stopReason = result.stopReason || 'Stop hook prevented continuation'
272          // Create attachment to track the stopped continuation (for structured data)
273          yield createAttachmentMessage({
274            type: 'hook_stopped_continuation',
275            message: stopReason,
276            hookName: 'Stop',
277            toolUseID: stopHookToolUseID,
278            hookEvent: 'Stop',
279          })
280        }
281  
282        // Check if we were aborted during hook execution
283        if (toolUseContext.abortController.signal.aborted) {
284          logEvent('tengu_pre_stop_hooks_cancelled', {
285            queryChainId: toolUseContext.queryTracking
286              ?.chainId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
287  
288            queryDepth: toolUseContext.queryTracking?.depth,
289          })
290          yield createUserInterruptionMessage({
291            toolUse: false,
292          })
293          return { blockingErrors: [], preventContinuation: true }
294        }
295      }
296  
297      // Create summary system message if hooks ran
298      if (hookCount > 0) {
299        yield createStopHookSummaryMessage(
300          hookCount,
301          hookInfos,
302          hookErrors,
303          preventedContinuation,
304          stopReason,
305          hasOutput,
306          'suggestion',
307          stopHookToolUseID,
308        )
309  
310        // Send notification about errors (shown in verbose/transcript mode via ctrl+o)
311        if (hookErrors.length > 0) {
312          const expandShortcut = getShortcutDisplay(
313            'app:toggleTranscript',
314            'Global',
315            'ctrl+o',
316          )
317          toolUseContext.addNotification?.({
318            key: 'stop-hook-error',
319            text: `Stop hook error occurred \u00b7 ${expandShortcut} to see`,
320            priority: 'immediate',
321          })
322        }
323      }
324  
325      if (preventedContinuation) {
326        return { blockingErrors: [], preventContinuation: true }
327      }
328  
329      // Collect blocking errors from stop hooks
330      if (blockingErrors.length > 0) {
331        return { blockingErrors, preventContinuation: false }
332      }
333  
334      // After Stop hooks pass, run TeammateIdle and TaskCompleted hooks if this is a teammate
335      if (isTeammate()) {
336        const teammateName = getAgentName() ?? ''
337        const teamName = getTeamName() ?? ''
338        const teammateBlockingErrors: Message[] = []
339        let teammatePreventedContinuation = false
340        let teammateStopReason: string | undefined
341        // Each hook executor generates its own toolUseID — capture from progress
342        // messages (same pattern as stopHookToolUseID at L142), not the Stop ID.
343        let teammateHookToolUseID = ''
344  
345        // Run TaskCompleted hooks for any in-progress tasks owned by this teammate
346        const taskListId = getTaskListId()
347        const tasks = await listTasks(taskListId)
348        const inProgressTasks = tasks.filter(
349          t => t.status === 'in_progress' && t.owner === teammateName,
350        )
351  
352        for (const task of inProgressTasks) {
353          const taskCompletedGenerator = executeTaskCompletedHooks(
354            task.id,
355            task.subject,
356            task.description,
357            teammateName,
358            teamName,
359            permissionMode,
360            toolUseContext.abortController.signal,
361            undefined,
362            toolUseContext,
363          )
364  
365          for await (const result of taskCompletedGenerator) {
366            if (result.message) {
367              if (
368                result.message.type === 'progress' &&
369                result.message.toolUseID
370              ) {
371                teammateHookToolUseID = result.message.toolUseID
372              }
373              yield result.message
374            }
375            if (result.blockingError) {
376              const userMessage = createUserMessage({
377                content: getTaskCompletedHookMessage(result.blockingError),
378                isMeta: true,
379              })
380              teammateBlockingErrors.push(userMessage)
381              yield userMessage
382            }
383            // Match Stop hook behavior: allow preventContinuation/stopReason
384            if (result.preventContinuation) {
385              teammatePreventedContinuation = true
386              teammateStopReason =
387                result.stopReason || 'TaskCompleted hook prevented continuation'
388              yield createAttachmentMessage({
389                type: 'hook_stopped_continuation',
390                message: teammateStopReason,
391                hookName: 'TaskCompleted',
392                toolUseID: teammateHookToolUseID,
393                hookEvent: 'TaskCompleted',
394              })
395            }
396            if (toolUseContext.abortController.signal.aborted) {
397              return { blockingErrors: [], preventContinuation: true }
398            }
399          }
400        }
401  
402        // Run TeammateIdle hooks
403        const teammateIdleGenerator = executeTeammateIdleHooks(
404          teammateName,
405          teamName,
406          permissionMode,
407          toolUseContext.abortController.signal,
408        )
409  
410        for await (const result of teammateIdleGenerator) {
411          if (result.message) {
412            if (result.message.type === 'progress' && result.message.toolUseID) {
413              teammateHookToolUseID = result.message.toolUseID
414            }
415            yield result.message
416          }
417          if (result.blockingError) {
418            const userMessage = createUserMessage({
419              content: getTeammateIdleHookMessage(result.blockingError),
420              isMeta: true,
421            })
422            teammateBlockingErrors.push(userMessage)
423            yield userMessage
424          }
425          // Match Stop hook behavior: allow preventContinuation/stopReason
426          if (result.preventContinuation) {
427            teammatePreventedContinuation = true
428            teammateStopReason =
429              result.stopReason || 'TeammateIdle hook prevented continuation'
430            yield createAttachmentMessage({
431              type: 'hook_stopped_continuation',
432              message: teammateStopReason,
433              hookName: 'TeammateIdle',
434              toolUseID: teammateHookToolUseID,
435              hookEvent: 'TeammateIdle',
436            })
437          }
438          if (toolUseContext.abortController.signal.aborted) {
439            return { blockingErrors: [], preventContinuation: true }
440          }
441        }
442  
443        if (teammatePreventedContinuation) {
444          return { blockingErrors: [], preventContinuation: true }
445        }
446  
447        if (teammateBlockingErrors.length > 0) {
448          return {
449            blockingErrors: teammateBlockingErrors,
450            preventContinuation: false,
451          }
452        }
453      }
454  
455      return { blockingErrors: [], preventContinuation: false }
456    } catch (error) {
457      const durationMs = Date.now() - hookStartTime
458      logEvent('tengu_stop_hook_error', {
459        duration: durationMs,
460  
461        queryChainId: toolUseContext.queryTracking
462          ?.chainId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
463        queryDepth: toolUseContext.queryTracking?.depth,
464      })
465      // Yield a system message that is not visible to the model for the user
466      // to debug their hook.
467      yield createSystemMessage(
468        `Stop hook failed: ${errorMessage(error)}`,
469        'warning',
470      )
471      return { blockingErrors: [], preventContinuation: false }
472    }
473  }