| import { feature } from 'bun:bundle' |
| import type { |
| ElicitResult, |
| JSONRPCMessage, |
| } from '@modelcontextprotocol/sdk/types.js' |
| import { randomUUID } from 'crypto' |
| import type { AssistantMessage } from 'src//types/message.js' |
| import type { |
| HookInput, |
| HookJSONOutput, |
| PermissionUpdate, |
| SDKMessage, |
| SDKUserMessage, |
| } from 'src/entrypoints/agentSdkTypes.js' |
| import { SDKControlElicitationResponseSchema } from 'src/entrypoints/sdk/controlSchemas.js' |
| import type { |
| SDKControlRequest, |
| SDKControlResponse, |
| StdinMessage, |
| StdoutMessage, |
| } from 'src/entrypoints/sdk/controlTypes.js' |
| import type { CanUseToolFn } from 'src/hooks/useCanUseTool.js' |
| import type { Tool, ToolUseContext } from 'src/Tool.js' |
| import { type HookCallback, hookJSONOutputSchema } from 'src/types/hooks.js' |
| import { logForDebugging } from 'src/utils/debug.js' |
| import { logForDiagnosticsNoPII } from 'src/utils/diagLogs.js' |
| import { AbortError } from 'src/utils/errors.js' |
| import { |
| type Output as PermissionToolOutput, |
| permissionPromptToolResultToPermissionDecision, |
| outputSchema as permissionToolOutputSchema, |
| } from 'src/utils/permissions/PermissionPromptToolResultSchema.js' |
| import type { |
| PermissionDecision, |
| PermissionDecisionReason, |
| } from 'src/utils/permissions/PermissionResult.js' |
| import { hasPermissionsToUseTool } from 'src/utils/permissions/permissions.js' |
| import { writeToStdout } from 'src/utils/process.js' |
| import { jsonStringify } from 'src/utils/slowOperations.js' |
| import { z } from 'zod/v4' |
| import { notifyCommandLifecycle } from '../utils/commandLifecycle.js' |
| import { normalizeControlMessageKeys } from '../utils/controlMessageCompat.js' |
| import { executePermissionRequestHooks } from '../utils/hooks.js' |
| import { |
| applyPermissionUpdates, |
| persistPermissionUpdates, |
| } from '../utils/permissions/PermissionUpdate.js' |
| import { |
| notifySessionStateChanged, |
| type RequiresActionDetails, |
| type SessionExternalMetadata, |
| } from '../utils/sessionState.js' |
| import { jsonParse } from '../utils/slowOperations.js' |
| import { Stream } from '../utils/stream.js' |
| import { ndjsonSafeStringify } from './ndjsonSafeStringify.js' |
|
|
| |
| |
| |
| |
| |
| export const SANDBOX_NETWORK_ACCESS_TOOL_NAME = 'SandboxNetworkAccess' |
|
|
| function serializeDecisionReason( |
| reason: PermissionDecisionReason | undefined, |
| ): string | undefined { |
| if (!reason) { |
| return undefined |
| } |
|
|
| if ( |
| (feature('BASH_CLASSIFIER') || feature('TRANSCRIPT_CLASSIFIER')) && |
| reason.type === 'classifier' |
| ) { |
| return reason.reason |
| } |
| switch (reason.type) { |
| case 'rule': |
| case 'mode': |
| case 'subcommandResults': |
| case 'permissionPromptTool': |
| return undefined |
| case 'hook': |
| case 'asyncAgent': |
| case 'sandboxOverride': |
| case 'workingDir': |
| case 'safetyCheck': |
| case 'other': |
| return reason.reason |
| } |
| } |
|
|
| function buildRequiresActionDetails( |
| tool: Tool, |
| input: Record<string, unknown>, |
| toolUseID: string, |
| requestId: string, |
| ): RequiresActionDetails { |
| |
| |
| let description: string |
| try { |
| description = |
| tool.getActivityDescription?.(input) ?? |
| tool.getToolUseSummary?.(input) ?? |
| tool.userFacingName(input) |
| } catch { |
| description = tool.name |
| } |
| return { |
| tool_name: tool.name, |
| action_description: description, |
| tool_use_id: toolUseID, |
| request_id: requestId, |
| input, |
| } |
| } |
|
|
| type PendingRequest<T> = { |
| resolve: (result: T) => void |
| reject: (error: unknown) => void |
| schema?: z.Schema |
| request: SDKControlRequest |
| } |
|
|
| |
| |
| |
| |
| |
| |
| |
| const MAX_RESOLVED_TOOL_USE_IDS = 1000 |
|
|
| export class StructuredIO { |
| readonly structuredInput: AsyncGenerator<StdinMessage | SDKMessage> |
| private readonly pendingRequests = new Map<string, PendingRequest<unknown>>() |
|
|
| |
| |
| restoredWorkerState: Promise<SessionExternalMetadata | null> = |
| Promise.resolve(null) |
|
|
| private inputClosed = false |
| private unexpectedResponseCallback?: ( |
| response: SDKControlResponse, |
| ) => Promise<void> |
|
|
| |
| |
| |
| |
| |
| |
| private readonly resolvedToolUseIds = new Set<string>() |
| private prependedLines: string[] = [] |
| private onControlRequestSent?: (request: SDKControlRequest) => void |
| private onControlRequestResolved?: (requestId: string) => void |
|
|
| |
| |
| readonly outbound = new Stream<StdoutMessage>() |
|
|
| constructor( |
| private readonly input: AsyncIterable<string>, |
| private readonly replayUserMessages?: boolean, |
| ) { |
| this.input = input |
| this.structuredInput = this.read() |
| } |
|
|
| |
| |
| |
| |
| private trackResolvedToolUseId(request: SDKControlRequest): void { |
| if (request.request.subtype === 'can_use_tool') { |
| this.resolvedToolUseIds.add(request.request.tool_use_id) |
| if (this.resolvedToolUseIds.size > MAX_RESOLVED_TOOL_USE_IDS) { |
| |
| const first = this.resolvedToolUseIds.values().next().value |
| if (first !== undefined) { |
| this.resolvedToolUseIds.delete(first) |
| } |
| } |
| } |
| } |
|
|
| |
| flushInternalEvents(): Promise<void> { |
| return Promise.resolve() |
| } |
|
|
| |
| get internalEventsPending(): number { |
| return 0 |
| } |
|
|
| |
| |
| |
| |
| |
| prependUserMessage(content: string): void { |
| this.prependedLines.push( |
| jsonStringify({ |
| type: 'user', |
| session_id: '', |
| message: { role: 'user', content }, |
| parent_tool_use_id: null, |
| } satisfies SDKUserMessage) + '\n', |
| ) |
| } |
|
|
| private async *read() { |
| let content = '' |
|
|
| |
| |
| |
| |
| const splitAndProcess = async function* (this: StructuredIO) { |
| for (;;) { |
| if (this.prependedLines.length > 0) { |
| content = this.prependedLines.join('') + content |
| this.prependedLines = [] |
| } |
| const newline = content.indexOf('\n') |
| if (newline === -1) break |
| const line = content.slice(0, newline) |
| content = content.slice(newline + 1) |
| const message = await this.processLine(line) |
| if (message) { |
| logForDiagnosticsNoPII('info', 'cli_stdin_message_parsed', { |
| type: message.type, |
| }) |
| yield message |
| } |
| } |
| }.bind(this) |
|
|
| yield* splitAndProcess() |
|
|
| for await (const block of this.input) { |
| content += block |
| yield* splitAndProcess() |
| } |
| if (content) { |
| const message = await this.processLine(content) |
| if (message) { |
| yield message |
| } |
| } |
| this.inputClosed = true |
| for (const request of this.pendingRequests.values()) { |
| |
| request.reject( |
| new Error('Tool permission stream closed before response received'), |
| ) |
| } |
| } |
|
|
| getPendingPermissionRequests() { |
| return Array.from(this.pendingRequests.values()) |
| .map(entry => entry.request) |
| .filter(pr => pr.request.subtype === 'can_use_tool') |
| } |
|
|
| setUnexpectedResponseCallback( |
| callback: (response: SDKControlResponse) => Promise<void>, |
| ): void { |
| this.unexpectedResponseCallback = callback |
| } |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| injectControlResponse(response: SDKControlResponse): void { |
| const requestId = response.response?.request_id |
| if (!requestId) return |
| const request = this.pendingRequests.get(requestId) |
| if (!request) return |
| this.trackResolvedToolUseId(request.request) |
| this.pendingRequests.delete(requestId) |
| |
| void this.write({ |
| type: 'control_cancel_request', |
| request_id: requestId, |
| }) |
| if (response.response.subtype === 'error') { |
| request.reject(new Error(response.response.error)) |
| } else { |
| const result = response.response.response |
| if (request.schema) { |
| try { |
| request.resolve(request.schema.parse(result)) |
| } catch (error) { |
| request.reject(error) |
| } |
| } else { |
| request.resolve({}) |
| } |
| } |
| } |
|
|
| |
| |
| |
| |
| |
| setOnControlRequestSent( |
| callback: ((request: SDKControlRequest) => void) | undefined, |
| ): void { |
| this.onControlRequestSent = callback |
| } |
|
|
| |
| |
| |
| |
| |
| setOnControlRequestResolved( |
| callback: ((requestId: string) => void) | undefined, |
| ): void { |
| this.onControlRequestResolved = callback |
| } |
|
|
| private async processLine( |
| line: string, |
| ): Promise<StdinMessage | SDKMessage | undefined> { |
| |
| if (!line) { |
| return undefined |
| } |
| try { |
| const message = normalizeControlMessageKeys(jsonParse(line)) as |
| | StdinMessage |
| | SDKMessage |
| if (message.type === 'keep_alive') { |
| |
| return undefined |
| } |
| if (message.type === 'update_environment_variables') { |
| |
| |
| |
| |
| const keys = Object.keys(message.variables) |
| for (const [key, value] of Object.entries(message.variables)) { |
| process.env[key] = value |
| } |
| logForDebugging( |
| `[structuredIO] applied update_environment_variables: ${keys.join(', ')}`, |
| ) |
| return undefined |
| } |
| if (message.type === 'control_response') { |
| |
| |
| |
| |
| const uuid = |
| 'uuid' in message && typeof message.uuid === 'string' |
| ? message.uuid |
| : undefined |
| if (uuid) { |
| notifyCommandLifecycle(uuid, 'completed') |
| } |
| const request = this.pendingRequests.get(message.response.request_id) |
| if (!request) { |
| |
| |
| |
| |
| |
| const responsePayload = |
| message.response.subtype === 'success' |
| ? message.response.response |
| : undefined |
| const toolUseID = responsePayload?.toolUseID |
| if ( |
| typeof toolUseID === 'string' && |
| this.resolvedToolUseIds.has(toolUseID) |
| ) { |
| logForDebugging( |
| `Ignoring duplicate control_response for already-resolved toolUseID=${toolUseID} request_id=${message.response.request_id}`, |
| ) |
| return undefined |
| } |
| if (this.unexpectedResponseCallback) { |
| await this.unexpectedResponseCallback(message) |
| } |
| return undefined |
| } |
| this.trackResolvedToolUseId(request.request) |
| this.pendingRequests.delete(message.response.request_id) |
| |
| |
| if ( |
| request.request.request.subtype === 'can_use_tool' && |
| this.onControlRequestResolved |
| ) { |
| this.onControlRequestResolved(message.response.request_id) |
| } |
|
|
| if (message.response.subtype === 'error') { |
| request.reject(new Error(message.response.error)) |
| return undefined |
| } |
| const result = message.response.response |
| if (request.schema) { |
| try { |
| request.resolve(request.schema.parse(result)) |
| } catch (error) { |
| request.reject(error) |
| } |
| } else { |
| request.resolve({}) |
| } |
| |
| if (this.replayUserMessages) { |
| return message |
| } |
| return undefined |
| } |
| if ( |
| message.type !== 'user' && |
| message.type !== 'control_request' && |
| message.type !== 'assistant' && |
| message.type !== 'system' |
| ) { |
| logForDebugging(`Ignoring unknown message type: ${message.type}`, { |
| level: 'warn', |
| }) |
| return undefined |
| } |
| if (message.type === 'control_request') { |
| if (!message.request) { |
| exitWithMessage(`Error: Missing request on control_request`) |
| } |
| return message |
| } |
| if (message.type === 'assistant' || message.type === 'system') { |
| return message |
| } |
| if (message.message.role !== 'user') { |
| exitWithMessage( |
| `Error: Expected message role 'user', got '${message.message.role}'`, |
| ) |
| } |
| return message |
| } catch (error) { |
| |
| console.error(`Error parsing streaming input line: ${line}: ${error}`) |
| |
| process.exit(1) |
| } |
| } |
|
|
| async write(message: StdoutMessage): Promise<void> { |
| writeToStdout(ndjsonSafeStringify(message) + '\n') |
| } |
|
|
| private async sendRequest<Response>( |
| request: SDKControlRequest['request'], |
| schema: z.Schema, |
| signal?: AbortSignal, |
| requestId: string = randomUUID(), |
| ): Promise<Response> { |
| const message: SDKControlRequest = { |
| type: 'control_request', |
| request_id: requestId, |
| request, |
| } |
| if (this.inputClosed) { |
| throw new Error('Stream closed') |
| } |
| if (signal?.aborted) { |
| throw new Error('Request aborted') |
| } |
| this.outbound.enqueue(message) |
| if (request.subtype === 'can_use_tool' && this.onControlRequestSent) { |
| this.onControlRequestSent(message) |
| } |
| const aborted = () => { |
| this.outbound.enqueue({ |
| type: 'control_cancel_request', |
| request_id: requestId, |
| }) |
| |
| |
| const request = this.pendingRequests.get(requestId) |
| if (request) { |
| |
| |
| this.trackResolvedToolUseId(request.request) |
| request.reject(new AbortError()) |
| } |
| } |
| if (signal) { |
| signal.addEventListener('abort', aborted, { |
| once: true, |
| }) |
| } |
| try { |
| return await new Promise<Response>((resolve, reject) => { |
| this.pendingRequests.set(requestId, { |
| request: { |
| type: 'control_request', |
| request_id: requestId, |
| request, |
| }, |
| resolve: result => { |
| resolve(result as Response) |
| }, |
| reject, |
| schema, |
| }) |
| }) |
| } finally { |
| if (signal) { |
| signal.removeEventListener('abort', aborted) |
| } |
| this.pendingRequests.delete(requestId) |
| } |
| } |
|
|
| createCanUseTool( |
| onPermissionPrompt?: (details: RequiresActionDetails) => void, |
| ): CanUseToolFn { |
| return async ( |
| tool: Tool, |
| input: { [key: string]: unknown }, |
| toolUseContext: ToolUseContext, |
| assistantMessage: AssistantMessage, |
| toolUseID: string, |
| forceDecision?: PermissionDecision, |
| ): Promise<PermissionDecision> => { |
| const mainPermissionResult = |
| forceDecision ?? |
| (await hasPermissionsToUseTool( |
| tool, |
| input, |
| toolUseContext, |
| assistantMessage, |
| toolUseID, |
| )) |
| |
| if ( |
| mainPermissionResult.behavior === 'allow' || |
| mainPermissionResult.behavior === 'deny' |
| ) { |
| return mainPermissionResult |
| } |
|
|
| |
| |
| |
| |
| |
| |
|
|
| |
| const hookAbortController = new AbortController() |
| const parentSignal = toolUseContext.abortController.signal |
| |
| const onParentAbort = () => hookAbortController.abort() |
| parentSignal.addEventListener('abort', onParentAbort, { once: true }) |
|
|
| try { |
| |
| const hookPromise = executePermissionRequestHooksForSDK( |
| tool.name, |
| toolUseID, |
| input, |
| toolUseContext, |
| mainPermissionResult.suggestions, |
| ).then(decision => ({ source: 'hook' as const, decision })) |
|
|
| |
| const requestId = randomUUID() |
| onPermissionPrompt?.( |
| buildRequiresActionDetails(tool, input, toolUseID, requestId), |
| ) |
| const sdkPromise = this.sendRequest<PermissionToolOutput>( |
| { |
| subtype: 'can_use_tool', |
| tool_name: tool.name, |
| input, |
| permission_suggestions: mainPermissionResult.suggestions, |
| blocked_path: mainPermissionResult.blockedPath, |
| decision_reason: serializeDecisionReason( |
| mainPermissionResult.decisionReason, |
| ), |
| tool_use_id: toolUseID, |
| agent_id: toolUseContext.agentId, |
| }, |
| permissionToolOutputSchema(), |
| hookAbortController.signal, |
| requestId, |
| ).then(result => ({ source: 'sdk' as const, result })) |
|
|
| |
| |
| |
| const winner = await Promise.race([hookPromise, sdkPromise]) |
|
|
| if (winner.source === 'hook') { |
| if (winner.decision) { |
| |
| |
| sdkPromise.catch(() => {}) |
| hookAbortController.abort() |
| return winner.decision |
| } |
| |
| const sdkResult = await sdkPromise |
| return permissionPromptToolResultToPermissionDecision( |
| sdkResult.result, |
| tool, |
| input, |
| toolUseContext, |
| ) |
| } |
|
|
| |
| |
| return permissionPromptToolResultToPermissionDecision( |
| winner.result, |
| tool, |
| input, |
| toolUseContext, |
| ) |
| } catch (error) { |
| return permissionPromptToolResultToPermissionDecision( |
| { |
| behavior: 'deny', |
| message: `Tool permission request failed: ${error}`, |
| toolUseID, |
| }, |
| tool, |
| input, |
| toolUseContext, |
| ) |
| } finally { |
| |
| |
| if (this.getPendingPermissionRequests().length === 0) { |
| notifySessionStateChanged('running') |
| } |
| parentSignal.removeEventListener('abort', onParentAbort) |
| } |
| } |
| } |
|
|
| createHookCallback(callbackId: string, timeout?: number): HookCallback { |
| return { |
| type: 'callback', |
| timeout, |
| callback: async ( |
| input: HookInput, |
| toolUseID: string | null, |
| abort: AbortSignal | undefined, |
| ): Promise<HookJSONOutput> => { |
| try { |
| const result = await this.sendRequest<HookJSONOutput>( |
| { |
| subtype: 'hook_callback', |
| callback_id: callbackId, |
| input, |
| tool_use_id: toolUseID || undefined, |
| }, |
| hookJSONOutputSchema(), |
| abort, |
| ) |
| return result |
| } catch (error) { |
| |
| console.error(`Error in hook callback ${callbackId}:`, error) |
| return {} |
| } |
| }, |
| } |
| } |
|
|
| |
| |
| |
| async handleElicitation( |
| serverName: string, |
| message: string, |
| requestedSchema?: Record<string, unknown>, |
| signal?: AbortSignal, |
| mode?: 'form' | 'url', |
| url?: string, |
| elicitationId?: string, |
| ): Promise<ElicitResult> { |
| try { |
| const result = await this.sendRequest<ElicitResult>( |
| { |
| subtype: 'elicitation', |
| mcp_server_name: serverName, |
| message, |
| mode, |
| url, |
| elicitation_id: elicitationId, |
| requested_schema: requestedSchema, |
| }, |
| SDKControlElicitationResponseSchema(), |
| signal, |
| ) |
| return result |
| } catch { |
| return { action: 'cancel' as const } |
| } |
| } |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| createSandboxAskCallback(): (hostPattern: { |
| host: string |
| port?: number |
| }) => Promise<boolean> { |
| return async (hostPattern): Promise<boolean> => { |
| try { |
| const result = await this.sendRequest<PermissionToolOutput>( |
| { |
| subtype: 'can_use_tool', |
| tool_name: SANDBOX_NETWORK_ACCESS_TOOL_NAME, |
| input: { host: hostPattern.host }, |
| tool_use_id: randomUUID(), |
| description: `Allow network connection to ${hostPattern.host}?`, |
| }, |
| permissionToolOutputSchema(), |
| ) |
| return result.behavior === 'allow' |
| } catch { |
| |
| return false |
| } |
| } |
| } |
|
|
| |
| |
| |
| async sendMcpMessage( |
| serverName: string, |
| message: JSONRPCMessage, |
| ): Promise<JSONRPCMessage> { |
| const response = await this.sendRequest<{ mcp_response: JSONRPCMessage }>( |
| { |
| subtype: 'mcp_message', |
| server_name: serverName, |
| message, |
| }, |
| z.object({ |
| mcp_response: z.any() as z.Schema<JSONRPCMessage>, |
| }), |
| ) |
| return response.mcp_response |
| } |
| } |
|
|
| function exitWithMessage(message: string): never { |
| |
| console.error(message) |
| |
| process.exit(1) |
| } |
|
|
| |
| |
| |
| |
| async function executePermissionRequestHooksForSDK( |
| toolName: string, |
| toolUseID: string, |
| input: Record<string, unknown>, |
| toolUseContext: ToolUseContext, |
| suggestions: PermissionUpdate[] | undefined, |
| ): Promise<PermissionDecision | undefined> { |
| const appState = toolUseContext.getAppState() |
| const permissionMode = appState.toolPermissionContext.mode |
|
|
| |
| const hookGenerator = executePermissionRequestHooks( |
| toolName, |
| toolUseID, |
| input, |
| toolUseContext, |
| permissionMode, |
| suggestions, |
| toolUseContext.abortController.signal, |
| ) |
|
|
| for await (const hookResult of hookGenerator) { |
| if ( |
| hookResult.permissionRequestResult && |
| (hookResult.permissionRequestResult.behavior === 'allow' || |
| hookResult.permissionRequestResult.behavior === 'deny') |
| ) { |
| const decision = hookResult.permissionRequestResult |
| if (decision.behavior === 'allow') { |
| const finalInput = decision.updatedInput || input |
|
|
| |
| const permissionUpdates = decision.updatedPermissions ?? [] |
| if (permissionUpdates.length > 0) { |
| persistPermissionUpdates(permissionUpdates) |
| const currentAppState = toolUseContext.getAppState() |
| const updatedContext = applyPermissionUpdates( |
| currentAppState.toolPermissionContext, |
| permissionUpdates, |
| ) |
| |
| toolUseContext.setAppState(prev => { |
| if (prev.toolPermissionContext === updatedContext) return prev |
| return { ...prev, toolPermissionContext: updatedContext } |
| }) |
| } |
|
|
| return { |
| behavior: 'allow', |
| updatedInput: finalInput, |
| userModified: false, |
| decisionReason: { |
| type: 'hook', |
| hookName: 'PermissionRequest', |
| }, |
| } |
| } else { |
| |
| return { |
| behavior: 'deny', |
| message: |
| decision.message || 'Permission denied by PermissionRequest hook', |
| decisionReason: { |
| type: 'hook', |
| hookName: 'PermissionRequest', |
| }, |
| } |
| } |
| } |
| } |
|
|
| return undefined |
| } |
|
|