File size: 5,537 Bytes
064bfd6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | import { feature } from 'bun:bundle'
import type { ContentBlockParam } from '@anthropic-ai/sdk/resources/messages.mjs'
import type { PendingClassifierCheck } from '../../../types/permissions.js'
import { isAgentSwarmsEnabled } from '../../../utils/agentSwarmsEnabled.js'
import { toError } from '../../../utils/errors.js'
import { logError } from '../../../utils/log.js'
import type { PermissionDecision } from '../../../utils/permissions/PermissionResult.js'
import type { PermissionUpdate } from '../../../utils/permissions/PermissionUpdateSchema.js'
import {
createPermissionRequest,
isSwarmWorker,
sendPermissionRequestViaMailbox,
} from '../../../utils/swarm/permissionSync.js'
import { registerPermissionCallback } from '../../useSwarmPermissionPoller.js'
import type { PermissionContext } from '../PermissionContext.js'
import { createResolveOnce } from '../PermissionContext.js'
type SwarmWorkerPermissionParams = {
ctx: PermissionContext
description: string
pendingClassifierCheck?: PendingClassifierCheck | undefined
updatedInput: Record<string, unknown> | undefined
suggestions: PermissionUpdate[] | undefined
}
/**
* Handles the swarm worker permission flow.
*
* When running as a swarm worker:
* 1. Tries classifier auto-approval for bash commands
* 2. Forwards the permission request to the leader via mailbox
* 3. Registers callbacks for when the leader responds
* 4. Sets the pending indicator while waiting
*
* Returns a PermissionDecision if the classifier auto-approves,
* or a Promise that resolves when the leader responds.
* Returns null if swarms are not enabled or this is not a swarm worker,
* so the caller can fall through to interactive handling.
*/
async function handleSwarmWorkerPermission(
params: SwarmWorkerPermissionParams,
): Promise<PermissionDecision | null> {
if (!isAgentSwarmsEnabled() || !isSwarmWorker()) {
return null
}
const { ctx, description, updatedInput, suggestions } = params
// For bash commands, try classifier auto-approval before forwarding to
// the leader. Agents await the classifier result (rather than racing it
// against user interaction like the main agent).
const classifierResult = feature('BASH_CLASSIFIER')
? await ctx.tryClassifier?.(params.pendingClassifierCheck, updatedInput)
: null
if (classifierResult) {
return classifierResult
}
// Forward permission request to the leader via mailbox
try {
const clearPendingRequest = (): void =>
ctx.toolUseContext.setAppState(prev => ({
...prev,
pendingWorkerRequest: null,
}))
const decision = await new Promise<PermissionDecision>(resolve => {
const { resolve: resolveOnce, claim } = createResolveOnce(resolve)
// Create the permission request
const request = createPermissionRequest({
toolName: ctx.tool.name,
toolUseId: ctx.toolUseID,
input: ctx.input,
description,
permissionSuggestions: suggestions,
})
// Register callback BEFORE sending the request to avoid race condition
// where leader responds before callback is registered
registerPermissionCallback({
requestId: request.id,
toolUseId: ctx.toolUseID,
async onAllow(
allowedInput: Record<string, unknown> | undefined,
permissionUpdates: PermissionUpdate[],
feedback?: string,
contentBlocks?: ContentBlockParam[],
) {
if (!claim()) return // atomic check-and-mark before await
clearPendingRequest()
// Merge the updated input with the original input
const finalInput =
allowedInput && Object.keys(allowedInput).length > 0
? allowedInput
: ctx.input
resolveOnce(
await ctx.handleUserAllow(
finalInput,
permissionUpdates,
feedback,
undefined,
contentBlocks,
),
)
},
onReject(feedback?: string, contentBlocks?: ContentBlockParam[]) {
if (!claim()) return
clearPendingRequest()
ctx.logDecision({
decision: 'reject',
source: { type: 'user_reject', hasFeedback: !!feedback },
})
resolveOnce(ctx.cancelAndAbort(feedback, undefined, contentBlocks))
},
})
// Now that callback is registered, send the request to the leader
void sendPermissionRequestViaMailbox(request)
// Show visual indicator that we're waiting for leader approval
ctx.toolUseContext.setAppState(prev => ({
...prev,
pendingWorkerRequest: {
toolName: ctx.tool.name,
toolUseId: ctx.toolUseID,
description,
},
}))
// If the abort signal fires while waiting for the leader response,
// resolve the promise with a cancel decision so it does not hang.
ctx.toolUseContext.abortController.signal.addEventListener(
'abort',
() => {
if (!claim()) return
clearPendingRequest()
ctx.logCancelled()
resolveOnce(ctx.cancelAndAbort(undefined, true))
},
{ once: true },
)
})
return decision
} catch (error) {
// If swarm permission submission fails, fall back to local handling
logError(toError(error))
// Continue to local UI handling below
return null
}
}
export { handleSwarmWorkerPermission }
export type { SwarmWorkerPermissionParams }
|