1  import { HOOK_EVENTS, type HookEvent } from 'src/entrypoints/agentSdkTypes.js'
  2  import type { AppState } from 'src/state/AppState.js'
  3  import type { Message } from 'src/types/message.js'
  4  import { logForDebugging } from '../debug.js'
  5  import type { AggregatedHookResult } from '../hooks.js'
  6  import type { HookCommand } from '../settings/types.js'
  7  import { isHookEqual } from './hooksSettings.js'
  8  
  9  type OnHookSuccess = (
 10    hook: HookCommand | FunctionHook,
 11    result: AggregatedHookResult,
 12  ) => void
 13  
 14  /** Function hook callback - returns true if check passes, false to block */
 15  export type FunctionHookCallback = (
 16    messages: Message[],
 17    signal?: AbortSignal,
 18  ) => boolean | Promise<boolean>
 19  
 20  /**
 21   * Function hook type with callback embedded.
 22   * Session-scoped only, cannot be persisted to settings.json.
 23   */
 24  export type FunctionHook = {
 25    type: 'function'
 26    id?: string // Optional unique ID for removal
 27    timeout?: number
 28    callback: FunctionHookCallback
 29    errorMessage: string
 30    statusMessage?: string
 31  }
 32  
 33  type SessionHookMatcher = {
 34    matcher: string
 35    skillRoot?: string
 36    hooks: Array<{
 37      hook: HookCommand | FunctionHook
 38      onHookSuccess?: OnHookSuccess
 39    }>
 40  }
 41  
 42  export type SessionStore = {
 43    hooks: {
 44      [event in HookEvent]?: SessionHookMatcher[]
 45    }
 46  }
 47  
 48  /**
 49   * Map (not Record) so .set/.delete don't change the container's identity.
 50   * Mutator functions mutate the Map and return prev unchanged, letting
 51   * store.ts's Object.is(next, prev) check short-circuit and skip listener
 52   * notification. Session hooks are ephemeral per-agent runtime callbacks,
 53   * never reactively read (only getAppState() snapshots in the query loop).
 54   * Same pattern as agentControllers on LocalWorkflowTaskState.
 55   *
 56   * This matters under high-concurrency workflows: parallel() with N
 57   * schema-mode agents fires N addFunctionHook calls in one synchronous
 58   * tick. With a Record + spread, each call cost O(N) to copy the growing
 59   * map (O(N²) total) plus fired all ~30 store listeners. With Map: .set()
 60   * is O(1), return prev means zero listener fires.
 61   */
 62  export type SessionHooksState = Map<string, SessionStore>
 63  
 64  /**
 65   * Add a command or prompt hook to the session.
 66   * Session hooks are temporary, in-memory only, and cleared when session ends.
 67   */
 68  export function addSessionHook(
 69    setAppState: (updater: (prev: AppState) => AppState) => void,
 70    sessionId: string,
 71    event: HookEvent,
 72    matcher: string,
 73    hook: HookCommand,
 74    onHookSuccess?: OnHookSuccess,
 75    skillRoot?: string,
 76  ): void {
 77    addHookToSession(
 78      setAppState,
 79      sessionId,
 80      event,
 81      matcher,
 82      hook,
 83      onHookSuccess,
 84      skillRoot,
 85    )
 86  }
 87  
 88  /**
 89   * Add a function hook to the session.
 90   * Function hooks execute TypeScript callbacks in-memory for validation.
 91   * @returns The hook ID (for removal)
 92   */
 93  export function addFunctionHook(
 94    setAppState: (updater: (prev: AppState) => AppState) => void,
 95    sessionId: string,
 96    event: HookEvent,
 97    matcher: string,
 98    callback: FunctionHookCallback,
 99    errorMessage: string,
100    options?: {
101      timeout?: number
102      id?: string
103    },
104  ): string {
105    const id = options?.id || `function-hook-${Date.now()}-${Math.random()}`
106    const hook: FunctionHook = {
107      type: 'function',
108      id,
109      timeout: options?.timeout || 5000,
110      callback,
111      errorMessage,
112    }
113    addHookToSession(setAppState, sessionId, event, matcher, hook)
114    return id
115  }
116  
117  /**
118   * Remove a function hook by ID from the session.
119   */
120  export function removeFunctionHook(
121    setAppState: (updater: (prev: AppState) => AppState) => void,
122    sessionId: string,
123    event: HookEvent,
124    hookId: string,
125  ): void {
126    setAppState(prev => {
127      const store = prev.sessionHooks.get(sessionId)
128      if (!store) {
129        return prev
130      }
131  
132      const eventMatchers = store.hooks[event] || []
133  
134      // Remove the hook with matching ID from all matchers
135      const updatedMatchers = eventMatchers
136        .map(matcher => {
137          const updatedHooks = matcher.hooks.filter(h => {
138            if (h.hook.type !== 'function') return true
139            return h.hook.id !== hookId
140          })
141  
142          return updatedHooks.length > 0
143            ? { ...matcher, hooks: updatedHooks }
144            : null
145        })
146        .filter((m): m is SessionHookMatcher => m !== null)
147  
148      const newHooks =
149        updatedMatchers.length > 0
150          ? { ...store.hooks, [event]: updatedMatchers }
151          : Object.fromEntries(
152              Object.entries(store.hooks).filter(([e]) => e !== event),
153            )
154  
155      prev.sessionHooks.set(sessionId, { hooks: newHooks })
156      return prev
157    })
158  
159    logForDebugging(
160      `Removed function hook ${hookId} for event ${event} in session ${sessionId}`,
161    )
162  }
163  
164  /**
165   * Internal helper to add a hook to session state
166   */
167  function addHookToSession(
168    setAppState: (updater: (prev: AppState) => AppState) => void,
169    sessionId: string,
170    event: HookEvent,
171    matcher: string,
172    hook: HookCommand | FunctionHook,
173    onHookSuccess?: OnHookSuccess,
174    skillRoot?: string,
175  ): void {
176    setAppState(prev => {
177      const store = prev.sessionHooks.get(sessionId) ?? { hooks: {} }
178      const eventMatchers = store.hooks[event] || []
179  
180      // Find existing matcher or create new one
181      const existingMatcherIndex = eventMatchers.findIndex(
182        m => m.matcher === matcher && m.skillRoot === skillRoot,
183      )
184  
185      let updatedMatchers: SessionHookMatcher[]
186      if (existingMatcherIndex >= 0) {
187        // Add to existing matcher
188        updatedMatchers = [...eventMatchers]
189        const existingMatcher = updatedMatchers[existingMatcherIndex]!
190        updatedMatchers[existingMatcherIndex] = {
191          matcher: existingMatcher.matcher,
192          skillRoot: existingMatcher.skillRoot,
193          hooks: [...existingMatcher.hooks, { hook, onHookSuccess }],
194        }
195      } else {
196        // Create new matcher
197        updatedMatchers = [
198          ...eventMatchers,
199          {
200            matcher,
201            skillRoot,
202            hooks: [{ hook, onHookSuccess }],
203          },
204        ]
205      }
206  
207      const newHooks = { ...store.hooks, [event]: updatedMatchers }
208  
209      prev.sessionHooks.set(sessionId, { hooks: newHooks })
210      return prev
211    })
212  
213    logForDebugging(
214      `Added session hook for event ${event} in session ${sessionId}`,
215    )
216  }
217  
218  /**
219   * Remove a specific hook from the session
220   * @param setAppState The function to update the app state
221   * @param sessionId The session ID
222   * @param event The hook event
223   * @param hook The hook command to remove
224   */
225  export function removeSessionHook(
226    setAppState: (updater: (prev: AppState) => AppState) => void,
227    sessionId: string,
228    event: HookEvent,
229    hook: HookCommand,
230  ): void {
231    setAppState(prev => {
232      const store = prev.sessionHooks.get(sessionId)
233      if (!store) {
234        return prev
235      }
236  
237      const eventMatchers = store.hooks[event] || []
238  
239      // Remove the hook from all matchers
240      const updatedMatchers = eventMatchers
241        .map(matcher => {
242          const updatedHooks = matcher.hooks.filter(
243            h => !isHookEqual(h.hook, hook),
244          )
245  
246          return updatedHooks.length > 0
247            ? { ...matcher, hooks: updatedHooks }
248            : null
249        })
250        .filter((m): m is SessionHookMatcher => m !== null)
251  
252      const newHooks =
253        updatedMatchers.length > 0
254          ? { ...store.hooks, [event]: updatedMatchers }
255          : { ...store.hooks }
256  
257      if (updatedMatchers.length === 0) {
258        delete newHooks[event]
259      }
260  
261      prev.sessionHooks.set(sessionId, { ...store, hooks: newHooks })
262      return prev
263    })
264  
265    logForDebugging(
266      `Removed session hook for event ${event} in session ${sessionId}`,
267    )
268  }
269  
270  // Extended hook matcher that includes optional skillRoot for skill-scoped hooks
271  export type SessionDerivedHookMatcher = {
272    matcher: string
273    hooks: HookCommand[]
274    skillRoot?: string
275  }
276  
277  /**
278   * Convert session hook matchers to regular hook matchers
279   * @param sessionMatchers The session hook matchers to convert
280   * @returns Regular hook matchers (with optional skillRoot preserved)
281   */
282  function convertToHookMatchers(
283    sessionMatchers: SessionHookMatcher[],
284  ): SessionDerivedHookMatcher[] {
285    return sessionMatchers.map(sm => ({
286      matcher: sm.matcher,
287      skillRoot: sm.skillRoot,
288      // Filter out function hooks - they can't be persisted to HookMatcher format
289      hooks: sm.hooks
290        .map(h => h.hook)
291        .filter((h): h is HookCommand => h.type !== 'function'),
292    }))
293  }
294  
295  /**
296   * Get all session hooks for a specific event (excluding function hooks)
297   * @param appState The app state
298   * @param sessionId The session ID
299   * @param event Optional event to filter by
300   * @returns Hook matchers for the event, or all hooks if no event specified
301   */
302  export function getSessionHooks(
303    appState: AppState,
304    sessionId: string,
305    event?: HookEvent,
306  ): Map<HookEvent, SessionDerivedHookMatcher[]> {
307    const store = appState.sessionHooks.get(sessionId)
308    if (!store) {
309      return new Map()
310    }
311  
312    const result = new Map<HookEvent, SessionDerivedHookMatcher[]>()
313  
314    if (event) {
315      const sessionMatchers = store.hooks[event]
316      if (sessionMatchers) {
317        result.set(event, convertToHookMatchers(sessionMatchers))
318      }
319      return result
320    }
321  
322    for (const evt of HOOK_EVENTS) {
323      const sessionMatchers = store.hooks[evt]
324      if (sessionMatchers) {
325        result.set(evt, convertToHookMatchers(sessionMatchers))
326      }
327    }
328  
329    return result
330  }
331  
332  type FunctionHookMatcher = {
333    matcher: string
334    hooks: FunctionHook[]
335  }
336  
337  /**
338   * Get all session function hooks for a specific event
339   * Function hooks are kept separate because they can't be persisted to HookMatcher format.
340   * @param appState The app state
341   * @param sessionId The session ID
342   * @param event Optional event to filter by
343   * @returns Function hook matchers for the event
344   */
345  export function getSessionFunctionHooks(
346    appState: AppState,
347    sessionId: string,
348    event?: HookEvent,
349  ): Map<HookEvent, FunctionHookMatcher[]> {
350    const store = appState.sessionHooks.get(sessionId)
351    if (!store) {
352      return new Map()
353    }
354  
355    const result = new Map<HookEvent, FunctionHookMatcher[]>()
356  
357    const extractFunctionHooks = (
358      sessionMatchers: SessionHookMatcher[],
359    ): FunctionHookMatcher[] => {
360      return sessionMatchers
361        .map(sm => ({
362          matcher: sm.matcher,
363          hooks: sm.hooks
364            .map(h => h.hook)
365            .filter((h): h is FunctionHook => h.type === 'function'),
366        }))
367        .filter(m => m.hooks.length > 0)
368    }
369  
370    if (event) {
371      const sessionMatchers = store.hooks[event]
372      if (sessionMatchers) {
373        const functionMatchers = extractFunctionHooks(sessionMatchers)
374        if (functionMatchers.length > 0) {
375          result.set(event, functionMatchers)
376        }
377      }
378      return result
379    }
380  
381    for (const evt of HOOK_EVENTS) {
382      const sessionMatchers = store.hooks[evt]
383      if (sessionMatchers) {
384        const functionMatchers = extractFunctionHooks(sessionMatchers)
385        if (functionMatchers.length > 0) {
386          result.set(evt, functionMatchers)
387        }
388      }
389    }
390  
391    return result
392  }
393  
394  /**
395   * Get the full hook entry (including callbacks) for a specific session hook
396   */
397  export function getSessionHookCallback(
398    appState: AppState,
399    sessionId: string,
400    event: HookEvent,
401    matcher: string,
402    hook: HookCommand | FunctionHook,
403  ):
404    | {
405        hook: HookCommand | FunctionHook
406        onHookSuccess?: OnHookSuccess
407      }
408    | undefined {
409    const store = appState.sessionHooks.get(sessionId)
410    if (!store) {
411      return undefined
412    }
413  
414    const eventMatchers = store.hooks[event]
415    if (!eventMatchers) {
416      return undefined
417    }
418  
419    // Find the hook in the matchers
420    for (const matcherEntry of eventMatchers) {
421      if (matcherEntry.matcher === matcher || matcher === '') {
422        const hookEntry = matcherEntry.hooks.find(h => isHookEqual(h.hook, hook))
423        if (hookEntry) {
424          return hookEntry
425        }
426      }
427    }
428  
429    return undefined
430  }
431  
432  /**
433   * Clear all session hooks for a specific session
434   * @param setAppState The function to update the app state
435   * @param sessionId The session ID
436   */
437  export function clearSessionHooks(
438    setAppState: (updater: (prev: AppState) => AppState) => void,
439    sessionId: string,
440  ): void {
441    setAppState(prev => {
442      prev.sessionHooks.delete(sessionId)
443      return prev
444    })
445  
446    logForDebugging(`Cleared all session hooks for session ${sessionId}`)
447  }