1  import type {
  2    AsyncHookJSONOutput,
  3    HookEvent,
  4    SyncHookJSONOutput,
  5  } from 'src/entrypoints/agentSdkTypes.js'
  6  import { logForDebugging } from '../debug.js'
  7  import type { ShellCommand } from '../ShellCommand.js'
  8  import { invalidateSessionEnvCache } from '../sessionEnvironment.js'
  9  import { jsonParse, jsonStringify } from '../slowOperations.js'
 10  import { emitHookResponse, startHookProgressInterval } from './hookEvents.js'
 11  
 12  export type PendingAsyncHook = {
 13    processId: string
 14    hookId: string
 15    hookName: string
 16    hookEvent: HookEvent | 'StatusLine' | 'FileSuggestion'
 17    toolName?: string
 18    pluginId?: string
 19    startTime: number
 20    timeout: number
 21    command: string
 22    responseAttachmentSent: boolean
 23    shellCommand?: ShellCommand
 24    stopProgressInterval: () => void
 25  }
 26  
 27  // Global registry state
 28  const pendingHooks = new Map<string, PendingAsyncHook>()
 29  
 30  export function registerPendingAsyncHook({
 31    processId,
 32    hookId,
 33    asyncResponse,
 34    hookName,
 35    hookEvent,
 36    command,
 37    shellCommand,
 38    toolName,
 39    pluginId,
 40  }: {
 41    processId: string
 42    hookId: string
 43    asyncResponse: AsyncHookJSONOutput
 44    hookName: string
 45    hookEvent: HookEvent | 'StatusLine' | 'FileSuggestion'
 46    command: string
 47    shellCommand: ShellCommand
 48    toolName?: string
 49    pluginId?: string
 50  }): void {
 51    const timeout = asyncResponse.asyncTimeout || 15000 // Default 15s
 52    logForDebugging(
 53      `Hooks: Registering async hook ${processId} (${hookName}) with timeout ${timeout}ms`,
 54    )
 55    const stopProgressInterval = startHookProgressInterval({
 56      hookId,
 57      hookName,
 58      hookEvent,
 59      getOutput: async () => {
 60        const taskOutput = pendingHooks.get(processId)?.shellCommand?.taskOutput
 61        if (!taskOutput) {
 62          return { stdout: '', stderr: '', output: '' }
 63        }
 64        const stdout = await taskOutput.getStdout()
 65        const stderr = taskOutput.getStderr()
 66        return { stdout, stderr, output: stdout + stderr }
 67      },
 68    })
 69    pendingHooks.set(processId, {
 70      processId,
 71      hookId,
 72      hookName,
 73      hookEvent,
 74      toolName,
 75      pluginId,
 76      command,
 77      startTime: Date.now(),
 78      timeout,
 79      responseAttachmentSent: false,
 80      shellCommand,
 81      stopProgressInterval,
 82    })
 83  }
 84  
 85  export function getPendingAsyncHooks(): PendingAsyncHook[] {
 86    return Array.from(pendingHooks.values()).filter(
 87      hook => !hook.responseAttachmentSent,
 88    )
 89  }
 90  
 91  async function finalizeHook(
 92    hook: PendingAsyncHook,
 93    exitCode: number,
 94    outcome: 'success' | 'error' | 'cancelled',
 95  ): Promise<void> {
 96    hook.stopProgressInterval()
 97    const taskOutput = hook.shellCommand?.taskOutput
 98    const stdout = taskOutput ? await taskOutput.getStdout() : ''
 99    const stderr = taskOutput?.getStderr() ?? ''
100    hook.shellCommand?.cleanup()
101    emitHookResponse({
102      hookId: hook.hookId,
103      hookName: hook.hookName,
104      hookEvent: hook.hookEvent,
105      output: stdout + stderr,
106      stdout,
107      stderr,
108      exitCode,
109      outcome,
110    })
111  }
112  
113  export async function checkForAsyncHookResponses(): Promise<
114    Array<{
115      processId: string
116      response: SyncHookJSONOutput
117      hookName: string
118      hookEvent: HookEvent | 'StatusLine' | 'FileSuggestion'
119      toolName?: string
120      pluginId?: string
121      stdout: string
122      stderr: string
123      exitCode?: number
124    }>
125  > {
126    const responses: {
127      processId: string
128      response: SyncHookJSONOutput
129      hookName: string
130      hookEvent: HookEvent | 'StatusLine' | 'FileSuggestion'
131      toolName?: string
132      pluginId?: string
133      stdout: string
134      stderr: string
135      exitCode?: number
136    }[] = []
137  
138    const pendingCount = pendingHooks.size
139    logForDebugging(`Hooks: Found ${pendingCount} total hooks in registry`)
140  
141    // Snapshot hooks before processing — we'll mutate the map after.
142    const hooks = Array.from(pendingHooks.values())
143  
144    const settled = await Promise.allSettled(
145      hooks.map(async hook => {
146        const stdout = (await hook.shellCommand?.taskOutput.getStdout()) ?? ''
147        const stderr = hook.shellCommand?.taskOutput.getStderr() ?? ''
148        logForDebugging(
149          `Hooks: Checking hook ${hook.processId} (${hook.hookName}) - attachmentSent: ${hook.responseAttachmentSent}, stdout length: ${stdout.length}`,
150        )
151  
152        if (!hook.shellCommand) {
153          logForDebugging(
154            `Hooks: Hook ${hook.processId} has no shell command, removing from registry`,
155          )
156          hook.stopProgressInterval()
157          return { type: 'remove' as const, processId: hook.processId }
158        }
159  
160        logForDebugging(`Hooks: Hook shell status ${hook.shellCommand.status}`)
161  
162        if (hook.shellCommand.status === 'killed') {
163          logForDebugging(
164            `Hooks: Hook ${hook.processId} is ${hook.shellCommand.status}, removing from registry`,
165          )
166          hook.stopProgressInterval()
167          hook.shellCommand.cleanup()
168          return { type: 'remove' as const, processId: hook.processId }
169        }
170  
171        if (hook.shellCommand.status !== 'completed') {
172          return { type: 'skip' as const }
173        }
174  
175        if (hook.responseAttachmentSent || !stdout.trim()) {
176          logForDebugging(
177            `Hooks: Skipping hook ${hook.processId} - already delivered/sent or no stdout`,
178          )
179          hook.stopProgressInterval()
180          return { type: 'remove' as const, processId: hook.processId }
181        }
182  
183        const lines = stdout.split('\n')
184        logForDebugging(
185          `Hooks: Processing ${lines.length} lines of stdout for ${hook.processId}`,
186        )
187  
188        const execResult = await hook.shellCommand.result
189        const exitCode = execResult.code
190  
191        let response: SyncHookJSONOutput = {}
192        for (const line of lines) {
193          if (line.trim().startsWith('{')) {
194            logForDebugging(
195              `Hooks: Found JSON line: ${line.trim().substring(0, 100)}...`,
196            )
197            try {
198              const parsed = jsonParse(line.trim())
199              if (!('async' in parsed)) {
200                logForDebugging(
201                  `Hooks: Found sync response from ${hook.processId}: ${jsonStringify(parsed)}`,
202                )
203                response = parsed
204                break
205              }
206            } catch {
207              logForDebugging(
208                `Hooks: Failed to parse JSON from ${hook.processId}: ${line.trim()}`,
209              )
210            }
211          }
212        }
213  
214        hook.responseAttachmentSent = true
215        await finalizeHook(hook, exitCode, exitCode === 0 ? 'success' : 'error')
216  
217        return {
218          type: 'response' as const,
219          processId: hook.processId,
220          isSessionStart: hook.hookEvent === 'SessionStart',
221          payload: {
222            processId: hook.processId,
223            response,
224            hookName: hook.hookName,
225            hookEvent: hook.hookEvent,
226            toolName: hook.toolName,
227            pluginId: hook.pluginId,
228            stdout,
229            stderr,
230            exitCode,
231          },
232        }
233      }),
234    )
235  
236    // allSettled — isolate failures so one throwing callback doesn't orphan
237    // already-applied side effects (responseAttachmentSent, finalizeHook) from others.
238    let sessionStartCompleted = false
239    for (const s of settled) {
240      if (s.status !== 'fulfilled') {
241        logForDebugging(
242          `Hooks: checkForAsyncHookResponses callback rejected: ${s.reason}`,
243          { level: 'error' },
244        )
245        continue
246      }
247      const r = s.value
248      if (r.type === 'remove') {
249        pendingHooks.delete(r.processId)
250      } else if (r.type === 'response') {
251        responses.push(r.payload)
252        pendingHooks.delete(r.processId)
253        if (r.isSessionStart) sessionStartCompleted = true
254      }
255    }
256  
257    if (sessionStartCompleted) {
258      logForDebugging(
259        `Invalidating session env cache after SessionStart hook completed`,
260      )
261      invalidateSessionEnvCache()
262    }
263  
264    logForDebugging(
265      `Hooks: checkForNewResponses returning ${responses.length} responses`,
266    )
267    return responses
268  }
269  
270  export function removeDeliveredAsyncHooks(processIds: string[]): void {
271    for (const processId of processIds) {
272      const hook = pendingHooks.get(processId)
273      if (hook && hook.responseAttachmentSent) {
274        logForDebugging(`Hooks: Removing delivered hook ${processId}`)
275        hook.stopProgressInterval()
276        pendingHooks.delete(processId)
277      }
278    }
279  }
280  
281  export async function finalizePendingAsyncHooks(): Promise<void> {
282    const hooks = Array.from(pendingHooks.values())
283    await Promise.all(
284      hooks.map(async hook => {
285        if (hook.shellCommand?.status === 'completed') {
286          const result = await hook.shellCommand.result
287          await finalizeHook(
288            hook,
289            result.code,
290            result.code === 0 ? 'success' : 'error',
291          )
292        } else {
293          if (hook.shellCommand && hook.shellCommand.status !== 'killed') {
294            hook.shellCommand.kill()
295          }
296          await finalizeHook(hook, 1, 'cancelled')
297        }
298      }),
299    )
300    pendingHooks.clear()
301  }
302  
303  // Test utility function to clear all hooks
304  export function clearAllAsyncHooks(): void {
305    for (const hook of pendingHooks.values()) {
306      hook.stopProgressInterval()
307    }
308    pendingHooks.clear()
309  }