From 90be9c35876d2c5de61261a5a32cefe3a5f1e6a1 Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Sun, 4 Jan 2026 15:51:23 -0500 Subject: [PATCH] feat(core): Standardize Tool and Agent Invocation constructors (Phase 2) (#15775) --- .../src/agents/delegate-to-agent-tool.test.ts | 2 ++ .../core/src/agents/delegate-to-agent-tool.ts | 16 +++++++-- packages/core/src/agents/local-invocation.ts | 9 ++++- packages/core/src/agents/remote-invocation.ts | 9 ++++- .../src/agents/subagent-tool-wrapper.test.ts | 4 +++ .../core/src/agents/subagent-tool-wrapper.ts | 17 +++++++-- .../core/src/core/coreToolScheduler.test.ts | 35 ++++++++++++++----- packages/core/src/test-utils/mock-tool.ts | 20 +++++++---- packages/core/src/tools/get-internal-docs.ts | 18 ++++++++-- packages/core/src/tools/glob.ts | 2 +- packages/core/src/tools/grep.ts | 2 +- packages/core/src/tools/ls.ts | 2 +- packages/core/src/tools/mcp-tool.ts | 4 +-- .../src/tools/message-bus-integration.test.ts | 16 +++++++-- packages/core/src/tools/read-file.ts | 2 +- packages/core/src/tools/read-many-files.ts | 2 +- packages/core/src/tools/ripGrep.ts | 2 +- packages/core/src/tools/shell.ts | 2 +- packages/core/src/tools/smart-edit.ts | 3 +- packages/core/src/tools/tool-registry.ts | 4 +-- packages/core/src/tools/web-fetch.ts | 2 +- packages/core/src/tools/web-search.ts | 2 +- packages/core/src/tools/write-todos.ts | 9 +++-- 23 files changed, 140 insertions(+), 44 deletions(-) diff --git a/packages/core/src/agents/delegate-to-agent-tool.test.ts b/packages/core/src/agents/delegate-to-agent-tool.test.ts index 0075cb03a7..9afaee7b87 100644 --- a/packages/core/src/agents/delegate-to-agent-tool.test.ts +++ b/packages/core/src/agents/delegate-to-agent-tool.test.ts @@ -100,6 +100,8 @@ describe('DelegateToAgentTool', () => { config, { arg1: 'valid' }, messageBus, + mockAgentDef.name, + mockAgentDef.name, ); }); diff --git a/packages/core/src/agents/delegate-to-agent-tool.ts b/packages/core/src/agents/delegate-to-agent-tool.ts index 7fa42c80a5..435f697c72 100644 --- a/packages/core/src/agents/delegate-to-agent-tool.ts +++ b/packages/core/src/agents/delegate-to-agent-tool.ts @@ -127,12 +127,17 @@ export class DelegateToAgentTool extends BaseDeclarativeTool< protected createInvocation( params: DelegateParams, + messageBus?: MessageBus, + _toolName?: string, + _toolDisplayName?: string, ): ToolInvocation { return new DelegateInvocation( params, this.registry, this.config, - this.messageBus, + messageBus ?? this.messageBus, + _toolName, + _toolDisplayName, ); } } @@ -146,8 +151,15 @@ class DelegateInvocation extends BaseToolInvocation< private readonly registry: AgentRegistry, private readonly config: Config, messageBus?: MessageBus, + _toolName?: string, + _toolDisplayName?: string, ) { - super(params, messageBus, DELEGATE_TO_AGENT_TOOL_NAME); + super( + params, + messageBus, + _toolName ?? DELEGATE_TO_AGENT_TOOL_NAME, + _toolDisplayName, + ); } getDescription(): string { diff --git a/packages/core/src/agents/local-invocation.ts b/packages/core/src/agents/local-invocation.ts index ad27a85a61..1ca315e1ca 100644 --- a/packages/core/src/agents/local-invocation.ts +++ b/packages/core/src/agents/local-invocation.ts @@ -44,8 +44,15 @@ export class LocalSubagentInvocation extends BaseToolInvocation< private readonly config: Config, params: AgentInputs, messageBus?: MessageBus, + _toolName?: string, + _toolDisplayName?: string, ) { - super(params, messageBus, definition.name, definition.displayName); + super( + params, + messageBus, + _toolName ?? definition.name, + _toolDisplayName ?? definition.displayName, + ); } /** diff --git a/packages/core/src/agents/remote-invocation.ts b/packages/core/src/agents/remote-invocation.ts index ee52f2f388..c74af79e95 100644 --- a/packages/core/src/agents/remote-invocation.ts +++ b/packages/core/src/agents/remote-invocation.ts @@ -26,8 +26,15 @@ export class RemoteAgentInvocation extends BaseToolInvocation< private readonly definition: RemoteAgentDefinition, params: AgentInputs, messageBus?: MessageBus, + _toolName?: string, + _toolDisplayName?: string, ) { - super(params, messageBus, definition.name, definition.displayName); + super( + params, + messageBus, + _toolName ?? definition.name, + _toolDisplayName ?? definition.displayName, + ); } getDescription(): string { diff --git a/packages/core/src/agents/subagent-tool-wrapper.test.ts b/packages/core/src/agents/subagent-tool-wrapper.test.ts index 382aafebf8..1b8f6269ec 100644 --- a/packages/core/src/agents/subagent-tool-wrapper.test.ts +++ b/packages/core/src/agents/subagent-tool-wrapper.test.ts @@ -120,6 +120,8 @@ describe('SubagentToolWrapper', () => { mockConfig, params, undefined, + mockDefinition.name, + mockDefinition.displayName, ); }); @@ -139,6 +141,8 @@ describe('SubagentToolWrapper', () => { mockConfig, params, mockMessageBus, + mockDefinition.name, + mockDefinition.displayName, ); }); diff --git a/packages/core/src/agents/subagent-tool-wrapper.ts b/packages/core/src/agents/subagent-tool-wrapper.ts index 69c96014cd..5acc1799a9 100644 --- a/packages/core/src/agents/subagent-tool-wrapper.ts +++ b/packages/core/src/agents/subagent-tool-wrapper.ts @@ -67,17 +67,30 @@ export class SubagentToolWrapper extends BaseDeclarativeTool< */ protected createInvocation( params: AgentInputs, + messageBus?: MessageBus, + _toolName?: string, + _toolDisplayName?: string, ): ToolInvocation { const definition = this.definition; + const effectiveMessageBus = messageBus ?? this.messageBus; + if (definition.kind === 'remote') { - return new RemoteAgentInvocation(definition, params, this.messageBus); + return new RemoteAgentInvocation( + definition, + params, + effectiveMessageBus, + _toolName, + _toolDisplayName, + ); } return new LocalSubagentInvocation( definition, this.config, params, - this.messageBus, + effectiveMessageBus, + _toolName, + _toolDisplayName, ); } } diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index 1da26dc6a3..3a8fad3aa3 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -20,6 +20,7 @@ import type { Config, ToolRegistry, AnyToolInvocation, + MessageBus, } from '../index.js'; import { DEFAULT_TRUNCATE_TOOL_OUTPUT_LINES, @@ -62,10 +63,17 @@ class TestApprovalTool extends BaseDeclarativeTool<{ id: string }, ToolResult> { ); } - protected createInvocation(params: { - id: string; - }): ToolInvocation<{ id: string }, ToolResult> { - return new TestApprovalInvocation(this.config, params); + protected createInvocation( + params: { id: string }, + messageBus?: MessageBus, + _toolName?: string, + _toolDisplayName?: string, + ): ToolInvocation<{ id: string }, ToolResult> { + return new TestApprovalInvocation( + this.config, + params, + messageBus ?? this.messageBus, + ); } } @@ -76,8 +84,9 @@ class TestApprovalInvocation extends BaseToolInvocation< constructor( private config: Config, params: { id: string }, + messageBus?: MessageBus, ) { - super(params); + super(params, messageBus); } getDescription(): string { @@ -124,8 +133,9 @@ class AbortDuringConfirmationInvocation extends BaseToolInvocation< private readonly abortController: AbortController, private readonly abortError: Error, params: Record, + messageBus?: MessageBus, ) { - super(params); + super(params, messageBus); } override async shouldConfirmExecute( @@ -166,11 +176,15 @@ class AbortDuringConfirmationTool extends BaseDeclarativeTool< protected createInvocation( params: Record, + messageBus?: MessageBus, + _toolName?: string, + _toolDisplayName?: string, ): ToolInvocation, ToolResult> { return new AbortDuringConfirmationInvocation( this.abortController, this.abortError, params, + messageBus ?? this.messageBus, ); } } @@ -727,8 +741,8 @@ class MockEditToolInvocation extends BaseToolInvocation< Record, ToolResult > { - constructor(params: Record) { - super(params); + constructor(params: Record, messageBus?: MessageBus) { + super(params, messageBus); } getDescription(): string { @@ -769,8 +783,11 @@ class MockEditTool extends BaseDeclarativeTool< protected createInvocation( params: Record, + messageBus?: MessageBus, + _toolName?: string, + _toolDisplayName?: string, ): ToolInvocation, ToolResult> { - return new MockEditToolInvocation(params); + return new MockEditToolInvocation(params, messageBus ?? this.messageBus); } } diff --git a/packages/core/src/test-utils/mock-tool.ts b/packages/core/src/test-utils/mock-tool.ts index 8f5d44d7fd..4a0eeccc2e 100644 --- a/packages/core/src/test-utils/mock-tool.ts +++ b/packages/core/src/test-utils/mock-tool.ts @@ -47,7 +47,7 @@ class MockToolInvocation extends BaseToolInvocation< constructor( private readonly tool: MockTool, params: { [key: string]: unknown }, - messageBus: MessageBus, + messageBus?: MessageBus, ) { super(params, messageBus, tool.name, tool.displayName); } @@ -122,9 +122,11 @@ export class MockTool extends BaseDeclarativeTool< protected createInvocation( params: { [key: string]: unknown }, - messageBus: MessageBus, + messageBus?: MessageBus, + _toolName?: string, + _toolDisplayName?: string, ): ToolInvocation<{ [key: string]: unknown }, ToolResult> { - return new MockToolInvocation(this, params, messageBus); + return new MockToolInvocation(this, params, messageBus ?? this.messageBus); } } @@ -144,7 +146,7 @@ export class MockModifiableToolInvocation extends BaseToolInvocation< constructor( private readonly tool: MockModifiableTool, params: Record, - messageBus: MessageBus, + messageBus?: MessageBus, ) { super(params, messageBus, tool.name, tool.displayName); } @@ -228,8 +230,14 @@ export class MockModifiableTool protected createInvocation( params: Record, - messageBus: MessageBus, + messageBus?: MessageBus, + _toolName?: string, + _toolDisplayName?: string, ): ToolInvocation, ToolResult> { - return new MockModifiableToolInvocation(this, params, messageBus); + return new MockModifiableToolInvocation( + this, + params, + messageBus ?? this.messageBus, + ); } } diff --git a/packages/core/src/tools/get-internal-docs.ts b/packages/core/src/tools/get-internal-docs.ts index aec44a5272..90637ffced 100644 --- a/packages/core/src/tools/get-internal-docs.ts +++ b/packages/core/src/tools/get-internal-docs.ts @@ -80,8 +80,13 @@ class GetInternalDocsInvocation extends BaseToolInvocation< GetInternalDocsParams, ToolResult > { - constructor(params: GetInternalDocsParams, messageBus?: MessageBus) { - super(params, messageBus, GET_INTERNAL_DOCS_TOOL_NAME); + constructor( + params: GetInternalDocsParams, + messageBus?: MessageBus, + _toolName?: string, + _toolDisplayName?: string, + ) { + super(params, messageBus, _toolName, _toolDisplayName); } override async shouldConfirmExecute( @@ -181,7 +186,14 @@ export class GetInternalDocsTool extends BaseDeclarativeTool< protected createInvocation( params: GetInternalDocsParams, messageBus?: MessageBus, + _toolName?: string, + _toolDisplayName?: string, ): ToolInvocation { - return new GetInternalDocsInvocation(params, messageBus); + return new GetInternalDocsInvocation( + params, + messageBus ?? this.messageBus, + _toolName ?? GetInternalDocsTool.Name, + _toolDisplayName, + ); } } diff --git a/packages/core/src/tools/glob.ts b/packages/core/src/tools/glob.ts index e505d3bf25..40a911d2f0 100644 --- a/packages/core/src/tools/glob.ts +++ b/packages/core/src/tools/glob.ts @@ -355,7 +355,7 @@ export class GlobTool extends BaseDeclarativeTool { return new GlobToolInvocation( this.config, params, - messageBus, + messageBus ?? this.messageBus, _toolName, _toolDisplayName, ); diff --git a/packages/core/src/tools/grep.ts b/packages/core/src/tools/grep.ts index 0df4dc5ec4..a4f35e07b4 100644 --- a/packages/core/src/tools/grep.ts +++ b/packages/core/src/tools/grep.ts @@ -681,7 +681,7 @@ export class GrepTool extends BaseDeclarativeTool { return new GrepToolInvocation( this.config, params, - messageBus, + messageBus ?? this.messageBus, _toolName, _toolDisplayName, ); diff --git a/packages/core/src/tools/ls.ts b/packages/core/src/tools/ls.ts index 34a78038f4..c7e02acaf0 100644 --- a/packages/core/src/tools/ls.ts +++ b/packages/core/src/tools/ls.ts @@ -337,7 +337,7 @@ export class LSTool extends BaseDeclarativeTool { return new LSToolInvocation( this.config, params, - messageBus, + messageBus ?? this.messageBus, _toolName, _toolDisplayName, ); diff --git a/packages/core/src/tools/mcp-tool.ts b/packages/core/src/tools/mcp-tool.ts index 280927a6e0..844a05ab09 100644 --- a/packages/core/src/tools/mcp-tool.ts +++ b/packages/core/src/tools/mcp-tool.ts @@ -290,11 +290,11 @@ export class DiscoveredMCPTool extends BaseDeclarativeTool< this.mcpTool, this.serverName, this.serverToolName, - this.displayName, + _displayName ?? this.displayName, this.trust, params, this.cliConfig, - _messageBus, + _messageBus ?? this.messageBus, ); } } diff --git a/packages/core/src/tools/message-bus-integration.test.ts b/packages/core/src/tools/message-bus-integration.test.ts index 3f40468001..bafa140f5d 100644 --- a/packages/core/src/tools/message-bus-integration.test.ts +++ b/packages/core/src/tools/message-bus-integration.test.ts @@ -87,8 +87,18 @@ class TestTool extends BaseDeclarativeTool { ); } - protected createInvocation(params: TestParams, messageBus: MessageBus) { - return new TestToolInvocation(params, messageBus); + protected createInvocation( + params: TestParams, + messageBus?: MessageBus, + _toolName?: string, + _toolDisplayName?: string, + ) { + return new TestToolInvocation( + params, + messageBus ?? this.messageBus, + _toolName, + _toolDisplayName, + ); } } @@ -128,7 +138,7 @@ describe('Message Bus Integration', () => { expect(publishSpy).toHaveBeenCalledWith({ type: MessageBusType.TOOL_CONFIRMATION_REQUEST, toolCall: { - name: 'TestToolInvocation', + name: 'test-tool', args: { testParam: 'test-value' }, }, correlationId: 'test-correlation-id', diff --git a/packages/core/src/tools/read-file.ts b/packages/core/src/tools/read-file.ts index 4c0aed9565..6b19878385 100644 --- a/packages/core/src/tools/read-file.ts +++ b/packages/core/src/tools/read-file.ts @@ -232,7 +232,7 @@ export class ReadFileTool extends BaseDeclarativeTool< return new ReadFileToolInvocation( this.config, params, - messageBus, + messageBus ?? this.messageBus, _toolName, _toolDisplayName, ); diff --git a/packages/core/src/tools/read-many-files.ts b/packages/core/src/tools/read-many-files.ts index 85c6c4b4aa..de96588337 100644 --- a/packages/core/src/tools/read-many-files.ts +++ b/packages/core/src/tools/read-many-files.ts @@ -535,7 +535,7 @@ Use this tool when the user's query implies needing the content of several files return new ReadManyFilesToolInvocation( this.config, params, - messageBus, + messageBus ?? this.messageBus, _toolName, _toolDisplayName, ); diff --git a/packages/core/src/tools/ripGrep.ts b/packages/core/src/tools/ripGrep.ts index 49a9398c16..973e0a5fa3 100644 --- a/packages/core/src/tools/ripGrep.ts +++ b/packages/core/src/tools/ripGrep.ts @@ -594,7 +594,7 @@ export class RipGrepTool extends BaseDeclarativeTool< this.config, this.geminiIgnoreParser, params, - messageBus, + messageBus ?? this.messageBus, _toolName, _toolDisplayName, ); diff --git a/packages/core/src/tools/shell.ts b/packages/core/src/tools/shell.ts index dc72975301..8554b2e081 100644 --- a/packages/core/src/tools/shell.ts +++ b/packages/core/src/tools/shell.ts @@ -485,7 +485,7 @@ export class ShellTool extends BaseDeclarativeTool< return new ShellToolInvocation( this.config, params, - messageBus, + messageBus ?? this.messageBus, _toolName, _toolDisplayName, ); diff --git a/packages/core/src/tools/smart-edit.ts b/packages/core/src/tools/smart-edit.ts index cd1a563863..b9a79b7218 100644 --- a/packages/core/src/tools/smart-edit.ts +++ b/packages/core/src/tools/smart-edit.ts @@ -955,11 +955,12 @@ A good instruction should concisely answer: protected createInvocation( params: EditToolParams, + messageBus?: MessageBus, ): ToolInvocation { return new EditToolInvocation( this.config, params, - this.messageBus, + messageBus ?? this.messageBus, this.name, this.displayName, ); diff --git a/packages/core/src/tools/tool-registry.ts b/packages/core/src/tools/tool-registry.ts index 60e8080fb0..ef481c42af 100644 --- a/packages/core/src/tools/tool-registry.ts +++ b/packages/core/src/tools/tool-registry.ts @@ -179,9 +179,9 @@ Signal: Signal number or \`(none)\` if no signal was received. return new DiscoveredToolInvocation( this.config, this.originalName, - this.name, + _toolName ?? this.name, params, - messageBus, + messageBus ?? this.messageBus, ); } } diff --git a/packages/core/src/tools/web-fetch.ts b/packages/core/src/tools/web-fetch.ts index 92c0ae9fea..b31f30ae53 100644 --- a/packages/core/src/tools/web-fetch.ts +++ b/packages/core/src/tools/web-fetch.ts @@ -459,7 +459,7 @@ export class WebFetchTool extends BaseDeclarativeTool< return new WebFetchToolInvocation( this.config, params, - messageBus, + messageBus ?? this.messageBus, _toolName, _toolDisplayName, ); diff --git a/packages/core/src/tools/web-search.ts b/packages/core/src/tools/web-search.ts index 1d81bf96d6..6a9e7c9751 100644 --- a/packages/core/src/tools/web-search.ts +++ b/packages/core/src/tools/web-search.ts @@ -238,7 +238,7 @@ export class WebSearchTool extends BaseDeclarativeTool< return new WebSearchToolInvocation( this.config, params, - messageBus, + messageBus ?? this.messageBus, _toolName, _toolDisplayName, ); diff --git a/packages/core/src/tools/write-todos.ts b/packages/core/src/tools/write-todos.ts index a418fac23d..57c1ad5048 100644 --- a/packages/core/src/tools/write-todos.ts +++ b/packages/core/src/tools/write-todos.ts @@ -145,7 +145,7 @@ export class WriteTodosTool extends BaseDeclarativeTool< > { static readonly Name = WRITE_TODOS_TOOL_NAME; - constructor() { + constructor(messageBus?: MessageBus) { super( WriteTodosTool.Name, 'WriteTodos', @@ -180,6 +180,9 @@ export class WriteTodosTool extends BaseDeclarativeTool< required: ['todos'], additionalProperties: false, }, + true, // isOutputMarkdown + false, // canUpdateOutput + messageBus, ); } @@ -248,13 +251,13 @@ export class WriteTodosTool extends BaseDeclarativeTool< protected createInvocation( params: WriteTodosToolParams, - _messageBus?: MessageBus, + messageBus?: MessageBus, _toolName?: string, _displayName?: string, ): ToolInvocation { return new WriteTodosToolInvocation( params, - _messageBus, + messageBus ?? this.messageBus, _toolName, _displayName, );