1  /**
  2   * Swarm Permission Poller Hook
  3   *
  4   * This hook polls for permission responses from the team leader when running
  5   * as a worker agent in a swarm. When a response is received, it calls the
  6   * appropriate callback (onAllow/onReject) to continue execution.
  7   *
  8   * This hook should be used in conjunction with the worker-side integration
  9   * in useCanUseTool.ts, which creates pending requests that this hook monitors.
 10   */
 11  
 12  import { useCallback, useEffect, useRef } from 'react'
 13  import { useInterval } from 'usehooks-ts'
 14  import { logForDebugging } from '../utils/debug.js'
 15  import { errorMessage } from '../utils/errors.js'
 16  import {
 17    type PermissionUpdate,
 18    permissionUpdateSchema,
 19  } from '../utils/permissions/PermissionUpdateSchema.js'
 20  import {
 21    isSwarmWorker,
 22    type PermissionResponse,
 23    pollForResponse,
 24    removeWorkerResponse,
 25  } from '../utils/swarm/permissionSync.js'
 26  import { getAgentName, getTeamName } from '../utils/teammate.js'
 27  
 28  const POLL_INTERVAL_MS = 500
 29  
 30  /**
 31   * Validate permissionUpdates from external sources (mailbox IPC, disk polling).
 32   * Malformed entries from buggy/old teammate processes are filtered out rather
 33   * than propagated unchecked into callback.onAllow().
 34   */
 35  function parsePermissionUpdates(raw: unknown): PermissionUpdate[] {
 36    if (!Array.isArray(raw)) {
 37      return []
 38    }
 39    const schema = permissionUpdateSchema()
 40    const valid: PermissionUpdate[] = []
 41    for (const entry of raw) {
 42      const result = schema.safeParse(entry)
 43      if (result.success) {
 44        valid.push(result.data)
 45      } else {
 46        logForDebugging(
 47          `[SwarmPermissionPoller] Dropping malformed permissionUpdate entry: ${result.error.message}`,
 48          { level: 'warn' },
 49        )
 50      }
 51    }
 52    return valid
 53  }
 54  
 55  /**
 56   * Callback signature for handling permission responses
 57   */
 58  export type PermissionResponseCallback = {
 59    requestId: string
 60    toolUseId: string
 61    onAllow: (
 62      updatedInput: Record<string, unknown> | undefined,
 63      permissionUpdates: PermissionUpdate[],
 64      feedback?: string,
 65    ) => void
 66    onReject: (feedback?: string) => void
 67  }
 68  
 69  /**
 70   * Registry for pending permission request callbacks
 71   * This allows the poller to find and invoke the right callbacks when responses arrive
 72   */
 73  type PendingCallbackRegistry = Map<string, PermissionResponseCallback>
 74  
 75  // Module-level registry that persists across renders
 76  const pendingCallbacks: PendingCallbackRegistry = new Map()
 77  
 78  /**
 79   * Register a callback for a pending permission request
 80   * Called by useCanUseTool when a worker submits a permission request
 81   */
 82  export function registerPermissionCallback(
 83    callback: PermissionResponseCallback,
 84  ): void {
 85    pendingCallbacks.set(callback.requestId, callback)
 86    logForDebugging(
 87      `[SwarmPermissionPoller] Registered callback for request ${callback.requestId}`,
 88    )
 89  }
 90  
 91  /**
 92   * Unregister a callback (e.g., when the request is resolved locally or times out)
 93   */
 94  export function unregisterPermissionCallback(requestId: string): void {
 95    pendingCallbacks.delete(requestId)
 96    logForDebugging(
 97      `[SwarmPermissionPoller] Unregistered callback for request ${requestId}`,
 98    )
 99  }
100  
101  /**
102   * Check if a request has a registered callback
103   */
104  export function hasPermissionCallback(requestId: string): boolean {
105    return pendingCallbacks.has(requestId)
106  }
107  
108  /**
109   * Clear all pending callbacks (both permission and sandbox).
110   * Called from clearSessionCaches() on /clear to reset stale state,
111   * and also used in tests for isolation.
112   */
113  export function clearAllPendingCallbacks(): void {
114    pendingCallbacks.clear()
115    pendingSandboxCallbacks.clear()
116  }
117  
118  /**
119   * Process a permission response from a mailbox message.
120   * This is called by the inbox poller when it detects a permission_response message.
121   *
122   * @returns true if the response was processed, false if no callback was registered
123   */
124  export function processMailboxPermissionResponse(params: {
125    requestId: string
126    decision: 'approved' | 'rejected'
127    feedback?: string
128    updatedInput?: Record<string, unknown>
129    permissionUpdates?: unknown
130  }): boolean {
131    const callback = pendingCallbacks.get(params.requestId)
132  
133    if (!callback) {
134      logForDebugging(
135        `[SwarmPermissionPoller] No callback registered for mailbox response ${params.requestId}`,
136      )
137      return false
138    }
139  
140    logForDebugging(
141      `[SwarmPermissionPoller] Processing mailbox response for request ${params.requestId}: ${params.decision}`,
142    )
143  
144    // Remove from registry before invoking callback
145    pendingCallbacks.delete(params.requestId)
146  
147    if (params.decision === 'approved') {
148      const permissionUpdates = parsePermissionUpdates(params.permissionUpdates)
149      const updatedInput = params.updatedInput
150      callback.onAllow(updatedInput, permissionUpdates)
151    } else {
152      callback.onReject(params.feedback)
153    }
154  
155    return true
156  }
157  
158  // ============================================================================
159  // Sandbox Permission Callback Registry
160  // ============================================================================
161  
162  /**
163   * Callback signature for handling sandbox permission responses
164   */
165  export type SandboxPermissionResponseCallback = {
166    requestId: string
167    host: string
168    resolve: (allow: boolean) => void
169  }
170  
171  // Module-level registry for sandbox permission callbacks
172  const pendingSandboxCallbacks: Map<string, SandboxPermissionResponseCallback> =
173    new Map()
174  
175  /**
176   * Register a callback for a pending sandbox permission request
177   * Called when a worker sends a sandbox permission request to the leader
178   */
179  export function registerSandboxPermissionCallback(
180    callback: SandboxPermissionResponseCallback,
181  ): void {
182    pendingSandboxCallbacks.set(callback.requestId, callback)
183    logForDebugging(
184      `[SwarmPermissionPoller] Registered sandbox callback for request ${callback.requestId}`,
185    )
186  }
187  
188  /**
189   * Check if a sandbox request has a registered callback
190   */
191  export function hasSandboxPermissionCallback(requestId: string): boolean {
192    return pendingSandboxCallbacks.has(requestId)
193  }
194  
195  /**
196   * Process a sandbox permission response from a mailbox message.
197   * Called by the inbox poller when it detects a sandbox_permission_response message.
198   *
199   * @returns true if the response was processed, false if no callback was registered
200   */
201  export function processSandboxPermissionResponse(params: {
202    requestId: string
203    host: string
204    allow: boolean
205  }): boolean {
206    const callback = pendingSandboxCallbacks.get(params.requestId)
207  
208    if (!callback) {
209      logForDebugging(
210        `[SwarmPermissionPoller] No sandbox callback registered for request ${params.requestId}`,
211      )
212      return false
213    }
214  
215    logForDebugging(
216      `[SwarmPermissionPoller] Processing sandbox response for request ${params.requestId}: allow=${params.allow}`,
217    )
218  
219    // Remove from registry before invoking callback
220    pendingSandboxCallbacks.delete(params.requestId)
221  
222    // Resolve the promise with the allow decision
223    callback.resolve(params.allow)
224  
225    return true
226  }
227  
228  /**
229   * Process a permission response by invoking the registered callback
230   */
231  function processResponse(response: PermissionResponse): boolean {
232    const callback = pendingCallbacks.get(response.requestId)
233  
234    if (!callback) {
235      logForDebugging(
236        `[SwarmPermissionPoller] No callback registered for request ${response.requestId}`,
237      )
238      return false
239    }
240  
241    logForDebugging(
242      `[SwarmPermissionPoller] Processing response for request ${response.requestId}: ${response.decision}`,
243    )
244  
245    // Remove from registry before invoking callback
246    pendingCallbacks.delete(response.requestId)
247  
248    if (response.decision === 'approved') {
249      const permissionUpdates = parsePermissionUpdates(response.permissionUpdates)
250      const updatedInput = response.updatedInput
251      callback.onAllow(updatedInput, permissionUpdates)
252    } else {
253      callback.onReject(response.feedback)
254    }
255  
256    return true
257  }
258  
259  /**
260   * Hook that polls for permission responses when running as a swarm worker.
261   *
262   * This hook:
263   * 1. Only activates when isSwarmWorker() returns true
264   * 2. Polls every 500ms for responses
265   * 3. When a response is found, invokes the registered callback
266   * 4. Cleans up the response file after processing
267   */
268  export function useSwarmPermissionPoller(): void {
269    const isProcessingRef = useRef(false)
270  
271    const poll = useCallback(async () => {
272      // Don't poll if not a swarm worker
273      if (!isSwarmWorker()) {
274        return
275      }
276  
277      // Prevent concurrent polling
278      if (isProcessingRef.current) {
279        return
280      }
281  
282      // Don't poll if no callbacks are registered
283      if (pendingCallbacks.size === 0) {
284        return
285      }
286  
287      isProcessingRef.current = true
288  
289      try {
290        const agentName = getAgentName()
291        const teamName = getTeamName()
292  
293        if (!agentName || !teamName) {
294          return
295        }
296  
297        // Check each pending request for a response
298        for (const [requestId, _callback] of pendingCallbacks) {
299          const response = await pollForResponse(requestId, agentName, teamName)
300  
301          if (response) {
302            // Process the response
303            const processed = processResponse(response)
304  
305            if (processed) {
306              // Clean up the response from the worker's inbox
307              await removeWorkerResponse(requestId, agentName, teamName)
308            }
309          }
310        }
311      } catch (error) {
312        logForDebugging(
313          `[SwarmPermissionPoller] Error during poll: ${errorMessage(error)}`,
314        )
315      } finally {
316        isProcessingRef.current = false
317      }
318    }, [])
319  
320    // Only poll if we're a swarm worker
321    const shouldPoll = isSwarmWorker()
322    useInterval(() => void poll(), shouldPoll ? POLL_INTERVAL_MS : null)
323  
324    // Initial poll on mount
325    useEffect(() => {
326      if (isSwarmWorker()) {
327        void poll()
328      }
329    }, [poll])
330  }