init
This commit is contained in:
@@ -0,0 +1,530 @@
|
||||
import type { ToolUseBlock } from '@anthropic-ai/sdk/resources/index.mjs'
|
||||
import {
|
||||
createUserMessage,
|
||||
REJECT_MESSAGE,
|
||||
withMemoryCorrectionHint,
|
||||
} from 'src/utils/messages.js'
|
||||
import type { CanUseToolFn } from '../../hooks/useCanUseTool.js'
|
||||
import { findToolByName, type Tools, type ToolUseContext } from '../../Tool.js'
|
||||
import { BASH_TOOL_NAME } from '../../tools/BashTool/toolName.js'
|
||||
import type { AssistantMessage, Message } from '../../types/message.js'
|
||||
import { createChildAbortController } from '../../utils/abortController.js'
|
||||
import { runToolUse } from './toolExecution.js'
|
||||
|
||||
type MessageUpdate = {
|
||||
message?: Message
|
||||
newContext?: ToolUseContext
|
||||
}
|
||||
|
||||
type ToolStatus = 'queued' | 'executing' | 'completed' | 'yielded'
|
||||
|
||||
type TrackedTool = {
|
||||
id: string
|
||||
block: ToolUseBlock
|
||||
assistantMessage: AssistantMessage
|
||||
status: ToolStatus
|
||||
isConcurrencySafe: boolean
|
||||
promise?: Promise<void>
|
||||
results?: Message[]
|
||||
// Progress messages are stored separately and yielded immediately
|
||||
pendingProgress: Message[]
|
||||
contextModifiers?: Array<(context: ToolUseContext) => ToolUseContext>
|
||||
}
|
||||
|
||||
/**
|
||||
* Executes tools as they stream in with concurrency control.
|
||||
* - Concurrent-safe tools can execute in parallel with other concurrent-safe tools
|
||||
* - Non-concurrent tools must execute alone (exclusive access)
|
||||
* - Results are buffered and emitted in the order tools were received
|
||||
*/
|
||||
export class StreamingToolExecutor {
|
||||
private tools: TrackedTool[] = []
|
||||
private toolUseContext: ToolUseContext
|
||||
private hasErrored = false
|
||||
private erroredToolDescription = ''
|
||||
// Child of toolUseContext.abortController. Fires when a Bash tool errors
|
||||
// so sibling subprocesses die immediately instead of running to completion.
|
||||
// Aborting this does NOT abort the parent — query.ts won't end the turn.
|
||||
private siblingAbortController: AbortController
|
||||
private discarded = false
|
||||
// Signal to wake up getRemainingResults when progress is available
|
||||
private progressAvailableResolve?: () => void
|
||||
|
||||
constructor(
|
||||
private readonly toolDefinitions: Tools,
|
||||
private readonly canUseTool: CanUseToolFn,
|
||||
toolUseContext: ToolUseContext,
|
||||
) {
|
||||
this.toolUseContext = toolUseContext
|
||||
this.siblingAbortController = createChildAbortController(
|
||||
toolUseContext.abortController,
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Discards all pending and in-progress tools. Called when streaming fallback
|
||||
* occurs and results from the failed attempt should be abandoned.
|
||||
* Queued tools won't start, and in-progress tools will receive synthetic errors.
|
||||
*/
|
||||
discard(): void {
|
||||
this.discarded = true
|
||||
}
|
||||
|
||||
/**
|
||||
* Add a tool to the execution queue. Will start executing immediately if conditions allow.
|
||||
*/
|
||||
addTool(block: ToolUseBlock, assistantMessage: AssistantMessage): void {
|
||||
const toolDefinition = findToolByName(this.toolDefinitions, block.name)
|
||||
if (!toolDefinition) {
|
||||
this.tools.push({
|
||||
id: block.id,
|
||||
block,
|
||||
assistantMessage,
|
||||
status: 'completed',
|
||||
isConcurrencySafe: true,
|
||||
pendingProgress: [],
|
||||
results: [
|
||||
createUserMessage({
|
||||
content: [
|
||||
{
|
||||
type: 'tool_result',
|
||||
content: `<tool_use_error>Error: No such tool available: ${block.name}</tool_use_error>`,
|
||||
is_error: true,
|
||||
tool_use_id: block.id,
|
||||
},
|
||||
],
|
||||
toolUseResult: `Error: No such tool available: ${block.name}`,
|
||||
sourceToolAssistantUUID: assistantMessage.uuid,
|
||||
}),
|
||||
],
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
const parsedInput = toolDefinition.inputSchema.safeParse(block.input)
|
||||
const isConcurrencySafe = parsedInput?.success
|
||||
? (() => {
|
||||
try {
|
||||
return Boolean(toolDefinition.isConcurrencySafe(parsedInput.data))
|
||||
} catch {
|
||||
return false
|
||||
}
|
||||
})()
|
||||
: false
|
||||
this.tools.push({
|
||||
id: block.id,
|
||||
block,
|
||||
assistantMessage,
|
||||
status: 'queued',
|
||||
isConcurrencySafe,
|
||||
pendingProgress: [],
|
||||
})
|
||||
|
||||
void this.processQueue()
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a tool can execute based on current concurrency state
|
||||
*/
|
||||
private canExecuteTool(isConcurrencySafe: boolean): boolean {
|
||||
const executingTools = this.tools.filter(t => t.status === 'executing')
|
||||
return (
|
||||
executingTools.length === 0 ||
|
||||
(isConcurrencySafe && executingTools.every(t => t.isConcurrencySafe))
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Process the queue, starting tools when concurrency conditions allow
|
||||
*/
|
||||
private async processQueue(): Promise<void> {
|
||||
for (const tool of this.tools) {
|
||||
if (tool.status !== 'queued') continue
|
||||
|
||||
if (this.canExecuteTool(tool.isConcurrencySafe)) {
|
||||
await this.executeTool(tool)
|
||||
} else {
|
||||
// Can't execute this tool yet, and since we need to maintain order for non-concurrent tools, stop here
|
||||
if (!tool.isConcurrencySafe) break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private createSyntheticErrorMessage(
|
||||
toolUseId: string,
|
||||
reason: 'sibling_error' | 'user_interrupted' | 'streaming_fallback',
|
||||
assistantMessage: AssistantMessage,
|
||||
): Message {
|
||||
// For user interruptions (ESC to reject), use REJECT_MESSAGE so the UI shows
|
||||
// "User rejected edit" instead of "Error editing file"
|
||||
if (reason === 'user_interrupted') {
|
||||
return createUserMessage({
|
||||
content: [
|
||||
{
|
||||
type: 'tool_result',
|
||||
content: withMemoryCorrectionHint(REJECT_MESSAGE),
|
||||
is_error: true,
|
||||
tool_use_id: toolUseId,
|
||||
},
|
||||
],
|
||||
toolUseResult: 'User rejected tool use',
|
||||
sourceToolAssistantUUID: assistantMessage.uuid,
|
||||
})
|
||||
}
|
||||
if (reason === 'streaming_fallback') {
|
||||
return createUserMessage({
|
||||
content: [
|
||||
{
|
||||
type: 'tool_result',
|
||||
content:
|
||||
'<tool_use_error>Error: Streaming fallback - tool execution discarded</tool_use_error>',
|
||||
is_error: true,
|
||||
tool_use_id: toolUseId,
|
||||
},
|
||||
],
|
||||
toolUseResult: 'Streaming fallback - tool execution discarded',
|
||||
sourceToolAssistantUUID: assistantMessage.uuid,
|
||||
})
|
||||
}
|
||||
const desc = this.erroredToolDescription
|
||||
const msg = desc
|
||||
? `Cancelled: parallel tool call ${desc} errored`
|
||||
: 'Cancelled: parallel tool call errored'
|
||||
return createUserMessage({
|
||||
content: [
|
||||
{
|
||||
type: 'tool_result',
|
||||
content: `<tool_use_error>${msg}</tool_use_error>`,
|
||||
is_error: true,
|
||||
tool_use_id: toolUseId,
|
||||
},
|
||||
],
|
||||
toolUseResult: msg,
|
||||
sourceToolAssistantUUID: assistantMessage.uuid,
|
||||
})
|
||||
}
|
||||
|
||||
/**
|
||||
* Determine why a tool should be cancelled.
|
||||
*/
|
||||
private getAbortReason(
|
||||
tool: TrackedTool,
|
||||
): 'sibling_error' | 'user_interrupted' | 'streaming_fallback' | null {
|
||||
if (this.discarded) {
|
||||
return 'streaming_fallback'
|
||||
}
|
||||
if (this.hasErrored) {
|
||||
return 'sibling_error'
|
||||
}
|
||||
if (this.toolUseContext.abortController.signal.aborted) {
|
||||
// 'interrupt' means the user typed a new message while tools were
|
||||
// running. Only cancel tools whose interruptBehavior is 'cancel';
|
||||
// 'block' tools shouldn't reach here (abort isn't fired).
|
||||
if (this.toolUseContext.abortController.signal.reason === 'interrupt') {
|
||||
return this.getToolInterruptBehavior(tool) === 'cancel'
|
||||
? 'user_interrupted'
|
||||
: null
|
||||
}
|
||||
return 'user_interrupted'
|
||||
}
|
||||
return null
|
||||
}
|
||||
|
||||
private getToolInterruptBehavior(tool: TrackedTool): 'cancel' | 'block' {
|
||||
const definition = findToolByName(this.toolDefinitions, tool.block.name)
|
||||
if (!definition?.interruptBehavior) return 'block'
|
||||
try {
|
||||
return definition.interruptBehavior()
|
||||
} catch {
|
||||
return 'block'
|
||||
}
|
||||
}
|
||||
|
||||
private getToolDescription(tool: TrackedTool): string {
|
||||
const input = tool.block.input as Record<string, unknown> | undefined
|
||||
const summary = input?.command ?? input?.file_path ?? input?.pattern ?? ''
|
||||
if (typeof summary === 'string' && summary.length > 0) {
|
||||
const truncated =
|
||||
summary.length > 40 ? summary.slice(0, 40) + '\u2026' : summary
|
||||
return `${tool.block.name}(${truncated})`
|
||||
}
|
||||
return tool.block.name
|
||||
}
|
||||
|
||||
private updateInterruptibleState(): void {
|
||||
const executing = this.tools.filter(t => t.status === 'executing')
|
||||
this.toolUseContext.setHasInterruptibleToolInProgress?.(
|
||||
executing.length > 0 &&
|
||||
executing.every(t => this.getToolInterruptBehavior(t) === 'cancel'),
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute a tool and collect its results
|
||||
*/
|
||||
private async executeTool(tool: TrackedTool): Promise<void> {
|
||||
tool.status = 'executing'
|
||||
this.toolUseContext.setInProgressToolUseIDs(prev =>
|
||||
new Set(prev).add(tool.id),
|
||||
)
|
||||
this.updateInterruptibleState()
|
||||
|
||||
const messages: Message[] = []
|
||||
const contextModifiers: Array<(context: ToolUseContext) => ToolUseContext> =
|
||||
[]
|
||||
|
||||
const collectResults = async () => {
|
||||
// If already aborted (by error or user), generate synthetic error block instead of running the tool
|
||||
const initialAbortReason = this.getAbortReason(tool)
|
||||
if (initialAbortReason) {
|
||||
messages.push(
|
||||
this.createSyntheticErrorMessage(
|
||||
tool.id,
|
||||
initialAbortReason,
|
||||
tool.assistantMessage,
|
||||
),
|
||||
)
|
||||
tool.results = messages
|
||||
tool.contextModifiers = contextModifiers
|
||||
tool.status = 'completed'
|
||||
this.updateInterruptibleState()
|
||||
return
|
||||
}
|
||||
|
||||
// Per-tool child controller. Lets siblingAbortController kill running
|
||||
// subprocesses (Bash spawns listen to this signal) when a Bash error
|
||||
// cascades. Permission-dialog rejection also aborts this controller
|
||||
// (PermissionContext.ts cancelAndAbort) — that abort must bubble up to
|
||||
// the query controller so the query loop's post-tool abort check ends
|
||||
// the turn. Without bubble-up, ExitPlanMode "clear context + auto"
|
||||
// sends REJECT_MESSAGE to the model instead of aborting (#21056 regression).
|
||||
const toolAbortController = createChildAbortController(
|
||||
this.siblingAbortController,
|
||||
)
|
||||
toolAbortController.signal.addEventListener(
|
||||
'abort',
|
||||
() => {
|
||||
if (
|
||||
toolAbortController.signal.reason !== 'sibling_error' &&
|
||||
!this.toolUseContext.abortController.signal.aborted &&
|
||||
!this.discarded
|
||||
) {
|
||||
this.toolUseContext.abortController.abort(
|
||||
toolAbortController.signal.reason,
|
||||
)
|
||||
}
|
||||
},
|
||||
{ once: true },
|
||||
)
|
||||
|
||||
const generator = runToolUse(
|
||||
tool.block,
|
||||
tool.assistantMessage,
|
||||
this.canUseTool,
|
||||
{ ...this.toolUseContext, abortController: toolAbortController },
|
||||
)
|
||||
|
||||
// Track if this specific tool has produced an error result.
|
||||
// This prevents the tool from receiving a duplicate "sibling error"
|
||||
// message when it is the one that caused the error.
|
||||
let thisToolErrored = false
|
||||
|
||||
for await (const update of generator) {
|
||||
// Check if we were aborted by a sibling tool error or user interruption.
|
||||
// Only add the synthetic error if THIS tool didn't produce the error.
|
||||
const abortReason = this.getAbortReason(tool)
|
||||
if (abortReason && !thisToolErrored) {
|
||||
messages.push(
|
||||
this.createSyntheticErrorMessage(
|
||||
tool.id,
|
||||
abortReason,
|
||||
tool.assistantMessage,
|
||||
),
|
||||
)
|
||||
break
|
||||
}
|
||||
|
||||
const isErrorResult =
|
||||
update.message.type === 'user' &&
|
||||
Array.isArray(update.message.message.content) &&
|
||||
update.message.message.content.some(
|
||||
_ => _.type === 'tool_result' && _.is_error === true,
|
||||
)
|
||||
|
||||
if (isErrorResult) {
|
||||
thisToolErrored = true
|
||||
// Only Bash errors cancel siblings. Bash commands often have implicit
|
||||
// dependency chains (e.g. mkdir fails → subsequent commands pointless).
|
||||
// Read/WebFetch/etc are independent — one failure shouldn't nuke the rest.
|
||||
if (tool.block.name === BASH_TOOL_NAME) {
|
||||
this.hasErrored = true
|
||||
this.erroredToolDescription = this.getToolDescription(tool)
|
||||
this.siblingAbortController.abort('sibling_error')
|
||||
}
|
||||
}
|
||||
|
||||
if (update.message) {
|
||||
// Progress messages go to pendingProgress for immediate yielding
|
||||
if (update.message.type === 'progress') {
|
||||
tool.pendingProgress.push(update.message)
|
||||
// Signal that progress is available
|
||||
if (this.progressAvailableResolve) {
|
||||
this.progressAvailableResolve()
|
||||
this.progressAvailableResolve = undefined
|
||||
}
|
||||
} else {
|
||||
messages.push(update.message)
|
||||
}
|
||||
}
|
||||
if (update.contextModifier) {
|
||||
contextModifiers.push(update.contextModifier.modifyContext)
|
||||
}
|
||||
}
|
||||
tool.results = messages
|
||||
tool.contextModifiers = contextModifiers
|
||||
tool.status = 'completed'
|
||||
this.updateInterruptibleState()
|
||||
|
||||
// NOTE: we currently don't support context modifiers for concurrent
|
||||
// tools. None are actively being used, but if we want to use
|
||||
// them in concurrent tools, we need to support that here.
|
||||
if (!tool.isConcurrencySafe && contextModifiers.length > 0) {
|
||||
for (const modifier of contextModifiers) {
|
||||
this.toolUseContext = modifier(this.toolUseContext)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const promise = collectResults()
|
||||
tool.promise = promise
|
||||
|
||||
// Process more queue when done
|
||||
void promise.finally(() => {
|
||||
void this.processQueue()
|
||||
})
|
||||
}
|
||||
|
||||
/**
|
||||
* Get any completed results that haven't been yielded yet (non-blocking)
|
||||
* Maintains order where necessary
|
||||
* Also yields any pending progress messages immediately
|
||||
*/
|
||||
*getCompletedResults(): Generator<MessageUpdate, void> {
|
||||
if (this.discarded) {
|
||||
return
|
||||
}
|
||||
|
||||
for (const tool of this.tools) {
|
||||
// Always yield pending progress messages immediately, regardless of tool status
|
||||
while (tool.pendingProgress.length > 0) {
|
||||
const progressMessage = tool.pendingProgress.shift()!
|
||||
yield { message: progressMessage, newContext: this.toolUseContext }
|
||||
}
|
||||
|
||||
if (tool.status === 'yielded') {
|
||||
continue
|
||||
}
|
||||
|
||||
if (tool.status === 'completed' && tool.results) {
|
||||
tool.status = 'yielded'
|
||||
|
||||
for (const message of tool.results) {
|
||||
yield { message, newContext: this.toolUseContext }
|
||||
}
|
||||
|
||||
markToolUseAsComplete(this.toolUseContext, tool.id)
|
||||
} else if (tool.status === 'executing' && !tool.isConcurrencySafe) {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if any tool has pending progress messages
|
||||
*/
|
||||
private hasPendingProgress(): boolean {
|
||||
return this.tools.some(t => t.pendingProgress.length > 0)
|
||||
}
|
||||
|
||||
/**
|
||||
* Wait for remaining tools and yield their results as they complete
|
||||
* Also yields progress messages as they become available
|
||||
*/
|
||||
async *getRemainingResults(): AsyncGenerator<MessageUpdate, void> {
|
||||
if (this.discarded) {
|
||||
return
|
||||
}
|
||||
|
||||
while (this.hasUnfinishedTools()) {
|
||||
await this.processQueue()
|
||||
|
||||
for (const result of this.getCompletedResults()) {
|
||||
yield result
|
||||
}
|
||||
|
||||
// If we still have executing tools but nothing completed, wait for any to complete
|
||||
// OR for progress to become available
|
||||
if (
|
||||
this.hasExecutingTools() &&
|
||||
!this.hasCompletedResults() &&
|
||||
!this.hasPendingProgress()
|
||||
) {
|
||||
const executingPromises = this.tools
|
||||
.filter(t => t.status === 'executing' && t.promise)
|
||||
.map(t => t.promise!)
|
||||
|
||||
// Also wait for progress to become available
|
||||
const progressPromise = new Promise<void>(resolve => {
|
||||
this.progressAvailableResolve = resolve
|
||||
})
|
||||
|
||||
if (executingPromises.length > 0) {
|
||||
await Promise.race([...executingPromises, progressPromise])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (const result of this.getCompletedResults()) {
|
||||
yield result
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if there are any completed results ready to yield
|
||||
*/
|
||||
private hasCompletedResults(): boolean {
|
||||
return this.tools.some(t => t.status === 'completed')
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if there are any tools still executing
|
||||
*/
|
||||
private hasExecutingTools(): boolean {
|
||||
return this.tools.some(t => t.status === 'executing')
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if there are any unfinished tools
|
||||
*/
|
||||
private hasUnfinishedTools(): boolean {
|
||||
return this.tools.some(t => t.status !== 'yielded')
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the current tool use context (may have been modified by context modifiers)
|
||||
*/
|
||||
getUpdatedContext(): ToolUseContext {
|
||||
return this.toolUseContext
|
||||
}
|
||||
}
|
||||
|
||||
function markToolUseAsComplete(
|
||||
toolUseContext: ToolUseContext,
|
||||
toolUseID: string,
|
||||
) {
|
||||
toolUseContext.setInProgressToolUseIDs(prev => {
|
||||
const next = new Set(prev)
|
||||
next.delete(toolUseID)
|
||||
return next
|
||||
})
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,650 @@
|
||||
import {
|
||||
type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
logEvent,
|
||||
} from 'src/services/analytics/index.js'
|
||||
import { sanitizeToolNameForAnalytics } from 'src/services/analytics/metadata.js'
|
||||
import type z from 'zod/v4'
|
||||
import type { CanUseToolFn } from '../../hooks/useCanUseTool.js'
|
||||
import type { AnyObject, Tool, ToolUseContext } from '../../Tool.js'
|
||||
import type { HookProgress } from '../../types/hooks.js'
|
||||
import type {
|
||||
AssistantMessage,
|
||||
AttachmentMessage,
|
||||
ProgressMessage,
|
||||
} from '../../types/message.js'
|
||||
import type { PermissionDecision } from '../../types/permissions.js'
|
||||
import { createAttachmentMessage } from '../../utils/attachments.js'
|
||||
import { logForDebugging } from '../../utils/debug.js'
|
||||
import {
|
||||
executePostToolHooks,
|
||||
executePostToolUseFailureHooks,
|
||||
executePreToolHooks,
|
||||
getPreToolHookBlockingMessage,
|
||||
} from '../../utils/hooks.js'
|
||||
import { logError } from '../../utils/log.js'
|
||||
import {
|
||||
getRuleBehaviorDescription,
|
||||
type PermissionDecisionReason,
|
||||
type PermissionResult,
|
||||
} from '../../utils/permissions/PermissionResult.js'
|
||||
import { checkRuleBasedPermissions } from '../../utils/permissions/permissions.js'
|
||||
import { formatError } from '../../utils/toolErrors.js'
|
||||
import { isMcpTool } from '../mcp/utils.js'
|
||||
import type { McpServerType, MessageUpdateLazy } from './toolExecution.js'
|
||||
|
||||
export type PostToolUseHooksResult<Output> =
|
||||
| MessageUpdateLazy<AttachmentMessage | ProgressMessage<HookProgress>>
|
||||
| { updatedMCPToolOutput: Output }
|
||||
|
||||
export async function* runPostToolUseHooks<Input extends AnyObject, Output>(
|
||||
toolUseContext: ToolUseContext,
|
||||
tool: Tool<Input, Output>,
|
||||
toolUseID: string,
|
||||
messageId: string,
|
||||
toolInput: Record<string, unknown>,
|
||||
toolResponse: Output,
|
||||
requestId: string | undefined,
|
||||
mcpServerType: McpServerType,
|
||||
mcpServerBaseUrl: string | undefined,
|
||||
): AsyncGenerator<PostToolUseHooksResult<Output>> {
|
||||
const postToolStartTime = Date.now()
|
||||
try {
|
||||
const appState = toolUseContext.getAppState()
|
||||
const permissionMode = appState.toolPermissionContext.mode
|
||||
|
||||
let toolOutput = toolResponse
|
||||
for await (const result of executePostToolHooks(
|
||||
tool.name,
|
||||
toolUseID,
|
||||
toolInput,
|
||||
toolOutput,
|
||||
toolUseContext,
|
||||
permissionMode,
|
||||
toolUseContext.abortController.signal,
|
||||
)) {
|
||||
try {
|
||||
// Check if we were aborted during hook execution
|
||||
// IMPORTANT: We emit a cancelled event per hook
|
||||
if (
|
||||
result.message?.type === 'attachment' &&
|
||||
result.message.attachment.type === 'hook_cancelled'
|
||||
) {
|
||||
logEvent('tengu_post_tool_hooks_cancelled', {
|
||||
toolName: sanitizeToolNameForAnalytics(tool.name),
|
||||
|
||||
queryChainId: toolUseContext.queryTracking
|
||||
?.chainId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
queryDepth: toolUseContext.queryTracking?.depth,
|
||||
})
|
||||
yield {
|
||||
message: createAttachmentMessage({
|
||||
type: 'hook_cancelled',
|
||||
hookName: `PostToolUse:${tool.name}`,
|
||||
toolUseID,
|
||||
hookEvent: 'PostToolUse',
|
||||
}),
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// For JSON {decision:"block"} hooks, executeHooks yields two results:
|
||||
// {blockingError} and {message: hook_blocking_error attachment}. The
|
||||
// blockingError path below creates that same attachment, so skip it
|
||||
// here to avoid displaying the block reason twice (#31301). The
|
||||
// exit-code-2 path only yields {blockingError}, so it's unaffected.
|
||||
if (
|
||||
result.message &&
|
||||
!(
|
||||
result.message.type === 'attachment' &&
|
||||
result.message.attachment.type === 'hook_blocking_error'
|
||||
)
|
||||
) {
|
||||
yield { message: result.message }
|
||||
}
|
||||
|
||||
if (result.blockingError) {
|
||||
yield {
|
||||
message: createAttachmentMessage({
|
||||
type: 'hook_blocking_error',
|
||||
hookName: `PostToolUse:${tool.name}`,
|
||||
toolUseID: toolUseID,
|
||||
hookEvent: 'PostToolUse',
|
||||
blockingError: result.blockingError,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
// If hook indicated to prevent continuation, yield a stop reason message
|
||||
if (result.preventContinuation) {
|
||||
yield {
|
||||
message: createAttachmentMessage({
|
||||
type: 'hook_stopped_continuation',
|
||||
message:
|
||||
result.stopReason || 'Execution stopped by PostToolUse hook',
|
||||
hookName: `PostToolUse:${tool.name}`,
|
||||
toolUseID: toolUseID,
|
||||
hookEvent: 'PostToolUse',
|
||||
}),
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// If hooks provided additional context, add it as a message
|
||||
if (result.additionalContexts && result.additionalContexts.length > 0) {
|
||||
yield {
|
||||
message: createAttachmentMessage({
|
||||
type: 'hook_additional_context',
|
||||
content: result.additionalContexts,
|
||||
hookName: `PostToolUse:${tool.name}`,
|
||||
toolUseID: toolUseID,
|
||||
hookEvent: 'PostToolUse',
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
// If hooks provided updatedMCPToolOutput, yield it if this is an MCP tool
|
||||
if (result.updatedMCPToolOutput && isMcpTool(tool)) {
|
||||
toolOutput = result.updatedMCPToolOutput as Output
|
||||
yield {
|
||||
updatedMCPToolOutput: toolOutput,
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
const postToolDurationMs = Date.now() - postToolStartTime
|
||||
logEvent('tengu_post_tool_hook_error', {
|
||||
messageID:
|
||||
messageId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
toolName: sanitizeToolNameForAnalytics(tool.name),
|
||||
isMcp: tool.isMcp ?? false,
|
||||
duration: postToolDurationMs,
|
||||
|
||||
queryChainId: toolUseContext.queryTracking
|
||||
?.chainId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
queryDepth: toolUseContext.queryTracking?.depth,
|
||||
...(mcpServerType
|
||||
? {
|
||||
mcpServerType:
|
||||
mcpServerType as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
}
|
||||
: {}),
|
||||
...(requestId
|
||||
? {
|
||||
requestId:
|
||||
requestId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
}
|
||||
: {}),
|
||||
})
|
||||
yield {
|
||||
message: createAttachmentMessage({
|
||||
type: 'hook_error_during_execution',
|
||||
content: formatError(error),
|
||||
hookName: `PostToolUse:${tool.name}`,
|
||||
toolUseID: toolUseID,
|
||||
hookEvent: 'PostToolUse',
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
logError(error)
|
||||
}
|
||||
}
|
||||
|
||||
export async function* runPostToolUseFailureHooks<Input extends AnyObject>(
|
||||
toolUseContext: ToolUseContext,
|
||||
tool: Tool<Input, unknown>,
|
||||
toolUseID: string,
|
||||
messageId: string,
|
||||
processedInput: z.infer<Input>,
|
||||
error: string,
|
||||
isInterrupt: boolean | undefined,
|
||||
requestId: string | undefined,
|
||||
mcpServerType: McpServerType,
|
||||
mcpServerBaseUrl: string | undefined,
|
||||
): AsyncGenerator<
|
||||
MessageUpdateLazy<AttachmentMessage | ProgressMessage<HookProgress>>
|
||||
> {
|
||||
const postToolStartTime = Date.now()
|
||||
try {
|
||||
const appState = toolUseContext.getAppState()
|
||||
const permissionMode = appState.toolPermissionContext.mode
|
||||
|
||||
for await (const result of executePostToolUseFailureHooks(
|
||||
tool.name,
|
||||
toolUseID,
|
||||
processedInput,
|
||||
error,
|
||||
toolUseContext,
|
||||
isInterrupt,
|
||||
permissionMode,
|
||||
toolUseContext.abortController.signal,
|
||||
)) {
|
||||
try {
|
||||
// Check if we were aborted during hook execution
|
||||
if (
|
||||
result.message?.type === 'attachment' &&
|
||||
result.message.attachment.type === 'hook_cancelled'
|
||||
) {
|
||||
logEvent('tengu_post_tool_failure_hooks_cancelled', {
|
||||
toolName: sanitizeToolNameForAnalytics(tool.name),
|
||||
queryChainId: toolUseContext.queryTracking
|
||||
?.chainId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
queryDepth: toolUseContext.queryTracking?.depth,
|
||||
})
|
||||
yield {
|
||||
message: createAttachmentMessage({
|
||||
type: 'hook_cancelled',
|
||||
hookName: `PostToolUseFailure:${tool.name}`,
|
||||
toolUseID,
|
||||
hookEvent: 'PostToolUseFailure',
|
||||
}),
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// Skip hook_blocking_error in result.message — blockingError path
|
||||
// below creates the same attachment (see #31301 / PostToolUse above).
|
||||
if (
|
||||
result.message &&
|
||||
!(
|
||||
result.message.type === 'attachment' &&
|
||||
result.message.attachment.type === 'hook_blocking_error'
|
||||
)
|
||||
) {
|
||||
yield { message: result.message }
|
||||
}
|
||||
|
||||
if (result.blockingError) {
|
||||
yield {
|
||||
message: createAttachmentMessage({
|
||||
type: 'hook_blocking_error',
|
||||
hookName: `PostToolUseFailure:${tool.name}`,
|
||||
toolUseID: toolUseID,
|
||||
hookEvent: 'PostToolUseFailure',
|
||||
blockingError: result.blockingError,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
// If hooks provided additional context, add it as a message
|
||||
if (result.additionalContexts && result.additionalContexts.length > 0) {
|
||||
yield {
|
||||
message: createAttachmentMessage({
|
||||
type: 'hook_additional_context',
|
||||
content: result.additionalContexts,
|
||||
hookName: `PostToolUseFailure:${tool.name}`,
|
||||
toolUseID: toolUseID,
|
||||
hookEvent: 'PostToolUseFailure',
|
||||
}),
|
||||
}
|
||||
}
|
||||
} catch (hookError) {
|
||||
const postToolDurationMs = Date.now() - postToolStartTime
|
||||
logEvent('tengu_post_tool_failure_hook_error', {
|
||||
messageID:
|
||||
messageId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
toolName: sanitizeToolNameForAnalytics(tool.name),
|
||||
isMcp: tool.isMcp ?? false,
|
||||
duration: postToolDurationMs,
|
||||
queryChainId: toolUseContext.queryTracking
|
||||
?.chainId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
queryDepth: toolUseContext.queryTracking?.depth,
|
||||
...(mcpServerType
|
||||
? {
|
||||
mcpServerType:
|
||||
mcpServerType as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
}
|
||||
: {}),
|
||||
...(requestId
|
||||
? {
|
||||
requestId:
|
||||
requestId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
}
|
||||
: {}),
|
||||
})
|
||||
yield {
|
||||
message: createAttachmentMessage({
|
||||
type: 'hook_error_during_execution',
|
||||
content: formatError(hookError),
|
||||
hookName: `PostToolUseFailure:${tool.name}`,
|
||||
toolUseID: toolUseID,
|
||||
hookEvent: 'PostToolUseFailure',
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
} catch (outerError) {
|
||||
logError(outerError)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve a PreToolUse hook's permission result into a final PermissionDecision.
|
||||
*
|
||||
* Encapsulates the invariant that hook 'allow' does NOT bypass settings.json
|
||||
* deny/ask rules — checkRuleBasedPermissions still applies (inc-4788 analog).
|
||||
* Also handles the requiresUserInteraction/requireCanUseTool guards and the
|
||||
* 'ask' forceDecision passthrough.
|
||||
*
|
||||
* Shared by toolExecution.ts (main query loop) and REPLTool/toolWrappers.ts
|
||||
* (REPL inner calls) so the permission semantics stay in lockstep.
|
||||
*/
|
||||
export async function resolveHookPermissionDecision(
|
||||
hookPermissionResult: PermissionResult | undefined,
|
||||
tool: Tool,
|
||||
input: Record<string, unknown>,
|
||||
toolUseContext: ToolUseContext,
|
||||
canUseTool: CanUseToolFn,
|
||||
assistantMessage: AssistantMessage,
|
||||
toolUseID: string,
|
||||
): Promise<{
|
||||
decision: PermissionDecision
|
||||
input: Record<string, unknown>
|
||||
}> {
|
||||
const requiresInteraction = tool.requiresUserInteraction?.()
|
||||
const requireCanUseTool = toolUseContext.requireCanUseTool
|
||||
|
||||
if (hookPermissionResult?.behavior === 'allow') {
|
||||
const hookInput = hookPermissionResult.updatedInput ?? input
|
||||
|
||||
// Hook provided updatedInput for an interactive tool — the hook IS the
|
||||
// user interaction (e.g. headless wrapper that collected AskUserQuestion
|
||||
// answers). Treat as non-interactive for the rule-check path.
|
||||
const interactionSatisfied =
|
||||
requiresInteraction && hookPermissionResult.updatedInput !== undefined
|
||||
|
||||
if ((requiresInteraction && !interactionSatisfied) || requireCanUseTool) {
|
||||
logForDebugging(
|
||||
`Hook approved tool use for ${tool.name}, but canUseTool is required`,
|
||||
)
|
||||
return {
|
||||
decision: await canUseTool(
|
||||
tool,
|
||||
hookInput,
|
||||
toolUseContext,
|
||||
assistantMessage,
|
||||
toolUseID,
|
||||
),
|
||||
input: hookInput,
|
||||
}
|
||||
}
|
||||
|
||||
// Hook allow skips the interactive prompt, but deny/ask rules still apply.
|
||||
const ruleCheck = await checkRuleBasedPermissions(
|
||||
tool,
|
||||
hookInput,
|
||||
toolUseContext,
|
||||
)
|
||||
if (ruleCheck === null) {
|
||||
logForDebugging(
|
||||
interactionSatisfied
|
||||
? `Hook satisfied user interaction for ${tool.name} via updatedInput`
|
||||
: `Hook approved tool use for ${tool.name}, bypassing permission prompt`,
|
||||
)
|
||||
return { decision: hookPermissionResult, input: hookInput }
|
||||
}
|
||||
if (ruleCheck.behavior === 'deny') {
|
||||
logForDebugging(
|
||||
`Hook approved tool use for ${tool.name}, but deny rule overrides: ${ruleCheck.message}`,
|
||||
)
|
||||
return { decision: ruleCheck, input: hookInput }
|
||||
}
|
||||
// ask rule — dialog required despite hook approval
|
||||
logForDebugging(
|
||||
`Hook approved tool use for ${tool.name}, but ask rule requires prompt`,
|
||||
)
|
||||
return {
|
||||
decision: await canUseTool(
|
||||
tool,
|
||||
hookInput,
|
||||
toolUseContext,
|
||||
assistantMessage,
|
||||
toolUseID,
|
||||
),
|
||||
input: hookInput,
|
||||
}
|
||||
}
|
||||
|
||||
if (hookPermissionResult?.behavior === 'deny') {
|
||||
logForDebugging(`Hook denied tool use for ${tool.name}`)
|
||||
return { decision: hookPermissionResult, input }
|
||||
}
|
||||
|
||||
// No hook decision or 'ask' — normal permission flow, possibly with
|
||||
// forceDecision so the dialog shows the hook's ask message.
|
||||
const forceDecision =
|
||||
hookPermissionResult?.behavior === 'ask' ? hookPermissionResult : undefined
|
||||
const askInput =
|
||||
hookPermissionResult?.behavior === 'ask' &&
|
||||
hookPermissionResult.updatedInput
|
||||
? hookPermissionResult.updatedInput
|
||||
: input
|
||||
return {
|
||||
decision: await canUseTool(
|
||||
tool,
|
||||
askInput,
|
||||
toolUseContext,
|
||||
assistantMessage,
|
||||
toolUseID,
|
||||
forceDecision,
|
||||
),
|
||||
input: askInput,
|
||||
}
|
||||
}
|
||||
|
||||
export async function* runPreToolUseHooks(
|
||||
toolUseContext: ToolUseContext,
|
||||
tool: Tool,
|
||||
processedInput: Record<string, unknown>,
|
||||
toolUseID: string,
|
||||
messageId: string,
|
||||
requestId: string | undefined,
|
||||
mcpServerType: McpServerType,
|
||||
mcpServerBaseUrl: string | undefined,
|
||||
): AsyncGenerator<
|
||||
| {
|
||||
type: 'message'
|
||||
message: MessageUpdateLazy<
|
||||
AttachmentMessage | ProgressMessage<HookProgress>
|
||||
>
|
||||
}
|
||||
| { type: 'hookPermissionResult'; hookPermissionResult: PermissionResult }
|
||||
| { type: 'hookUpdatedInput'; updatedInput: Record<string, unknown> }
|
||||
| { type: 'preventContinuation'; shouldPreventContinuation: boolean }
|
||||
| { type: 'stopReason'; stopReason: string }
|
||||
| {
|
||||
type: 'additionalContext'
|
||||
message: MessageUpdateLazy<AttachmentMessage>
|
||||
}
|
||||
// stop execution
|
||||
| { type: 'stop' }
|
||||
> {
|
||||
const hookStartTime = Date.now()
|
||||
try {
|
||||
const appState = toolUseContext.getAppState()
|
||||
|
||||
for await (const result of executePreToolHooks(
|
||||
tool.name,
|
||||
toolUseID,
|
||||
processedInput,
|
||||
toolUseContext,
|
||||
appState.toolPermissionContext.mode,
|
||||
toolUseContext.abortController.signal,
|
||||
undefined, // timeoutMs - use default
|
||||
toolUseContext.requestPrompt,
|
||||
tool.getToolUseSummary?.(processedInput),
|
||||
)) {
|
||||
try {
|
||||
if (result.message) {
|
||||
yield { type: 'message', message: { message: result.message } }
|
||||
}
|
||||
if (result.blockingError) {
|
||||
const denialMessage = getPreToolHookBlockingMessage(
|
||||
`PreToolUse:${tool.name}`,
|
||||
result.blockingError,
|
||||
)
|
||||
yield {
|
||||
type: 'hookPermissionResult',
|
||||
hookPermissionResult: {
|
||||
behavior: 'deny',
|
||||
message: denialMessage,
|
||||
decisionReason: {
|
||||
type: 'hook',
|
||||
hookName: `PreToolUse:${tool.name}`,
|
||||
reason: denialMessage,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
// Check if hook wants to prevent continuation
|
||||
if (result.preventContinuation) {
|
||||
yield {
|
||||
type: 'preventContinuation',
|
||||
shouldPreventContinuation: true,
|
||||
}
|
||||
if (result.stopReason) {
|
||||
yield { type: 'stopReason', stopReason: result.stopReason }
|
||||
}
|
||||
}
|
||||
// Check for hook-defined permission behavior
|
||||
if (result.permissionBehavior !== undefined) {
|
||||
logForDebugging(
|
||||
`Hook result has permissionBehavior=${result.permissionBehavior}`,
|
||||
)
|
||||
const decisionReason: PermissionDecisionReason = {
|
||||
type: 'hook',
|
||||
hookName: `PreToolUse:${tool.name}`,
|
||||
hookSource: result.hookSource,
|
||||
reason: result.hookPermissionDecisionReason,
|
||||
}
|
||||
if (result.permissionBehavior === 'allow') {
|
||||
yield {
|
||||
type: 'hookPermissionResult',
|
||||
hookPermissionResult: {
|
||||
behavior: 'allow',
|
||||
updatedInput: result.updatedInput,
|
||||
decisionReason,
|
||||
},
|
||||
}
|
||||
} else if (result.permissionBehavior === 'ask') {
|
||||
yield {
|
||||
type: 'hookPermissionResult',
|
||||
hookPermissionResult: {
|
||||
behavior: 'ask',
|
||||
updatedInput: result.updatedInput,
|
||||
message:
|
||||
result.hookPermissionDecisionReason ||
|
||||
`Hook PreToolUse:${tool.name} ${getRuleBehaviorDescription(result.permissionBehavior)} this tool`,
|
||||
decisionReason,
|
||||
},
|
||||
}
|
||||
} else {
|
||||
// deny - updatedInput is irrelevant since tool won't run
|
||||
yield {
|
||||
type: 'hookPermissionResult',
|
||||
hookPermissionResult: {
|
||||
behavior: result.permissionBehavior,
|
||||
message:
|
||||
result.hookPermissionDecisionReason ||
|
||||
`Hook PreToolUse:${tool.name} ${getRuleBehaviorDescription(result.permissionBehavior)} this tool`,
|
||||
decisionReason,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Yield updatedInput for passthrough case (no permission decision)
|
||||
// This allows hooks to modify input while letting normal permission flow continue
|
||||
if (result.updatedInput && result.permissionBehavior === undefined) {
|
||||
yield {
|
||||
type: 'hookUpdatedInput',
|
||||
updatedInput: result.updatedInput,
|
||||
}
|
||||
}
|
||||
|
||||
// If hooks provided additional context, add it as a message
|
||||
if (result.additionalContexts && result.additionalContexts.length > 0) {
|
||||
yield {
|
||||
type: 'additionalContext',
|
||||
message: {
|
||||
message: createAttachmentMessage({
|
||||
type: 'hook_additional_context',
|
||||
content: result.additionalContexts,
|
||||
hookName: `PreToolUse:${tool.name}`,
|
||||
toolUseID,
|
||||
hookEvent: 'PreToolUse',
|
||||
}),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Check if we were aborted during hook execution
|
||||
if (toolUseContext.abortController.signal.aborted) {
|
||||
logEvent('tengu_pre_tool_hooks_cancelled', {
|
||||
toolName: sanitizeToolNameForAnalytics(tool.name),
|
||||
|
||||
queryChainId: toolUseContext.queryTracking
|
||||
?.chainId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
queryDepth: toolUseContext.queryTracking?.depth,
|
||||
})
|
||||
yield {
|
||||
type: 'message',
|
||||
message: {
|
||||
message: createAttachmentMessage({
|
||||
type: 'hook_cancelled',
|
||||
hookName: `PreToolUse:${tool.name}`,
|
||||
toolUseID,
|
||||
hookEvent: 'PreToolUse',
|
||||
}),
|
||||
},
|
||||
}
|
||||
yield { type: 'stop' }
|
||||
return
|
||||
}
|
||||
} catch (error) {
|
||||
logError(error)
|
||||
const durationMs = Date.now() - hookStartTime
|
||||
logEvent('tengu_pre_tool_hook_error', {
|
||||
messageID:
|
||||
messageId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
toolName: sanitizeToolNameForAnalytics(tool.name),
|
||||
isMcp: tool.isMcp ?? false,
|
||||
duration: durationMs,
|
||||
|
||||
queryChainId: toolUseContext.queryTracking
|
||||
?.chainId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
queryDepth: toolUseContext.queryTracking?.depth,
|
||||
...(mcpServerType
|
||||
? {
|
||||
mcpServerType:
|
||||
mcpServerType as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
}
|
||||
: {}),
|
||||
...(requestId
|
||||
? {
|
||||
requestId:
|
||||
requestId as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
|
||||
}
|
||||
: {}),
|
||||
})
|
||||
yield {
|
||||
type: 'message',
|
||||
message: {
|
||||
message: createAttachmentMessage({
|
||||
type: 'hook_error_during_execution',
|
||||
content: formatError(error),
|
||||
hookName: `PreToolUse:${tool.name}`,
|
||||
toolUseID: toolUseID,
|
||||
hookEvent: 'PreToolUse',
|
||||
}),
|
||||
},
|
||||
}
|
||||
yield { type: 'stop' }
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
logError(error)
|
||||
yield { type: 'stop' }
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,188 @@
|
||||
import type { ToolUseBlock } from '@anthropic-ai/sdk/resources/index.mjs'
|
||||
import type { CanUseToolFn } from '../../hooks/useCanUseTool.js'
|
||||
import { findToolByName, type ToolUseContext } from '../../Tool.js'
|
||||
import type { AssistantMessage, Message } from '../../types/message.js'
|
||||
import { all } from '../../utils/generators.js'
|
||||
import { type MessageUpdateLazy, runToolUse } from './toolExecution.js'
|
||||
|
||||
function getMaxToolUseConcurrency(): number {
|
||||
return (
|
||||
parseInt(process.env.CLAUDE_CODE_MAX_TOOL_USE_CONCURRENCY || '', 10) || 10
|
||||
)
|
||||
}
|
||||
|
||||
export type MessageUpdate = {
|
||||
message?: Message
|
||||
newContext: ToolUseContext
|
||||
}
|
||||
|
||||
export async function* runTools(
|
||||
toolUseMessages: ToolUseBlock[],
|
||||
assistantMessages: AssistantMessage[],
|
||||
canUseTool: CanUseToolFn,
|
||||
toolUseContext: ToolUseContext,
|
||||
): AsyncGenerator<MessageUpdate, void> {
|
||||
let currentContext = toolUseContext
|
||||
for (const { isConcurrencySafe, blocks } of partitionToolCalls(
|
||||
toolUseMessages,
|
||||
currentContext,
|
||||
)) {
|
||||
if (isConcurrencySafe) {
|
||||
const queuedContextModifiers: Record<
|
||||
string,
|
||||
((context: ToolUseContext) => ToolUseContext)[]
|
||||
> = {}
|
||||
// Run read-only batch concurrently
|
||||
for await (const update of runToolsConcurrently(
|
||||
blocks,
|
||||
assistantMessages,
|
||||
canUseTool,
|
||||
currentContext,
|
||||
)) {
|
||||
if (update.contextModifier) {
|
||||
const { toolUseID, modifyContext } = update.contextModifier
|
||||
if (!queuedContextModifiers[toolUseID]) {
|
||||
queuedContextModifiers[toolUseID] = []
|
||||
}
|
||||
queuedContextModifiers[toolUseID].push(modifyContext)
|
||||
}
|
||||
yield {
|
||||
message: update.message,
|
||||
newContext: currentContext,
|
||||
}
|
||||
}
|
||||
for (const block of blocks) {
|
||||
const modifiers = queuedContextModifiers[block.id]
|
||||
if (!modifiers) {
|
||||
continue
|
||||
}
|
||||
for (const modifier of modifiers) {
|
||||
currentContext = modifier(currentContext)
|
||||
}
|
||||
}
|
||||
yield { newContext: currentContext }
|
||||
} else {
|
||||
// Run non-read-only batch serially
|
||||
for await (const update of runToolsSerially(
|
||||
blocks,
|
||||
assistantMessages,
|
||||
canUseTool,
|
||||
currentContext,
|
||||
)) {
|
||||
if (update.newContext) {
|
||||
currentContext = update.newContext
|
||||
}
|
||||
yield {
|
||||
message: update.message,
|
||||
newContext: currentContext,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type Batch = { isConcurrencySafe: boolean; blocks: ToolUseBlock[] }
|
||||
|
||||
/**
|
||||
* Partition tool calls into batches where each batch is either:
|
||||
* 1. A single non-read-only tool, or
|
||||
* 2. Multiple consecutive read-only tools
|
||||
*/
|
||||
function partitionToolCalls(
|
||||
toolUseMessages: ToolUseBlock[],
|
||||
toolUseContext: ToolUseContext,
|
||||
): Batch[] {
|
||||
return toolUseMessages.reduce((acc: Batch[], toolUse) => {
|
||||
const tool = findToolByName(toolUseContext.options.tools, toolUse.name)
|
||||
const parsedInput = tool?.inputSchema.safeParse(toolUse.input)
|
||||
const isConcurrencySafe = parsedInput?.success
|
||||
? (() => {
|
||||
try {
|
||||
return Boolean(tool?.isConcurrencySafe(parsedInput.data))
|
||||
} catch {
|
||||
// If isConcurrencySafe throws (e.g., due to shell-quote parse failure),
|
||||
// treat as not concurrency-safe to be conservative
|
||||
return false
|
||||
}
|
||||
})()
|
||||
: false
|
||||
if (isConcurrencySafe && acc[acc.length - 1]?.isConcurrencySafe) {
|
||||
acc[acc.length - 1]!.blocks.push(toolUse)
|
||||
} else {
|
||||
acc.push({ isConcurrencySafe, blocks: [toolUse] })
|
||||
}
|
||||
return acc
|
||||
}, [])
|
||||
}
|
||||
|
||||
async function* runToolsSerially(
|
||||
toolUseMessages: ToolUseBlock[],
|
||||
assistantMessages: AssistantMessage[],
|
||||
canUseTool: CanUseToolFn,
|
||||
toolUseContext: ToolUseContext,
|
||||
): AsyncGenerator<MessageUpdate, void> {
|
||||
let currentContext = toolUseContext
|
||||
|
||||
for (const toolUse of toolUseMessages) {
|
||||
toolUseContext.setInProgressToolUseIDs(prev =>
|
||||
new Set(prev).add(toolUse.id),
|
||||
)
|
||||
for await (const update of runToolUse(
|
||||
toolUse,
|
||||
assistantMessages.find(_ =>
|
||||
_.message.content.some(
|
||||
_ => _.type === 'tool_use' && _.id === toolUse.id,
|
||||
),
|
||||
)!,
|
||||
canUseTool,
|
||||
currentContext,
|
||||
)) {
|
||||
if (update.contextModifier) {
|
||||
currentContext = update.contextModifier.modifyContext(currentContext)
|
||||
}
|
||||
yield {
|
||||
message: update.message,
|
||||
newContext: currentContext,
|
||||
}
|
||||
}
|
||||
markToolUseAsComplete(toolUseContext, toolUse.id)
|
||||
}
|
||||
}
|
||||
|
||||
async function* runToolsConcurrently(
|
||||
toolUseMessages: ToolUseBlock[],
|
||||
assistantMessages: AssistantMessage[],
|
||||
canUseTool: CanUseToolFn,
|
||||
toolUseContext: ToolUseContext,
|
||||
): AsyncGenerator<MessageUpdateLazy, void> {
|
||||
yield* all(
|
||||
toolUseMessages.map(async function* (toolUse) {
|
||||
toolUseContext.setInProgressToolUseIDs(prev =>
|
||||
new Set(prev).add(toolUse.id),
|
||||
)
|
||||
yield* runToolUse(
|
||||
toolUse,
|
||||
assistantMessages.find(_ =>
|
||||
_.message.content.some(
|
||||
_ => _.type === 'tool_use' && _.id === toolUse.id,
|
||||
),
|
||||
)!,
|
||||
canUseTool,
|
||||
toolUseContext,
|
||||
)
|
||||
markToolUseAsComplete(toolUseContext, toolUse.id)
|
||||
}),
|
||||
getMaxToolUseConcurrency(),
|
||||
)
|
||||
}
|
||||
|
||||
function markToolUseAsComplete(
|
||||
toolUseContext: ToolUseContext,
|
||||
toolUseID: string,
|
||||
) {
|
||||
toolUseContext.setInProgressToolUseIDs(prev => {
|
||||
const next = new Set(prev)
|
||||
next.delete(toolUseID)
|
||||
return next
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user