1  import { useCallback } from 'react'
  2  import { hasPermissionsToUseTool } from '../permissions.js'
  3  import { logEvent } from '../services/statsig.js'
  4  import { BashTool, inputSchema } from '../tools/BashTool/BashTool.js'
  5  import { getCommandSubcommandPrefix } from '../utils/commands.js'
  6  import { REJECT_MESSAGE } from '../utils/messages.js'
  7  import type { Tool as ToolType, ToolUseContext } from '../Tool.js'
  8  import { AssistantMessage } from '../query.js'
  9  import { ToolUseConfirm } from '../components/permissions/PermissionRequest.js'
 10  import { AbortError } from '../utils/errors.js'
 11  import { logError } from '../utils/log.js'
 12  
 13  type SetState<T> = React.Dispatch<React.SetStateAction<T>>
 14  
 15  export type CanUseToolFn = (
 16    tool: ToolType,
 17    input: { [key: string]: unknown },
 18    toolUseContext: ToolUseContext,
 19    assistantMessage: AssistantMessage,
 20  ) => Promise<{ result: true } | { result: false; message: string }>
 21  
 22  function useCanUseTool(
 23    setToolUseConfirm: SetState<ToolUseConfirm | null>,
 24  ): CanUseToolFn {
 25    return useCallback<CanUseToolFn>(
 26      async (tool, input, toolUseContext, assistantMessage) => {
 27        return new Promise(resolve => {
 28          function logCancelledEvent() {
 29            logEvent('tengu_tool_use_cancelled', {
 30              messageID: assistantMessage.message.id,
 31              toolName: tool.name,
 32            })
 33          }
 34  
 35          function resolveWithCancelledAndAbortAllToolCalls() {
 36            resolve({
 37              result: false,
 38              message: REJECT_MESSAGE,
 39            })
 40            // Trigger a synthetic assistant message in query(), to cancel
 41            // any other pending tool uses and stop further requests to the
 42            // API and wait for user input.
 43            toolUseContext.abortController.abort()
 44          }
 45  
 46          if (toolUseContext.abortController.signal.aborted) {
 47            logCancelledEvent()
 48            resolveWithCancelledAndAbortAllToolCalls()
 49            return
 50          }
 51  
 52          return hasPermissionsToUseTool(
 53            tool,
 54            input,
 55            toolUseContext,
 56            assistantMessage,
 57          )
 58            .then(async result => {
 59              // Has permissions to use tool, granted in config
 60              if (result.result) {
 61                logEvent('tengu_tool_use_granted_in_config', {
 62                  messageID: assistantMessage.message.id,
 63                  toolName: tool.name,
 64                })
 65                resolve({ result: true })
 66                return
 67              }
 68  
 69              const [description, commandPrefix] = await Promise.all([
 70                tool.description(input as never),
 71                tool === BashTool
 72                  ? getCommandSubcommandPrefix(
 73                      inputSchema.parse(input).command, // already validated upstream, so ok to parse (as opposed to safeParse)
 74                      toolUseContext.abortController.signal,
 75                    )
 76                  : Promise.resolve(null),
 77              ])
 78  
 79              if (toolUseContext.abortController.signal.aborted) {
 80                logCancelledEvent()
 81                resolveWithCancelledAndAbortAllToolCalls()
 82                return
 83              }
 84  
 85              // Does not have permissions to use tool, ask the user
 86              setToolUseConfirm({
 87                assistantMessage,
 88                tool,
 89                description,
 90                input,
 91                commandPrefix,
 92                riskScore: null,
 93                onAbort() {
 94                  logCancelledEvent()
 95                  logEvent('tengu_tool_use_rejected_in_prompt', {
 96                    messageID: assistantMessage.message.id,
 97                    toolName: tool.name,
 98                  })
 99                  resolveWithCancelledAndAbortAllToolCalls()
100                },
101                onAllow(type) {
102                  if (type === 'permanent') {
103                    logEvent('tengu_tool_use_granted_in_prompt_permanent', {
104                      messageID: assistantMessage.message.id,
105                      toolName: tool.name,
106                    })
107                  } else {
108                    logEvent('tengu_tool_use_granted_in_prompt_temporary', {
109                      messageID: assistantMessage.message.id,
110                      toolName: tool.name,
111                    })
112                  }
113                  resolve({ result: true })
114                },
115                onReject() {
116                  logEvent('tengu_tool_use_rejected_in_prompt', {
117                    messageID: assistantMessage.message.id,
118                    toolName: tool.name,
119                  })
120                  resolveWithCancelledAndAbortAllToolCalls()
121                },
122              })
123            })
124            .catch(error => {
125              if (error instanceof AbortError) {
126                logCancelledEvent()
127                resolveWithCancelledAndAbortAllToolCalls()
128              } else {
129                logError(error)
130              }
131            })
132        })
133      },
134      [setToolUseConfirm],
135    )
136  }
137  
138  export default useCanUseTool