diff --git a/packages/a2a-server/src/commands/memory.test.ts b/packages/a2a-server/src/commands/memory.test.ts index 2d3a5fef91..de5a09fcb2 100644 --- a/packages/a2a-server/src/commands/memory.test.ts +++ b/packages/a2a-server/src/commands/memory.test.ts @@ -177,10 +177,13 @@ describe('a2a-server memory commands', () => { expect.any(AbortSignal), undefined, { - sanitizationConfig: { - allowedEnvironmentVariables: [], - blockedEnvironmentVariables: [], - enableEnvironmentVariableRedaction: false, + shellExecutionConfig: { + sanitizationConfig: { + allowedEnvironmentVariables: [], + blockedEnvironmentVariables: [], + enableEnvironmentVariableRedaction: false, + }, + sandboxManager: undefined, }, }, ); diff --git a/packages/a2a-server/src/commands/memory.ts b/packages/a2a-server/src/commands/memory.ts index f7c3dfa896..f84d57b3fc 100644 --- a/packages/a2a-server/src/commands/memory.ts +++ b/packages/a2a-server/src/commands/memory.ts @@ -103,8 +103,10 @@ export class AddMemoryCommand implements Command { const abortController = new AbortController(); const signal = abortController.signal; await tool.buildAndExecute(result.toolArgs, signal, undefined, { - sanitizationConfig: DEFAULT_SANITIZATION_CONFIG, - sandboxManager: loopContext.sandboxManager, + shellExecutionConfig: { + sanitizationConfig: DEFAULT_SANITIZATION_CONFIG, + sandboxManager: loopContext.sandboxManager, + }, }); await refreshMemory(context.config); return { diff --git a/packages/cli/src/acp/commands/memory.ts b/packages/cli/src/acp/commands/memory.ts index 1154c852a1..f88aaac4f2 100644 --- a/packages/cli/src/acp/commands/memory.ts +++ b/packages/cli/src/acp/commands/memory.ts @@ -104,8 +104,10 @@ export class AddMemoryCommand implements Command { await context.sendMessage(`Saving memory via ${result.toolName}...`); await tool.buildAndExecute(result.toolArgs, signal, undefined, { - sanitizationConfig: DEFAULT_SANITIZATION_CONFIG, - sandboxManager: context.config.sandboxManager, + shellExecutionConfig: { + sanitizationConfig: DEFAULT_SANITIZATION_CONFIG, + sandboxManager: context.config.sandboxManager, + }, }); await refreshMemory(context.config); return { diff --git a/packages/core/src/core/coreToolHookTriggers.test.ts b/packages/core/src/core/coreToolHookTriggers.test.ts index ff9601fc33..414064ff85 100644 --- a/packages/core/src/core/coreToolHookTriggers.test.ts +++ b/packages/core/src/core/coreToolHookTriggers.test.ts @@ -51,10 +51,9 @@ class MockBackgroundableInvocation extends BaseToolInvocation< async execute( _signal: AbortSignal, _updateOutput?: (output: ToolLiveOutput) => void, - _shellExecutionConfig?: unknown, - setExecutionIdCallback?: (executionId: number) => void, + options?: { setExecutionIdCallback?: (executionId: number) => void }, ) { - setExecutionIdCallback?.(4242); + options?.setExecutionIdCallback?.(4242); return { llmContent: 'pid', returnDisplay: 'pid', @@ -111,7 +110,6 @@ describe('executeToolWithHooks', () => { mockTool, undefined, undefined, - undefined, mockConfig, ); @@ -136,7 +134,6 @@ describe('executeToolWithHooks', () => { mockTool, undefined, undefined, - undefined, mockConfig, ); @@ -168,7 +165,6 @@ describe('executeToolWithHooks', () => { mockTool, undefined, undefined, - undefined, mockConfig, ); @@ -200,7 +196,6 @@ describe('executeToolWithHooks', () => { mockTool, undefined, undefined, - undefined, mockConfig, ); @@ -234,7 +229,6 @@ describe('executeToolWithHooks', () => { mockTool, undefined, undefined, - undefined, mockConfig, ); @@ -275,7 +269,6 @@ describe('executeToolWithHooks', () => { mockTool, undefined, undefined, - undefined, mockConfig, ); @@ -298,8 +291,7 @@ describe('executeToolWithHooks', () => { abortSignal, mockTool, undefined, - undefined, - setExecutionIdCallback, + { setExecutionIdCallback }, mockConfig, ); diff --git a/packages/core/src/core/coreToolHookTriggers.ts b/packages/core/src/core/coreToolHookTriggers.ts index 464cfc5f04..6bff4cfdd5 100644 --- a/packages/core/src/core/coreToolHookTriggers.ts +++ b/packages/core/src/core/coreToolHookTriggers.ts @@ -11,10 +11,10 @@ import type { AnyDeclarativeTool, AnyToolInvocation, ToolLiveOutput, + ExecuteOptions, } from '../tools/tools.js'; import { ToolErrorType } from '../tools/tool-error.js'; import { debugLogger } from '../utils/debugLogger.js'; -import type { ShellExecutionConfig } from '../index.js'; import { DiscoveredMCPToolInvocation } from '../tools/mcp-tool.js'; /** @@ -61,8 +61,7 @@ function extractMcpContext( * @param toolName The name of the tool * @param signal Abort signal for cancellation * @param liveOutputCallback Optional callback for live output updates - * @param shellExecutionConfig Optional shell execution config - * @param setExecutionIdCallback Optional callback to set an execution ID for backgroundable invocations + * @param options Optional execution options (shell config, execution ID callback, etc.) * @param config Config to look up MCP server details for hook context * @returns The tool result */ @@ -72,8 +71,7 @@ export async function executeToolWithHooks( signal: AbortSignal, tool: AnyDeclarativeTool, liveOutputCallback?: (outputChunk: ToolLiveOutput) => void, - shellExecutionConfig?: ShellExecutionConfig, - setExecutionIdCallback?: (executionId: number) => void, + options?: ExecuteOptions, config?: Config, originalRequestName?: string, ): Promise { @@ -158,8 +156,7 @@ export async function executeToolWithHooks( const toolResult: ToolResult = await invocation.execute( signal, liveOutputCallback, - shellExecutionConfig, - setExecutionIdCallback, + options, ); // Append notification if parameters were modified diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 40d5ef9411..a76e7aa2d4 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -156,6 +156,12 @@ export * from './services/executionLifecycleService.js'; // Export Injection Service export * from './config/injectionService.js'; +// Export Execution Lifecycle Service +export * from './services/executionLifecycleService.js'; + +// Export Injection Service +export * from './config/injectionService.js'; + // Export base tool definitions export * from './tools/tools.js'; export * from './tools/tool-error.js'; diff --git a/packages/core/src/scheduler/tool-executor.test.ts b/packages/core/src/scheduler/tool-executor.test.ts index 6f3c54d358..ff9edd83f3 100644 --- a/packages/core/src/scheduler/tool-executor.test.ts +++ b/packages/core/src/scheduler/tool-executor.test.ts @@ -570,14 +570,13 @@ describe('ToolExecutor', () => { _sig, _tool, _liveCb, - _shellCfg, - setExecutionIdCallback, + options, _config, _originalRequestName, ) => { // Simulate the tool reporting an execution ID - if (setExecutionIdCallback) { - setExecutionIdCallback(testPid); + if (options?.setExecutionIdCallback) { + options.setExecutionIdCallback(testPid); } return { llmContent: 'done', returnDisplay: 'done' }; }, @@ -624,16 +623,8 @@ describe('ToolExecutor', () => { const testExecutionId = 67890; vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockImplementation( - async ( - _inv, - _name, - _sig, - _tool, - _liveCb, - _shellCfg, - setExecutionIdCallback, - ) => { - setExecutionIdCallback?.(testExecutionId); + async (_inv, _name, _sig, _tool, _liveCb, options) => { + options?.setExecutionIdCallback?.(testExecutionId); return { llmContent: 'done', returnDisplay: 'done' }; }, ); diff --git a/packages/core/src/scheduler/tool-executor.ts b/packages/core/src/scheduler/tool-executor.ts index 83d77c5a0b..81232d39d9 100644 --- a/packages/core/src/scheduler/tool-executor.ts +++ b/packages/core/src/scheduler/tool-executor.ts @@ -112,8 +112,7 @@ export class ToolExecutor { signal, tool, liveOutputCallback, - shellExecutionConfig, - setExecutionIdCallback, + { shellExecutionConfig, setExecutionIdCallback }, this.config, request.originalRequestName, ); diff --git a/packages/core/src/tools/shell.ts b/packages/core/src/tools/shell.ts index 069bcd5981..8917d281bd 100644 --- a/packages/core/src/tools/shell.ts +++ b/packages/core/src/tools/shell.ts @@ -22,13 +22,13 @@ import { type ToolExecuteConfirmationDetails, type PolicyUpdateOptions, type ToolLiveOutput, + type ExecuteOptions, } from './tools.js'; import { getErrorMessage } from '../utils/errors.js'; import { summarizeToolOutput } from '../utils/summarizer.js'; import { ShellExecutionService, - type ShellExecutionConfig, type ShellOutputEvent, } from '../services/shellExecutionService.js'; import { formatBytes } from '../utils/formatters.js'; @@ -150,9 +150,9 @@ export class ShellToolInvocation extends BaseToolInvocation< async execute( signal: AbortSignal, updateOutput?: (output: ToolLiveOutput) => void, - shellExecutionConfig?: ShellExecutionConfig, - setExecutionIdCallback?: (executionId: number) => void, + options?: ExecuteOptions, ): Promise { + const { shellExecutionConfig, setExecutionIdCallback } = options ?? {}; const strippedCommand = stripShellWrapper(this.params.command); if (signal.aborted) { diff --git a/packages/core/src/tools/tools.ts b/packages/core/src/tools/tools.ts index c58396adb8..03dddf4b8f 100644 --- a/packages/core/src/tools/tools.ts +++ b/packages/core/src/tools/tools.ts @@ -22,6 +22,15 @@ import { import { type ApprovalMode } from '../policy/types.js'; import type { SubagentProgress } from '../agents/types.js'; +/** + * Options bag for tool execution, replacing positional parameters that are + * only relevant to specific tool types. + */ +export interface ExecuteOptions { + shellExecutionConfig?: ShellExecutionConfig; + setExecutionIdCallback?: (executionId: number) => void; +} + /** * Represents a validated and ready-to-execute tool call. * An instance of this is created by a `ToolBuilder`. @@ -68,8 +77,7 @@ export interface ToolInvocation< execute( signal: AbortSignal, updateOutput?: (output: ToolLiveOutput) => void, - shellExecutionConfig?: ShellExecutionConfig, - setExecutionIdCallback?: (executionId: number) => void, + options?: ExecuteOptions, ): Promise; /** @@ -325,7 +333,7 @@ export abstract class BaseToolInvocation< abstract execute( signal: AbortSignal, updateOutput?: (output: ToolLiveOutput) => void, - shellExecutionConfig?: ShellExecutionConfig, + options?: ExecuteOptions, ): Promise; } @@ -522,10 +530,10 @@ export abstract class DeclarativeTool< params: TParams, signal: AbortSignal, updateOutput?: (output: ToolLiveOutput) => void, - shellExecutionConfig?: ShellExecutionConfig, + options?: ExecuteOptions, ): Promise { const invocation = this.build(params); - return invocation.execute(signal, updateOutput, shellExecutionConfig); + return invocation.execute(signal, updateOutput, options); } /**