diff --git a/packages/agent-core-v2/src/features/flow/flowService.ts b/packages/agent-core-v2/src/features/flow/flowService.ts index 3980bc0a2..7084d8f8f 100644 --- a/packages/agent-core-v2/src/features/flow/flowService.ts +++ b/packages/agent-core-v2/src/features/flow/flowService.ts @@ -12,6 +12,7 @@ import { ContextUndone } from '#/agent/undo/undoService'; import { AgentStatusUpdated, type AgentFlowRunStatus } from '#/agent/usage/usageEvents'; import { IAgentStateService } from '#/agent/state/agentState'; import { IAgentToolApprovalService } from '#/agent/toolApproval/toolApproval'; +import { IAgentToolRegistryService } from '#/agent/toolRegistry/toolRegistry'; import { denyToolExecution } from '#/agent/toolExecutor/beforeToolExecuteEvent'; import { IAgentToolExecutorService } from '#/agent/toolExecutor/toolExecutor'; import { IEventBus } from '#/app/event/eventBus'; @@ -67,6 +68,7 @@ export class AgentFlowService extends Disposable implements IAgentFlowService { @IEventDispatcher private readonly dispatcher: IEventDispatcher, @IAgentStateService private readonly agentState: IAgentStateService, @IAgentToolApprovalService private readonly toolApproval: IAgentToolApprovalService, + @IAgentToolRegistryService private readonly toolRegistry: IAgentToolRegistryService, @IAgentToolExecutorService toolExecutor: IAgentToolExecutorService, @IAgentPermissionModeService private readonly modeService: IAgentPermissionModeService, @IFlagService private readonly flags: IFlagService, @@ -91,8 +93,8 @@ export class AgentFlowService extends Disposable implements IAgentFlowService { this._register( toolExecutor.onBeforeExecuteTool((event) => { if (!this.flags.enabled(FLOW_FLAG_ID)) return; - if (!FLOW_TOOL_NAMES.has(event.toolCall.name)) return; - const firstFlowCall = event.toolCalls.find((call) => FLOW_TOOL_NAMES.has(call.name)); + if (!this.isBuiltinFlowTool(event.toolCall.name)) return; + const firstFlowCall = event.toolCalls.find((call) => this.isBuiltinFlowTool(call.name)); if (firstFlowCall !== undefined && firstFlowCall !== event.toolCall) { event.veto( denyToolExecution( @@ -174,6 +176,13 @@ export class AgentFlowService extends Disposable implements IAgentFlowService { ); } + private isBuiltinFlowTool(name: string): boolean { + if (!FLOW_TOOL_NAMES.has(name)) return false; + return this.toolRegistry + .listReferences() + .some((reference) => reference.name === name && reference.source === 'builtin'); + } + private summary(): AgentFlowRunStatus | null { const run = this.run(); const stage = this.currentStage(); diff --git a/packages/agent-core-v2/test/features/flow/flowService.test.ts b/packages/agent-core-v2/test/features/flow/flowService.test.ts index 055d55798..62620e005 100644 --- a/packages/agent-core-v2/test/features/flow/flowService.test.ts +++ b/packages/agent-core-v2/test/features/flow/flowService.test.ts @@ -11,6 +11,7 @@ import { import { IAgentPermissionModeService } from '#/agent/permissionMode/permissionMode'; import { IAgentStateService } from '#/agent/state/agentState'; import { IAgentToolApprovalService } from '#/agent/toolApproval/toolApproval'; +import { IAgentToolRegistryService } from '#/agent/toolRegistry/toolRegistry'; import { IAgentToolExecutorService } from '#/agent/toolExecutor/toolExecutor'; import type { ResolvedToolExecutionHookContext } from '#/agent/toolExecutor/toolHooks'; import { SkillActivate, skillKey } from '#/agent/skill/skillOps'; @@ -71,6 +72,7 @@ describe('AgentFlowService', () => { let activationDataStore: Map; let contextMessages: ContextMessage[]; let configHandlers: ((e: ConfigChangedEvent) => void)[]; + let flowToolSource: 'builtin' | 'user'; beforeEach(() => { disposables = new DisposableStore(); @@ -89,6 +91,14 @@ describe('AgentFlowService', () => { workDir: '/ws', additionalDirs: [], } as unknown as ISessionWorkspaceContext); + ix.stub(IAgentToolRegistryService, { + listReferences: () => + ['FlowStart', 'FlowAdvance', 'FlowAbort', 'FlowJump'].map((name) => ({ + name, + source: flowToolSource, + })), + } as unknown as IAgentToolRegistryService); + flowToolSource = 'builtin'; activationDataStore = new Map(); contextMessages = []; configHandlers = []; @@ -825,6 +835,25 @@ describe('AgentFlowService', () => { expect(requestToolApproval).not.toHaveBeenCalled(); }); + it('ignores a shadowing registration of a flow tool name', async () => { + flowToolSource = 'user'; + service.start(DEFINITION, 'task'); + const context = advanceContext(GATE_DISPLAY); + const sibling: ToolCall = { + type: 'function', + id: 'call_other', + name: 'Bash', + arguments: '{}', + }; + const batched = { + ...context, + toolCalls: [context.toolCall, sibling], + } as ResolvedToolExecutionHookContext; + const decision = await executorEvents.fireBeforeExecute(batched); + expect(decision?.veto).toBeUndefined(); + expect(requestToolApproval).not.toHaveBeenCalled(); + }); + it('vetoes an AI-gated FlowAdvance batched with a non-flow sibling call', async () => { service.start(DEFINITION, 'task'); const context = advanceContext(undefined);