diff --git a/packages/core/src/agents/agent-scheduler.ts b/packages/core/src/agents/agent-scheduler.ts index 983f814b0f..48282f23a4 100644 --- a/packages/core/src/agents/agent-scheduler.ts +++ b/packages/core/src/agents/agent-scheduler.ts @@ -12,6 +12,7 @@ import type { } from '../scheduler/types.js'; import type { ToolRegistry } from '../tools/tool-registry.js'; import type { EditorType } from '../utils/editor.js'; +import { PolicyDecision } from '../policy/types.js'; /** * Options for scheduling agent tools. @@ -29,6 +30,8 @@ export interface AgentSchedulingOptions { getPreferredEditor?: () => EditorType | undefined; /** Optional function to be notified when the scheduler is waiting for user confirmation. */ onWaitingForConfirmation?: (waiting: boolean) => void; + /** Optional list of tools to automatically approve for this agent. */ + allowedTools?: string[]; } /** @@ -51,6 +54,7 @@ export async function scheduleAgentTools( signal, getPreferredEditor, onWaitingForConfirmation, + allowedTools, } = options; // Create a proxy/override of the config to provide the agent-specific tool registry. @@ -59,6 +63,38 @@ export async function scheduleAgentTools( agentConfig.getToolRegistry = () => toolRegistry; agentConfig.getMessageBus = () => toolRegistry.getMessageBus(); + if (allowedTools && allowedTools.length > 0) { + const existingAllowed = config.getAllowedTools() || []; + const mergedAllowed = Array.from( + new Set([...existingAllowed, ...allowedTools]), + ); + agentConfig.getAllowedTools = () => mergedAllowed; + + const originalPolicyEngine = config.getPolicyEngine(); + // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment + const proxyPolicyEngine = Object.create(originalPolicyEngine); + proxyPolicyEngine.check = async ( + toolCall: { name: string; args?: Record }, + serverName?: string, + toolAnnotations?: Record, + ) => { + if (allowedTools.includes(toolCall.name)) { + return { + decision: PolicyDecision.ALLOW, + rule: { + toolName: toolCall.name, + decision: PolicyDecision.ALLOW, + priority: 999, + source: 'Agent Allowed Tools', + }, + }; + } + return originalPolicyEngine.check(toolCall, serverName, toolAnnotations); + }; + // eslint-disable-next-line @typescript-eslint/no-unsafe-return + agentConfig.getPolicyEngine = () => proxyPolicyEngine; + } + const scheduler = new Scheduler({ config: agentConfig, messageBus: toolRegistry.getMessageBus(), diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index dd5b78a9a6..17656fce6d 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -880,10 +880,6 @@ export class LocalAgentExecutor { taskCompleted: boolean; aborted: boolean; }> { - const allowedToolNames = new Set(this.toolRegistry.getAllToolNames()); - // Always allow the completion tool - allowedToolNames.add(TASK_COMPLETE_TOOL_NAME); - let submittedOutput: string | null = null; let taskCompleted = false; let aborted = false; @@ -1057,7 +1053,7 @@ export class LocalAgentExecutor { } // Handle standard tools - if (!allowedToolNames.has(toolName)) { + if (!this.toolRegistry.getTool(toolName)) { const error = createUnauthorizedToolError(toolName); debugLogger.warn(`[LocalAgentExecutor] Blocked call: ${error}`); @@ -1100,6 +1096,7 @@ export class LocalAgentExecutor { toolRegistry: this.toolRegistry, signal, onWaitingForConfirmation, + allowedTools: this.definition.toolConfig?.allowedTools, }, ); diff --git a/packages/core/src/agents/reflect-agent.ts b/packages/core/src/agents/reflect-agent.ts index 2229ac969e..1d6f9fcfa0 100644 --- a/packages/core/src/agents/reflect-agent.ts +++ b/packages/core/src/agents/reflect-agent.ts @@ -82,6 +82,7 @@ export const ReflectAgent = ( WRITE_FILE_TOOL_NAME, EDIT_TOOL_NAME, ], + allowedTools: [GET_SESSION_HISTORY_TOOL_NAME], }, promptConfig: { diff --git a/packages/core/src/agents/types.ts b/packages/core/src/agents/types.ts index ceac0909df..cc4a8a1b83 100644 --- a/packages/core/src/agents/types.ts +++ b/packages/core/src/agents/types.ts @@ -186,6 +186,7 @@ export interface PromptConfig { */ export interface ToolConfig { tools: Array; + allowedTools?: string[]; } /**