mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-05-13 13:22:35 -07:00
Merge branch 'main' into memory_usage3
This commit is contained in:
@@ -89,6 +89,7 @@ import {
|
||||
buildUserSteeringHintPrompt,
|
||||
logBillingEvent,
|
||||
ApiKeyUpdatedEvent,
|
||||
LegacyAgentProtocol,
|
||||
type InjectionSource,
|
||||
startMemoryService,
|
||||
} from '@google/gemini-cli-core';
|
||||
@@ -118,6 +119,7 @@ import { computeTerminalTitle } from '../utils/windowTitle.js';
|
||||
import { useTextBuffer } from './components/shared/text-buffer.js';
|
||||
import { useLogger } from './hooks/useLogger.js';
|
||||
import { useGeminiStream } from './hooks/useGeminiStream.js';
|
||||
import { useAgentStream } from './hooks/useAgentStream.js';
|
||||
import { type BackgroundTask } from './hooks/useExecutionLifecycle.js';
|
||||
import { useVim } from './hooks/vim.js';
|
||||
import { type LoadableSettingScope, SettingScope } from '../config/settings.js';
|
||||
@@ -1161,6 +1163,46 @@ Logging in with Google... Restarting Gemini CLI to continue.
|
||||
};
|
||||
}, [config]);
|
||||
|
||||
const streamAgent = useMemo(
|
||||
() =>
|
||||
config?.getAgentSessionInteractiveEnabled()
|
||||
? new LegacyAgentProtocol({ config, getPreferredEditor })
|
||||
: undefined,
|
||||
[config, getPreferredEditor],
|
||||
);
|
||||
|
||||
const activeStream = streamAgent
|
||||
? // eslint-disable-next-line react-hooks/rules-of-hooks
|
||||
useAgentStream({
|
||||
agent: streamAgent,
|
||||
addItem: historyManager.addItem,
|
||||
onCancelSubmit,
|
||||
isShellFocused: embeddedShellFocused,
|
||||
logger,
|
||||
})
|
||||
: // eslint-disable-next-line react-hooks/rules-of-hooks
|
||||
useGeminiStream(
|
||||
config.getGeminiClient(),
|
||||
historyManager.history,
|
||||
historyManager.addItem,
|
||||
config,
|
||||
settings,
|
||||
setDebugMessage,
|
||||
handleSlashCommand,
|
||||
shellModeActive,
|
||||
getPreferredEditor,
|
||||
onAuthError,
|
||||
performMemoryRefresh,
|
||||
modelSwitchedFromQuotaError,
|
||||
setModelSwitchedFromQuotaError,
|
||||
onCancelSubmit,
|
||||
setEmbeddedShellFocused,
|
||||
terminalWidth,
|
||||
terminalHeight,
|
||||
embeddedShellFocused,
|
||||
consumePendingHints,
|
||||
);
|
||||
|
||||
const {
|
||||
streamingState,
|
||||
submitQuery,
|
||||
@@ -1180,27 +1222,7 @@ Logging in with Google... Restarting Gemini CLI to continue.
|
||||
backgroundTasks,
|
||||
dismissBackgroundTask,
|
||||
retryStatus,
|
||||
} = useGeminiStream(
|
||||
config.getGeminiClient(),
|
||||
historyManager.history,
|
||||
historyManager.addItem,
|
||||
config,
|
||||
settings,
|
||||
setDebugMessage,
|
||||
handleSlashCommand,
|
||||
shellModeActive,
|
||||
getPreferredEditor,
|
||||
onAuthError,
|
||||
performMemoryRefresh,
|
||||
modelSwitchedFromQuotaError,
|
||||
setModelSwitchedFromQuotaError,
|
||||
onCancelSubmit,
|
||||
setEmbeddedShellFocused,
|
||||
terminalWidth,
|
||||
terminalHeight,
|
||||
embeddedShellFocused,
|
||||
consumePendingHints,
|
||||
);
|
||||
} = activeStream;
|
||||
|
||||
const pendingHistoryItems = useMemo(
|
||||
() => [...pendingSlashCommandHistoryItems, ...pendingGeminiHistoryItems],
|
||||
@@ -1783,7 +1805,7 @@ Logging in with Google... Restarting Gemini CLI to continue.
|
||||
if (keyMatchers[Command.QUIT](key)) {
|
||||
// If the user presses Ctrl+C, we want to cancel any ongoing requests.
|
||||
// This should happen regardless of the count.
|
||||
cancelOngoingRequest?.();
|
||||
void cancelOngoingRequest?.();
|
||||
|
||||
handleCtrlCPress();
|
||||
return true;
|
||||
|
||||
@@ -1,187 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import { SubagentToolWrapper } from './subagent-tool-wrapper.js';
|
||||
import { LocalSubagentInvocation } from './local-invocation.js';
|
||||
import { makeFakeConfig } from '../test-utils/config.js';
|
||||
import type { LocalAgentDefinition, AgentInputs } from './types.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import { Kind } from '../tools/tools.js';
|
||||
import type { MessageBus } from '../confirmation-bus/message-bus.js';
|
||||
import { createMockMessageBus } from '../test-utils/mock-message-bus.js';
|
||||
|
||||
// Mock dependencies to isolate the SubagentToolWrapper class
|
||||
vi.mock('./local-invocation.js');
|
||||
|
||||
const MockedLocalSubagentInvocation = vi.mocked(LocalSubagentInvocation);
|
||||
|
||||
// Define reusable test data
|
||||
let mockConfig: Config;
|
||||
let mockMessageBus: MessageBus;
|
||||
|
||||
const mockDefinition: LocalAgentDefinition = {
|
||||
kind: 'local',
|
||||
name: 'TestAgent',
|
||||
displayName: 'Test Agent Display Name',
|
||||
description: 'An agent for testing.',
|
||||
inputConfig: {
|
||||
inputSchema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
goal: { type: 'string', description: 'The goal.' },
|
||||
priority: {
|
||||
type: 'number',
|
||||
description: 'The priority.',
|
||||
},
|
||||
},
|
||||
required: ['goal'],
|
||||
},
|
||||
},
|
||||
modelConfig: {
|
||||
model: 'gemini-test-model',
|
||||
generateContentConfig: {
|
||||
temperature: 0,
|
||||
topP: 1,
|
||||
},
|
||||
},
|
||||
runConfig: { maxTimeMinutes: 5 },
|
||||
promptConfig: { systemPrompt: 'You are a test agent.' },
|
||||
};
|
||||
|
||||
describe('SubagentToolWrapper', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockConfig = makeFakeConfig();
|
||||
// .config is already set correctly by the getter on the instance.
|
||||
Object.defineProperty(mockConfig, 'promptId', {
|
||||
get: () => 'test-prompt-id',
|
||||
configurable: true,
|
||||
});
|
||||
mockMessageBus = createMockMessageBus();
|
||||
});
|
||||
|
||||
describe('constructor', () => {
|
||||
it('should correctly configure the tool properties from the agent definition', () => {
|
||||
const wrapper = new SubagentToolWrapper(
|
||||
mockDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
|
||||
expect(wrapper.name).toBe(mockDefinition.name);
|
||||
expect(wrapper.displayName).toBe(mockDefinition.displayName);
|
||||
expect(wrapper.description).toBe(mockDefinition.description);
|
||||
expect(wrapper.kind).toBe(Kind.Agent);
|
||||
expect(wrapper.isOutputMarkdown).toBe(true);
|
||||
expect(wrapper.canUpdateOutput).toBe(true);
|
||||
});
|
||||
|
||||
it('should fall back to the agent name for displayName if it is not provided', () => {
|
||||
const definitionWithoutDisplayName = {
|
||||
...mockDefinition,
|
||||
displayName: undefined,
|
||||
};
|
||||
const wrapper = new SubagentToolWrapper(
|
||||
definitionWithoutDisplayName,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
expect(wrapper.displayName).toBe(definitionWithoutDisplayName.name);
|
||||
});
|
||||
|
||||
it('should generate a valid tool schema using the definition and converted schema', () => {
|
||||
const wrapper = new SubagentToolWrapper(
|
||||
mockDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
const schema = wrapper.schema;
|
||||
|
||||
expect(schema.name).toBe(mockDefinition.name);
|
||||
expect(schema.description).toBe(mockDefinition.description);
|
||||
expect(schema.parametersJsonSchema).toEqual({
|
||||
...(mockDefinition.inputConfig.inputSchema as Record<string, unknown>),
|
||||
properties: {
|
||||
...((
|
||||
mockDefinition.inputConfig.inputSchema as Record<string, unknown>
|
||||
)['properties'] as Record<string, unknown>),
|
||||
wait_for_previous: {
|
||||
type: 'boolean',
|
||||
description:
|
||||
'Set to true to wait for all previously requested tools in this turn to complete before starting. Set to false (or omit) to run in parallel. Use true when this tool depends on the output of previous tools.',
|
||||
},
|
||||
},
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('createInvocation', () => {
|
||||
it('should create a LocalSubagentInvocation with the correct parameters', () => {
|
||||
const wrapper = new SubagentToolWrapper(
|
||||
mockDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
const params: AgentInputs = { goal: 'Test the invocation', priority: 1 };
|
||||
|
||||
// The public `build` method calls the protected `createInvocation` after validation
|
||||
const invocation = wrapper.build(params);
|
||||
|
||||
expect(invocation).toBeInstanceOf(LocalSubagentInvocation);
|
||||
expect(MockedLocalSubagentInvocation).toHaveBeenCalledExactlyOnceWith(
|
||||
mockDefinition,
|
||||
mockConfig,
|
||||
params,
|
||||
mockMessageBus,
|
||||
mockDefinition.name,
|
||||
mockDefinition.displayName,
|
||||
);
|
||||
});
|
||||
|
||||
it('should pass the messageBus to the LocalSubagentInvocation constructor', () => {
|
||||
const specificMessageBus = {
|
||||
publish: vi.fn(),
|
||||
subscribe: vi.fn(),
|
||||
unsubscribe: vi.fn(),
|
||||
} as unknown as MessageBus;
|
||||
const wrapper = new SubagentToolWrapper(
|
||||
mockDefinition,
|
||||
mockConfig,
|
||||
specificMessageBus,
|
||||
);
|
||||
const params: AgentInputs = { goal: 'Test the invocation', priority: 1 };
|
||||
|
||||
wrapper.build(params);
|
||||
|
||||
expect(MockedLocalSubagentInvocation).toHaveBeenCalledWith(
|
||||
mockDefinition,
|
||||
mockConfig,
|
||||
params,
|
||||
specificMessageBus,
|
||||
mockDefinition.name,
|
||||
mockDefinition.displayName,
|
||||
);
|
||||
});
|
||||
|
||||
it('should throw a validation error for invalid parameters before creating an invocation', () => {
|
||||
const wrapper = new SubagentToolWrapper(
|
||||
mockDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
// Missing the required 'goal' parameter
|
||||
const invalidParams = { priority: 1 };
|
||||
|
||||
// The `build` method in the base class performs JSON schema validation
|
||||
// before calling the protected `createInvocation` method.
|
||||
expect(() => wrapper.build(invalidParams)).toThrow(
|
||||
"params must have required property 'goal'",
|
||||
);
|
||||
expect(MockedLocalSubagentInvocation).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -1,106 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
BaseDeclarativeTool,
|
||||
Kind,
|
||||
type ToolInvocation,
|
||||
type ToolResult,
|
||||
} from '../tools/tools.js';
|
||||
|
||||
import { type AgentLoopContext } from '../config/agent-loop-context.js';
|
||||
import type { AgentDefinition, AgentInputs } from './types.js';
|
||||
import { LocalSubagentInvocation } from './local-invocation.js';
|
||||
import { RemoteAgentInvocation } from './remote-invocation.js';
|
||||
import { BrowserAgentInvocation } from './browser/browserAgentInvocation.js';
|
||||
import { BROWSER_AGENT_NAME } from './browser/browserAgentDefinition.js';
|
||||
import type { MessageBus } from '../confirmation-bus/message-bus.js';
|
||||
|
||||
/**
|
||||
* A tool wrapper that dynamically exposes a subagent as a standard,
|
||||
* strongly-typed `DeclarativeTool`.
|
||||
*/
|
||||
export class SubagentToolWrapper extends BaseDeclarativeTool<
|
||||
AgentInputs,
|
||||
ToolResult
|
||||
> {
|
||||
/**
|
||||
* Constructs the tool wrapper.
|
||||
*
|
||||
* The constructor dynamically generates the JSON schema for the tool's
|
||||
* parameters based on the subagent's input configuration.
|
||||
*
|
||||
* @param definition The `AgentDefinition` of the subagent to wrap.
|
||||
* @param context The execution context.
|
||||
* @param messageBus Optional message bus for policy enforcement.
|
||||
*/
|
||||
constructor(
|
||||
private readonly definition: AgentDefinition,
|
||||
private readonly context: AgentLoopContext,
|
||||
messageBus: MessageBus,
|
||||
) {
|
||||
super(
|
||||
definition.name,
|
||||
definition.displayName ?? definition.name,
|
||||
definition.description,
|
||||
Kind.Agent,
|
||||
definition.inputConfig.inputSchema,
|
||||
messageBus,
|
||||
/* isOutputMarkdown */ true,
|
||||
/* canUpdateOutput */ true,
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates an invocation instance for executing the subagent.
|
||||
*
|
||||
* This method is called by the tool framework when the parent agent decides
|
||||
* to use this tool.
|
||||
*
|
||||
* @param params The validated input parameters from the parent agent's call.
|
||||
* @returns A `ToolInvocation` instance ready for execution.
|
||||
*/
|
||||
protected createInvocation(
|
||||
params: AgentInputs,
|
||||
messageBus: MessageBus,
|
||||
_toolName?: string,
|
||||
_toolDisplayName?: string,
|
||||
): ToolInvocation<AgentInputs, ToolResult> {
|
||||
const definition = this.definition;
|
||||
const effectiveMessageBus = messageBus;
|
||||
|
||||
if (definition.kind === 'remote') {
|
||||
return new RemoteAgentInvocation(
|
||||
definition,
|
||||
this.context,
|
||||
params,
|
||||
effectiveMessageBus,
|
||||
_toolName,
|
||||
_toolDisplayName,
|
||||
);
|
||||
}
|
||||
|
||||
// Special handling for browser agent - needs async MCP setup
|
||||
if (definition.name === BROWSER_AGENT_NAME) {
|
||||
return new BrowserAgentInvocation(
|
||||
this.context,
|
||||
params,
|
||||
effectiveMessageBus,
|
||||
_toolName,
|
||||
_toolDisplayName,
|
||||
);
|
||||
}
|
||||
|
||||
return new LocalSubagentInvocation(
|
||||
definition,
|
||||
this.context,
|
||||
params,
|
||||
effectiveMessageBus,
|
||||
_toolName,
|
||||
_toolDisplayName,
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,424 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import { SubagentTool } from './subagent-tool.js';
|
||||
import { SubagentToolWrapper } from './subagent-tool-wrapper.js';
|
||||
import {
|
||||
Kind,
|
||||
type DeclarativeTool,
|
||||
type ToolCallConfirmationDetails,
|
||||
type ToolInvocation,
|
||||
type ToolResult,
|
||||
} from '../tools/tools.js';
|
||||
import type {
|
||||
LocalAgentDefinition,
|
||||
RemoteAgentDefinition,
|
||||
AgentInputs,
|
||||
} from './types.js';
|
||||
import { makeFakeConfig } from '../test-utils/config.js';
|
||||
import { createMockMessageBus } from '../test-utils/mock-message-bus.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import type { MessageBus } from '../confirmation-bus/message-bus.js';
|
||||
import {
|
||||
GeminiCliOperation,
|
||||
GEN_AI_AGENT_DESCRIPTION,
|
||||
GEN_AI_AGENT_NAME,
|
||||
} from '../telemetry/constants.js';
|
||||
import type { ToolRegistry } from 'src/tools/tool-registry.js';
|
||||
|
||||
vi.mock('./subagent-tool-wrapper.js');
|
||||
|
||||
// Mock runInDevTraceSpan
|
||||
const runInDevTraceSpan = vi.hoisted(() =>
|
||||
vi.fn(async (opts, fn) => {
|
||||
const metadata = { attributes: opts.attributes || {} };
|
||||
return fn({
|
||||
metadata,
|
||||
});
|
||||
}),
|
||||
);
|
||||
|
||||
vi.mock('../telemetry/trace.js', () => ({
|
||||
runInDevTraceSpan,
|
||||
}));
|
||||
|
||||
const MockSubagentToolWrapper = vi.mocked(SubagentToolWrapper);
|
||||
|
||||
const testDefinition: LocalAgentDefinition = {
|
||||
kind: 'local',
|
||||
name: 'LocalAgent',
|
||||
description: 'A local agent.',
|
||||
inputConfig: { inputSchema: { type: 'object', properties: {} } },
|
||||
modelConfig: { model: 'test', generateContentConfig: {} },
|
||||
runConfig: { maxTimeMinutes: 1 },
|
||||
promptConfig: { systemPrompt: 'test' },
|
||||
};
|
||||
|
||||
const testRemoteDefinition: RemoteAgentDefinition = {
|
||||
kind: 'remote',
|
||||
name: 'RemoteAgent',
|
||||
description: 'A remote agent.',
|
||||
inputConfig: {
|
||||
inputSchema: { type: 'object', properties: { query: { type: 'string' } } },
|
||||
},
|
||||
agentCardUrl: 'http://example.com/agent',
|
||||
};
|
||||
|
||||
describe('SubAgentInvocation', () => {
|
||||
let mockConfig: Config;
|
||||
let mockMessageBus: MessageBus;
|
||||
let mockInnerInvocation: ToolInvocation<AgentInputs, ToolResult>;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockConfig = makeFakeConfig();
|
||||
// .config is already set correctly by the getter on the instance.
|
||||
Object.defineProperty(mockConfig, 'promptId', {
|
||||
get: () => 'test-prompt-id',
|
||||
configurable: true,
|
||||
});
|
||||
mockMessageBus = createMockMessageBus();
|
||||
mockInnerInvocation = {
|
||||
shouldConfirmExecute: vi.fn(),
|
||||
execute: vi.fn(),
|
||||
params: {},
|
||||
getDescription: vi.fn(),
|
||||
toolLocations: vi.fn(),
|
||||
};
|
||||
|
||||
MockSubagentToolWrapper.prototype.build = vi
|
||||
.fn()
|
||||
.mockReturnValue(mockInnerInvocation);
|
||||
});
|
||||
|
||||
it('should have Kind.Agent', () => {
|
||||
const tool = new SubagentTool(testDefinition, mockConfig, mockMessageBus);
|
||||
expect(tool.kind).toBe(Kind.Agent);
|
||||
});
|
||||
|
||||
it('should delegate shouldConfirmExecute to the inner sub-invocation (local)', async () => {
|
||||
const tool = new SubagentTool(testDefinition, mockConfig, mockMessageBus);
|
||||
const params = {};
|
||||
// @ts-expect-error - accessing protected method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
vi.mocked(mockInnerInvocation.shouldConfirmExecute).mockResolvedValue(
|
||||
false,
|
||||
);
|
||||
|
||||
const abortSignal = new AbortController().signal;
|
||||
const result = await invocation.shouldConfirmExecute(abortSignal);
|
||||
|
||||
expect(result).toBe(false);
|
||||
expect(mockInnerInvocation.shouldConfirmExecute).toHaveBeenCalledWith(
|
||||
abortSignal,
|
||||
);
|
||||
expect(MockSubagentToolWrapper).toHaveBeenCalledWith(
|
||||
testDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
});
|
||||
|
||||
it('should return the correct description', () => {
|
||||
const tool = new SubagentTool(testDefinition, mockConfig, mockMessageBus);
|
||||
const params = {};
|
||||
// @ts-expect-error - accessing protected method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
expect(invocation.getDescription()).toBe(
|
||||
"Delegating to agent 'LocalAgent'",
|
||||
);
|
||||
});
|
||||
|
||||
it('should delegate shouldConfirmExecute to the inner sub-invocation (remote)', async () => {
|
||||
const tool = new SubagentTool(
|
||||
testRemoteDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
const params = { query: 'test' };
|
||||
// @ts-expect-error - accessing protected method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
const confirmationDetails = {
|
||||
type: 'info',
|
||||
title: 'Confirm',
|
||||
prompt: 'Prompt',
|
||||
onConfirm: vi.fn(),
|
||||
} as const;
|
||||
vi.mocked(mockInnerInvocation.shouldConfirmExecute).mockResolvedValue(
|
||||
confirmationDetails as unknown as ToolCallConfirmationDetails,
|
||||
);
|
||||
|
||||
const abortSignal = new AbortController().signal;
|
||||
const result = await invocation.shouldConfirmExecute(abortSignal);
|
||||
|
||||
expect(result).toBe(confirmationDetails);
|
||||
expect(mockInnerInvocation.shouldConfirmExecute).toHaveBeenCalledWith(
|
||||
abortSignal,
|
||||
);
|
||||
expect(MockSubagentToolWrapper).toHaveBeenCalledWith(
|
||||
testRemoteDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
});
|
||||
|
||||
it('should delegate execute to the inner sub-invocation', async () => {
|
||||
const tool = new SubagentTool(testDefinition, mockConfig, mockMessageBus);
|
||||
const params = {};
|
||||
// @ts-expect-error - accessing protected method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
const mockResult: ToolResult = {
|
||||
llmContent: 'success',
|
||||
returnDisplay: 'success',
|
||||
};
|
||||
vi.mocked(mockInnerInvocation.execute).mockResolvedValue(mockResult);
|
||||
|
||||
const abortSignal = new AbortController().signal;
|
||||
const updateOutput = vi.fn();
|
||||
const result = await invocation.execute(abortSignal, updateOutput);
|
||||
|
||||
expect(result).toBe(mockResult);
|
||||
expect(mockInnerInvocation.execute).toHaveBeenCalledWith(
|
||||
abortSignal,
|
||||
updateOutput,
|
||||
);
|
||||
|
||||
expect(runInDevTraceSpan).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
operation: GeminiCliOperation.AgentCall,
|
||||
attributes: expect.objectContaining({
|
||||
[GEN_AI_AGENT_NAME]: testDefinition.name,
|
||||
[GEN_AI_AGENT_DESCRIPTION]: testDefinition.description,
|
||||
}),
|
||||
}),
|
||||
expect.any(Function),
|
||||
);
|
||||
|
||||
// Verify metadata was set on the span
|
||||
const spanCallback = vi.mocked(runInDevTraceSpan).mock.calls[0][1];
|
||||
const mockMetadata = { input: undefined, output: undefined };
|
||||
const mockSpan = { metadata: mockMetadata };
|
||||
await spanCallback(mockSpan as Parameters<typeof spanCallback>[0]);
|
||||
expect(mockMetadata.input).toBe(params);
|
||||
expect(mockMetadata.output).toBe(mockResult);
|
||||
});
|
||||
|
||||
describe('withUserHints', () => {
|
||||
it('should NOT modify query for local agents', async () => {
|
||||
mockConfig = makeFakeConfig({ modelSteering: true });
|
||||
mockConfig.injectionService.addInjection('Test Hint', 'user_steering');
|
||||
|
||||
const tool = new SubagentTool(testDefinition, mockConfig, mockMessageBus);
|
||||
const params = { query: 'original query' };
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const hintedParams = invocation.withUserHints(params);
|
||||
|
||||
expect(hintedParams.query).toBe('original query');
|
||||
});
|
||||
|
||||
it('should NOT modify query for remote agents if model steering is disabled', async () => {
|
||||
mockConfig = makeFakeConfig({ modelSteering: false });
|
||||
mockConfig.injectionService.addInjection('Test Hint', 'user_steering');
|
||||
|
||||
const tool = new SubagentTool(
|
||||
testRemoteDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
const params = { query: 'original query' };
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const hintedParams = invocation.withUserHints(params);
|
||||
|
||||
expect(hintedParams.query).toBe('original query');
|
||||
});
|
||||
|
||||
it('should NOT modify query for remote agents if there are no hints', async () => {
|
||||
mockConfig = makeFakeConfig({ modelSteering: true });
|
||||
|
||||
const tool = new SubagentTool(
|
||||
testRemoteDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
const params = { query: 'original query' };
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const hintedParams = invocation.withUserHints(params);
|
||||
|
||||
expect(hintedParams.query).toBe('original query');
|
||||
});
|
||||
|
||||
it('should prepend hints to query for remote agents when hints exist and steering is enabled', async () => {
|
||||
mockConfig = makeFakeConfig({ modelSteering: true });
|
||||
|
||||
const tool = new SubagentTool(
|
||||
testRemoteDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
const params = { query: 'original query' };
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
mockConfig.injectionService.addInjection('Hint 1', 'user_steering');
|
||||
mockConfig.injectionService.addInjection('Hint 2', 'user_steering');
|
||||
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const hintedParams = invocation.withUserHints(params);
|
||||
|
||||
expect(hintedParams.query).toContain('Hint 1');
|
||||
expect(hintedParams.query).toContain('Hint 2');
|
||||
expect(hintedParams.query).toMatch(/original query$/);
|
||||
});
|
||||
|
||||
it('should NOT include legacy hints added before the invocation was created', async () => {
|
||||
mockConfig = makeFakeConfig({ modelSteering: true });
|
||||
mockConfig.injectionService.addInjection('Legacy Hint', 'user_steering');
|
||||
|
||||
const tool = new SubagentTool(
|
||||
testRemoteDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
const params = { query: 'original query' };
|
||||
|
||||
// Creation of invocation captures the current hint state
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
// Verify no hints are present yet
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
let hintedParams = invocation.withUserHints(params);
|
||||
expect(hintedParams.query).toBe('original query');
|
||||
|
||||
// Add a new hint after creation
|
||||
mockConfig.injectionService.addInjection('New Hint', 'user_steering');
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
hintedParams = invocation.withUserHints(params);
|
||||
|
||||
expect(hintedParams.query).toContain('New Hint');
|
||||
expect(hintedParams.query).not.toContain('Legacy Hint');
|
||||
});
|
||||
|
||||
it('should NOT modify query if query is missing or not a string', async () => {
|
||||
mockConfig = makeFakeConfig({ modelSteering: true });
|
||||
mockConfig.injectionService.addInjection('Hint', 'user_steering');
|
||||
|
||||
const tool = new SubagentTool(
|
||||
testRemoteDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
const params = { other: 'param' };
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const hintedParams = invocation.withUserHints(params);
|
||||
|
||||
expect(hintedParams).toEqual(params);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('SubagentTool Read-Only logic', () => {
|
||||
let mockConfig: Config;
|
||||
let mockMessageBus: MessageBus;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockConfig = makeFakeConfig();
|
||||
// .config is already set correctly by the getter on the instance.
|
||||
Object.defineProperty(mockConfig, 'promptId', {
|
||||
get: () => 'test-prompt-id',
|
||||
configurable: true,
|
||||
});
|
||||
mockMessageBus = createMockMessageBus();
|
||||
});
|
||||
|
||||
it('should be false for remote agents', () => {
|
||||
const tool = new SubagentTool(
|
||||
testRemoteDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
expect(tool.isReadOnly).toBe(false);
|
||||
});
|
||||
|
||||
it('should be true for local agent with only read-only tools', () => {
|
||||
const readOnlyTool = {
|
||||
name: 'read',
|
||||
isReadOnly: true,
|
||||
} as unknown as DeclarativeTool<object, ToolResult>;
|
||||
const registry = {
|
||||
getTool: (name: string) => (name === 'read' ? readOnlyTool : undefined),
|
||||
};
|
||||
vi.spyOn(mockConfig, 'toolRegistry', 'get').mockReturnValue(
|
||||
registry as unknown as ToolRegistry,
|
||||
);
|
||||
|
||||
const defWithTools: LocalAgentDefinition = {
|
||||
...testDefinition,
|
||||
toolConfig: { tools: ['read'] },
|
||||
};
|
||||
const tool = new SubagentTool(defWithTools, mockConfig, mockMessageBus);
|
||||
expect(tool.isReadOnly).toBe(true);
|
||||
});
|
||||
|
||||
it('should be false for local agent with at least one non-read-only tool', () => {
|
||||
const readOnlyTool = {
|
||||
name: 'read',
|
||||
isReadOnly: true,
|
||||
} as unknown as DeclarativeTool<object, ToolResult>;
|
||||
const mutatorTool = {
|
||||
name: 'write',
|
||||
isReadOnly: false,
|
||||
} as unknown as DeclarativeTool<object, ToolResult>;
|
||||
const registry = {
|
||||
getTool: (name: string) => {
|
||||
if (name === 'read') return readOnlyTool;
|
||||
if (name === 'write') return mutatorTool;
|
||||
return undefined;
|
||||
},
|
||||
};
|
||||
vi.spyOn(mockConfig, 'toolRegistry', 'get').mockReturnValue(
|
||||
registry as unknown as ToolRegistry,
|
||||
);
|
||||
|
||||
const defWithTools: LocalAgentDefinition = {
|
||||
...testDefinition,
|
||||
toolConfig: { tools: ['read', 'write'] },
|
||||
};
|
||||
const tool = new SubagentTool(defWithTools, mockConfig, mockMessageBus);
|
||||
expect(tool.isReadOnly).toBe(false);
|
||||
});
|
||||
|
||||
it('should be true for local agent with no tools', () => {
|
||||
const registry = { getTool: () => undefined };
|
||||
vi.spyOn(mockConfig, 'toolRegistry', 'get').mockReturnValue(
|
||||
registry as unknown as ToolRegistry,
|
||||
);
|
||||
|
||||
const defNoTools: LocalAgentDefinition = {
|
||||
...testDefinition,
|
||||
toolConfig: { tools: [] },
|
||||
};
|
||||
const tool = new SubagentTool(defNoTools, mockConfig, mockMessageBus);
|
||||
expect(tool.isReadOnly).toBe(true);
|
||||
});
|
||||
});
|
||||
@@ -1,237 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
BaseDeclarativeTool,
|
||||
Kind,
|
||||
type ToolInvocation,
|
||||
type ToolResult,
|
||||
BaseToolInvocation,
|
||||
type ToolCallConfirmationDetails,
|
||||
isTool,
|
||||
type ToolLiveOutput,
|
||||
} from '../tools/tools.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import { type AgentLoopContext } from '../config/agent-loop-context.js';
|
||||
import type { MessageBus } from '../confirmation-bus/message-bus.js';
|
||||
import type { AgentDefinition, AgentInputs } from './types.js';
|
||||
import { SubagentToolWrapper } from './subagent-tool-wrapper.js';
|
||||
import { SchemaValidator } from '../utils/schemaValidator.js';
|
||||
import { formatUserHintsForModel } from '../utils/fastAckHelper.js';
|
||||
import { runInDevTraceSpan } from '../telemetry/trace.js';
|
||||
import {
|
||||
GeminiCliOperation,
|
||||
GEN_AI_AGENT_DESCRIPTION,
|
||||
GEN_AI_AGENT_NAME,
|
||||
} from '../telemetry/constants.js';
|
||||
|
||||
export class SubagentTool extends BaseDeclarativeTool<AgentInputs, ToolResult> {
|
||||
constructor(
|
||||
private readonly definition: AgentDefinition,
|
||||
private readonly context: AgentLoopContext,
|
||||
messageBus: MessageBus,
|
||||
) {
|
||||
const inputSchema = definition.inputConfig.inputSchema;
|
||||
|
||||
// Validate schema on construction
|
||||
const schemaError = SchemaValidator.validateSchema(inputSchema);
|
||||
if (schemaError) {
|
||||
throw new Error(
|
||||
`Invalid schema for agent ${definition.name}: ${schemaError}`,
|
||||
);
|
||||
}
|
||||
|
||||
super(
|
||||
definition.name,
|
||||
definition.displayName ?? definition.name,
|
||||
definition.description,
|
||||
Kind.Agent,
|
||||
inputSchema,
|
||||
messageBus,
|
||||
/* isOutputMarkdown */ true,
|
||||
/* canUpdateOutput */ true,
|
||||
);
|
||||
}
|
||||
|
||||
private _memoizedIsReadOnly: boolean | undefined;
|
||||
|
||||
override get isReadOnly(): boolean {
|
||||
if (this._memoizedIsReadOnly !== undefined) {
|
||||
return this._memoizedIsReadOnly;
|
||||
}
|
||||
// No try-catch here. If getToolRegistry() throws, we let it throw.
|
||||
// This is an invariant: you can't check read-only status if the system isn't initialized.
|
||||
this._memoizedIsReadOnly = SubagentTool.checkIsReadOnly(
|
||||
this.definition,
|
||||
this.context,
|
||||
);
|
||||
return this._memoizedIsReadOnly;
|
||||
}
|
||||
|
||||
private static checkIsReadOnly(
|
||||
definition: AgentDefinition,
|
||||
context: AgentLoopContext,
|
||||
): boolean {
|
||||
if (definition.kind === 'remote') {
|
||||
return false;
|
||||
}
|
||||
const tools = definition.toolConfig?.tools ?? [];
|
||||
const registry = context.toolRegistry;
|
||||
|
||||
if (!registry) {
|
||||
return false;
|
||||
}
|
||||
|
||||
for (const tool of tools) {
|
||||
if (typeof tool === 'string') {
|
||||
const resolvedTool = registry.getTool(tool);
|
||||
if (!resolvedTool || !resolvedTool.isReadOnly) {
|
||||
return false;
|
||||
}
|
||||
} else if (isTool(tool)) {
|
||||
if (!tool.isReadOnly) {
|
||||
return false;
|
||||
}
|
||||
} else {
|
||||
// FunctionDeclaration - we don't know, so assume NOT read-only
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
protected createInvocation(
|
||||
params: AgentInputs,
|
||||
messageBus: MessageBus,
|
||||
_toolName?: string,
|
||||
_toolDisplayName?: string,
|
||||
): ToolInvocation<AgentInputs, ToolResult> {
|
||||
return new SubAgentInvocation(
|
||||
params,
|
||||
this.definition,
|
||||
this.context,
|
||||
messageBus,
|
||||
_toolName,
|
||||
_toolDisplayName,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
class SubAgentInvocation extends BaseToolInvocation<AgentInputs, ToolResult> {
|
||||
private readonly startIndex: number;
|
||||
|
||||
constructor(
|
||||
params: AgentInputs,
|
||||
private readonly definition: AgentDefinition,
|
||||
private readonly context: AgentLoopContext,
|
||||
messageBus: MessageBus,
|
||||
_toolName?: string,
|
||||
_toolDisplayName?: string,
|
||||
) {
|
||||
super(
|
||||
params,
|
||||
messageBus,
|
||||
_toolName ?? definition.name,
|
||||
_toolDisplayName ?? definition.displayName ?? definition.name,
|
||||
);
|
||||
this.startIndex = context.config.injectionService.getLatestInjectionIndex();
|
||||
}
|
||||
|
||||
private get config(): Config {
|
||||
return this.context.config;
|
||||
}
|
||||
|
||||
getDescription(): string {
|
||||
return `Delegating to agent '${this.definition.name}'`;
|
||||
}
|
||||
|
||||
override async shouldConfirmExecute(
|
||||
abortSignal: AbortSignal,
|
||||
): Promise<ToolCallConfirmationDetails | false> {
|
||||
const invocation = this.buildSubInvocation(
|
||||
this.definition,
|
||||
this.withUserHints(this.params),
|
||||
);
|
||||
return invocation.shouldConfirmExecute(abortSignal);
|
||||
}
|
||||
|
||||
async execute(
|
||||
signal: AbortSignal,
|
||||
updateOutput?: (output: ToolLiveOutput) => void,
|
||||
): Promise<ToolResult> {
|
||||
const validationError = SchemaValidator.validate(
|
||||
this.definition.inputConfig.inputSchema,
|
||||
this.params,
|
||||
);
|
||||
|
||||
if (validationError) {
|
||||
throw new Error(
|
||||
`Invalid arguments for agent '${this.definition.name}': ${validationError}. Input schema: ${JSON.stringify(this.definition.inputConfig.inputSchema)}.`,
|
||||
);
|
||||
}
|
||||
|
||||
const invocation = this.buildSubInvocation(
|
||||
this.definition,
|
||||
this.withUserHints(this.params),
|
||||
);
|
||||
|
||||
return runInDevTraceSpan(
|
||||
{
|
||||
operation: GeminiCliOperation.AgentCall,
|
||||
logPrompts: this.context.config.getTelemetryLogPromptsEnabled(),
|
||||
sessionId: this.context.config.getSessionId(),
|
||||
attributes: {
|
||||
[GEN_AI_AGENT_NAME]: this.definition.name,
|
||||
[GEN_AI_AGENT_DESCRIPTION]: this.definition.description,
|
||||
},
|
||||
},
|
||||
async ({ metadata }) => {
|
||||
metadata.input = this.params;
|
||||
const result = await invocation.execute(signal, updateOutput);
|
||||
metadata.output = result;
|
||||
return result;
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
private withUserHints(agentArgs: AgentInputs): AgentInputs {
|
||||
if (this.definition.kind !== 'remote') {
|
||||
return agentArgs;
|
||||
}
|
||||
|
||||
const userHints = this.config.injectionService.getInjectionsAfter(
|
||||
this.startIndex,
|
||||
'user_steering',
|
||||
);
|
||||
const formattedHints = formatUserHintsForModel(userHints);
|
||||
if (!formattedHints) {
|
||||
return agentArgs;
|
||||
}
|
||||
|
||||
const query = agentArgs['query'];
|
||||
if (typeof query !== 'string' || query.trim().length === 0) {
|
||||
return agentArgs;
|
||||
}
|
||||
|
||||
return {
|
||||
...agentArgs,
|
||||
query: `${formattedHints}\n\n${query}`,
|
||||
};
|
||||
}
|
||||
|
||||
private buildSubInvocation(
|
||||
definition: AgentDefinition,
|
||||
agentArgs: AgentInputs,
|
||||
): ToolInvocation<AgentInputs, ToolResult> {
|
||||
const wrapper = new SubagentToolWrapper(
|
||||
definition,
|
||||
this.context,
|
||||
this.messageBus,
|
||||
);
|
||||
|
||||
return wrapper.build(agentArgs);
|
||||
}
|
||||
}
|
||||
@@ -192,10 +192,6 @@ vi.mock('../agents/registry.js', () => {
|
||||
return { AgentRegistry: AgentRegistryMock };
|
||||
});
|
||||
|
||||
vi.mock('../agents/subagent-tool.js', () => ({
|
||||
SubagentTool: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('../resources/resource-registry.js', () => ({
|
||||
ResourceRegistry: vi.fn(),
|
||||
}));
|
||||
|
||||
@@ -348,4 +348,66 @@ describe('MessageBus', () => {
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('subscribe with AbortSignal', () => {
|
||||
it('should remove listener when signal is aborted', async () => {
|
||||
const handler = vi.fn();
|
||||
const controller = new AbortController();
|
||||
|
||||
messageBus.subscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler, {
|
||||
signal: controller.signal,
|
||||
});
|
||||
|
||||
const message: ToolExecutionSuccess<string> = {
|
||||
type: MessageBusType.TOOL_EXECUTION_SUCCESS as const,
|
||||
toolCall: { name: 'test' },
|
||||
result: 'test',
|
||||
};
|
||||
|
||||
controller.abort();
|
||||
|
||||
await messageBus.publish(message);
|
||||
|
||||
expect(handler).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should not add listener if signal is already aborted', async () => {
|
||||
const handler = vi.fn();
|
||||
const controller = new AbortController();
|
||||
controller.abort();
|
||||
|
||||
messageBus.subscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler, {
|
||||
signal: controller.signal,
|
||||
});
|
||||
|
||||
const message: ToolExecutionSuccess<string> = {
|
||||
type: MessageBusType.TOOL_EXECUTION_SUCCESS as const,
|
||||
toolCall: { name: 'test' },
|
||||
result: 'test',
|
||||
};
|
||||
|
||||
await messageBus.publish(message);
|
||||
|
||||
expect(handler).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should remove abort listener when unsubscribe is called', async () => {
|
||||
const handler = vi.fn();
|
||||
const controller = new AbortController();
|
||||
const signal = controller.signal;
|
||||
|
||||
const removeEventListenerSpy = vi.spyOn(signal, 'removeEventListener');
|
||||
|
||||
messageBus.subscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler, {
|
||||
signal,
|
||||
});
|
||||
|
||||
messageBus.unsubscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler);
|
||||
|
||||
expect(removeEventListenerSpy).toHaveBeenCalledWith(
|
||||
'abort',
|
||||
expect.any(Function),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -13,6 +13,11 @@ import { safeJsonStringify } from '../utils/safeJsonStringify.js';
|
||||
import { debugLogger } from '../utils/debugLogger.js';
|
||||
|
||||
export class MessageBus extends EventEmitter {
|
||||
private listenerToAbortCleanup = new WeakMap<
|
||||
object,
|
||||
Map<string, () => void>
|
||||
>();
|
||||
|
||||
constructor(
|
||||
private readonly policyEngine: PolicyEngine,
|
||||
private readonly debug = false,
|
||||
@@ -145,7 +150,36 @@ export class MessageBus extends EventEmitter {
|
||||
subscribe<T extends Message>(
|
||||
type: T['type'],
|
||||
listener: (message: T) => void,
|
||||
options?: { signal?: AbortSignal },
|
||||
): void {
|
||||
if (options?.signal) {
|
||||
const signal = options.signal;
|
||||
if (signal.aborted) return;
|
||||
|
||||
if (this.listenerToAbortCleanup.get(listener)?.has(type)) return;
|
||||
|
||||
const abortHandler = () => {
|
||||
this.off(type, listener);
|
||||
const typeToCleanup = this.listenerToAbortCleanup.get(listener);
|
||||
if (typeToCleanup) {
|
||||
typeToCleanup.delete(type);
|
||||
if (typeToCleanup.size === 0) {
|
||||
this.listenerToAbortCleanup.delete(listener);
|
||||
}
|
||||
}
|
||||
};
|
||||
signal.addEventListener('abort', abortHandler, { once: true });
|
||||
|
||||
let typeToCleanup = this.listenerToAbortCleanup.get(listener);
|
||||
if (!typeToCleanup) {
|
||||
typeToCleanup = new Map<string, () => void>();
|
||||
this.listenerToAbortCleanup.set(listener, typeToCleanup);
|
||||
}
|
||||
typeToCleanup.set(type, () => {
|
||||
signal.removeEventListener('abort', abortHandler);
|
||||
});
|
||||
}
|
||||
|
||||
this.on(type, listener);
|
||||
}
|
||||
|
||||
@@ -154,6 +188,17 @@ export class MessageBus extends EventEmitter {
|
||||
listener: (message: T) => void,
|
||||
): void {
|
||||
this.off(type, listener);
|
||||
const typeToCleanup = this.listenerToAbortCleanup.get(listener);
|
||||
if (typeToCleanup) {
|
||||
const cleanup = typeToCleanup.get(type);
|
||||
if (cleanup) {
|
||||
cleanup();
|
||||
typeToCleanup.delete(type);
|
||||
}
|
||||
if (typeToCleanup.size === 0) {
|
||||
this.listenerToAbortCleanup.delete(listener);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -80,6 +80,9 @@ priority = 40
|
||||
modes = ["plan"]
|
||||
denyMessage = "You are in Plan Mode with access to read-only tools. Execution of scripts (including those from skills) is blocked."
|
||||
|
||||
# Explicitly allowed tools in Plan Mode (interactive: ask user, non-interactive: deny)
|
||||
# Priority 50 overrides the catch-all (40) and also ensures we override default tier ALLOW rules (e.g. from read-only.toml).
|
||||
|
||||
[[rule]]
|
||||
toolName = "*"
|
||||
mcpName = "*"
|
||||
@@ -89,15 +92,6 @@ priority = 50
|
||||
modes = ["plan"]
|
||||
interactive = true
|
||||
|
||||
[[rule]]
|
||||
toolName = "*"
|
||||
mcpName = "*"
|
||||
toolAnnotations = { readOnlyHint = true }
|
||||
decision = "deny"
|
||||
priority = 50
|
||||
modes = ["plan"]
|
||||
interactive = false
|
||||
|
||||
# Allow specific subagents in Plan mode.
|
||||
# We use argsPattern to match the agent_name argument for invoke_agent.
|
||||
[[rule]]
|
||||
@@ -115,13 +109,6 @@ priority = 50
|
||||
modes = ["plan"]
|
||||
interactive = true
|
||||
|
||||
[[rule]]
|
||||
toolName = ["ask_user", "save_memory", "web_fetch", "activate_skill"]
|
||||
decision = "deny"
|
||||
priority = 50
|
||||
modes = ["plan"]
|
||||
interactive = false
|
||||
|
||||
# Allow write_file and replace for .md files in the plans directory (cross-platform)
|
||||
# We split this into two rules to avoid ReDoS checker issues with nested optional segments.
|
||||
# This rule handles the case where there is a session ID in the plan file path
|
||||
|
||||
@@ -49,6 +49,7 @@ import { resolveConfirmation } from './confirmation.js';
|
||||
import { checkPolicy, updatePolicy } from './policy.js';
|
||||
import { ToolExecutor } from './tool-executor.js';
|
||||
import { ToolModificationHandler } from './tool-modifier.js';
|
||||
import { MessageBusType, type Message } from '../confirmation-bus/types.js';
|
||||
|
||||
vi.mock('./state-manager.js');
|
||||
vi.mock('./confirmation.js');
|
||||
@@ -1299,6 +1300,64 @@ describe('Scheduler (Orchestrator)', () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe('Fallback Handlers', () => {
|
||||
it('should respond to TOOL_CONFIRMATION_REQUEST with requiresUserConfirmation: true', async () => {
|
||||
const listeners: Record<
|
||||
string,
|
||||
Array<(message: Message) => void | Promise<void>>
|
||||
> = {};
|
||||
|
||||
const mockBus = {
|
||||
subscribe: vi.fn(
|
||||
(
|
||||
type: string,
|
||||
handler: (message: Message) => void | Promise<void>,
|
||||
) => {
|
||||
listeners[type] = listeners[type] || [];
|
||||
listeners[type].push(handler);
|
||||
},
|
||||
),
|
||||
publish: vi.fn(async (message: Message) => {
|
||||
const type = message.type as string;
|
||||
if (listeners[type]) {
|
||||
for (const handler of listeners[type]) {
|
||||
await handler(message);
|
||||
}
|
||||
}
|
||||
}),
|
||||
} as unknown as MessageBus;
|
||||
|
||||
const scheduler = new Scheduler({
|
||||
context: mockConfig,
|
||||
messageBus: mockBus,
|
||||
getPreferredEditor,
|
||||
schedulerId: 'fallback-test',
|
||||
});
|
||||
|
||||
const handler = vi.fn();
|
||||
mockBus.subscribe(MessageBusType.TOOL_CONFIRMATION_RESPONSE, handler);
|
||||
|
||||
await mockBus.publish({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_REQUEST,
|
||||
correlationId: 'test-correlation-id',
|
||||
toolCall: { name: 'test-tool' },
|
||||
});
|
||||
|
||||
// Wait for async handler to fire
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
|
||||
expect(handler).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
correlationId: 'test-correlation-id',
|
||||
confirmed: false,
|
||||
requiresUserConfirmation: true,
|
||||
}),
|
||||
);
|
||||
|
||||
scheduler.dispose();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Cleanup', () => {
|
||||
it('should unregister McpProgress listener on dispose()', () => {
|
||||
const onSpy = vi.spyOn(coreEvents, 'on');
|
||||
@@ -1323,6 +1382,40 @@ describe('Scheduler (Orchestrator)', () => {
|
||||
expect.any(Function),
|
||||
);
|
||||
});
|
||||
|
||||
it('should abort disposeController signal on dispose()', () => {
|
||||
const mockSubscribe =
|
||||
vi.fn<
|
||||
(
|
||||
type: unknown,
|
||||
listener: unknown,
|
||||
options?: { signal?: AbortSignal },
|
||||
) => void
|
||||
>();
|
||||
const mockBus = {
|
||||
subscribe: mockSubscribe,
|
||||
publish: vi.fn(),
|
||||
} as unknown as MessageBus;
|
||||
|
||||
let capturedSignal: AbortSignal | undefined;
|
||||
mockSubscribe.mockImplementation((type, listener, options) => {
|
||||
capturedSignal = options?.signal;
|
||||
});
|
||||
|
||||
const s = new Scheduler({
|
||||
context: mockConfig,
|
||||
messageBus: mockBus,
|
||||
getPreferredEditor,
|
||||
schedulerId: 'cleanup-test-2',
|
||||
});
|
||||
|
||||
expect(capturedSignal).toBeDefined();
|
||||
expect(capturedSignal?.aborted).toBe(false);
|
||||
|
||||
s.dispose();
|
||||
|
||||
expect(capturedSignal?.aborted).toBe(true);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
@@ -93,8 +93,7 @@ const createErrorResponse = (
|
||||
* Coordinates execution via state updates and event listening.
|
||||
*/
|
||||
export class Scheduler {
|
||||
// Tracks which MessageBus instances have the legacy listener attached to prevent duplicates.
|
||||
private static subscribedMessageBuses = new WeakSet<MessageBus>();
|
||||
private readonly disposeController = new AbortController();
|
||||
|
||||
private readonly state: SchedulerStateManager;
|
||||
private readonly executor: ToolExecutor;
|
||||
@@ -136,6 +135,7 @@ export class Scheduler {
|
||||
|
||||
dispose(): void {
|
||||
coreEvents.off(CoreEvent.McpProgress, this.handleMcpProgress);
|
||||
this.disposeController.abort();
|
||||
}
|
||||
|
||||
private readonly handleMcpProgress = (payload: McpProgressPayload) => {
|
||||
@@ -163,26 +163,25 @@ export class Scheduler {
|
||||
});
|
||||
};
|
||||
|
||||
private setupMessageBusListener(messageBus: MessageBus): void {
|
||||
if (Scheduler.subscribedMessageBuses.has(messageBus)) {
|
||||
return;
|
||||
}
|
||||
private readonly handleToolConfirmationRequest = async (
|
||||
request: ToolConfirmationRequest,
|
||||
) => {
|
||||
await this.messageBus.publish({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: request.correlationId,
|
||||
confirmed: false,
|
||||
requiresUserConfirmation: true,
|
||||
});
|
||||
};
|
||||
|
||||
private setupMessageBusListener(messageBus: MessageBus): void {
|
||||
// TODO: Optimize policy checks. Currently, tools check policy via
|
||||
// MessageBus even though the Scheduler already checked it.
|
||||
messageBus.subscribe(
|
||||
MessageBusType.TOOL_CONFIRMATION_REQUEST,
|
||||
async (request: ToolConfirmationRequest) => {
|
||||
await messageBus.publish({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: request.correlationId,
|
||||
confirmed: false,
|
||||
requiresUserConfirmation: true,
|
||||
});
|
||||
},
|
||||
this.handleToolConfirmationRequest,
|
||||
{ signal: this.disposeController.signal },
|
||||
);
|
||||
|
||||
Scheduler.subscribedMessageBuses.add(messageBus);
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -81,6 +81,32 @@ describe('classifyGoogleError', () => {
|
||||
}
|
||||
});
|
||||
|
||||
it('should return RetryableQuotaError with delay for 503 Service Unavailable with RetryInfo', () => {
|
||||
const apiError: GoogleApiError = {
|
||||
code: 503,
|
||||
message:
|
||||
'No capacity available for model gemini-3.1-pro-preview on the server',
|
||||
details: [
|
||||
{
|
||||
'@type': 'type.googleapis.com/google.rpc.ErrorInfo',
|
||||
reason: 'MODEL_CAPACITY_EXHAUSTED',
|
||||
domain: 'cloudcode-pa.googleapis.com',
|
||||
metadata: {
|
||||
model: 'gemini-3.1-pro-preview',
|
||||
},
|
||||
},
|
||||
{
|
||||
'@type': 'type.googleapis.com/google.rpc.RetryInfo',
|
||||
retryDelay: '9s',
|
||||
},
|
||||
],
|
||||
};
|
||||
vi.spyOn(errorParser, 'parseGoogleApiError').mockReturnValue(apiError);
|
||||
const result = classifyGoogleError(new Error());
|
||||
expect(result).toBeInstanceOf(RetryableQuotaError);
|
||||
expect((result as RetryableQuotaError).retryDelayMs).toBe(9000);
|
||||
});
|
||||
|
||||
it('should return original error if code is not 429, 499 or 503', () => {
|
||||
const apiError: GoogleApiError = {
|
||||
code: 500,
|
||||
|
||||
@@ -14,6 +14,14 @@ import {
|
||||
} from './googleErrors.js';
|
||||
import { getErrorStatus, ModelNotFoundError } from './httpErrors.js';
|
||||
|
||||
// Enum for Google API type strings
|
||||
enum GoogleApiType {
|
||||
ERROR_INFO = 'type.googleapis.com/google.rpc.ErrorInfo',
|
||||
HELP = 'type.googleapis.com/google.rpc.Help',
|
||||
QUOTA_FAILURE = 'type.googleapis.com/google.rpc.QuotaFailure',
|
||||
RETRY_INFO = 'type.googleapis.com/google.rpc.RetryInfo',
|
||||
}
|
||||
|
||||
/**
|
||||
* A non-retryable error indicating a hard quota limit has been reached (e.g., daily limit).
|
||||
*/
|
||||
@@ -136,8 +144,7 @@ function classifyValidationRequiredError(
|
||||
googleApiError: GoogleApiError,
|
||||
): ValidationRequiredError | null {
|
||||
const errorInfo = googleApiError.details.find(
|
||||
(d): d is ErrorInfo =>
|
||||
d['@type'] === 'type.googleapis.com/google.rpc.ErrorInfo',
|
||||
(d): d is ErrorInfo => d['@type'] === GoogleApiType.ERROR_INFO,
|
||||
);
|
||||
|
||||
if (!errorInfo) {
|
||||
@@ -154,7 +161,7 @@ function classifyValidationRequiredError(
|
||||
|
||||
// Try to extract validation info from Help detail first
|
||||
const helpDetail = googleApiError.details.find(
|
||||
(d): d is Help => d['@type'] === 'type.googleapis.com/google.rpc.Help',
|
||||
(d): d is Help => d['@type'] === GoogleApiType.HELP,
|
||||
);
|
||||
|
||||
let validationLink: string | undefined;
|
||||
@@ -198,12 +205,13 @@ function classifyValidationRequiredError(
|
||||
* - 404 errors are classified as `ModelNotFoundError`.
|
||||
* - 403 errors with `VALIDATION_REQUIRED` from cloudcode-pa domains are classified
|
||||
* as `ValidationRequiredError`.
|
||||
* - 429 errors are classified as either `TerminalQuotaError` or `RetryableQuotaError`:
|
||||
* - 429 or 499 errors are classified as either `TerminalQuotaError` or `RetryableQuotaError`:
|
||||
* - CloudCode API: `RATE_LIMIT_EXCEEDED` → `RetryableQuotaError`, `QUOTA_EXHAUSTED` → `TerminalQuotaError`.
|
||||
* - If the error indicates a daily limit (in QuotaFailure), it's a `TerminalQuotaError`.
|
||||
* - If the error has a retry delay, it's a `RetryableQuotaError`.
|
||||
* - If the error indicates a per-minute limit, it's a `RetryableQuotaError`.
|
||||
* - If the error message contains the phrase "Please retry in X[s|ms]", it's a `RetryableQuotaError`.
|
||||
* - 503 errors are classified as `RetryableQuotaError`.
|
||||
*
|
||||
* @param error The error to classify.
|
||||
* @returns A classified error or the original `unknown` error.
|
||||
@@ -227,24 +235,11 @@ export function classifyGoogleError(error: unknown): unknown {
|
||||
}
|
||||
}
|
||||
|
||||
// Check for 503 Service Unavailable errors
|
||||
if (status === 503) {
|
||||
const errorMessage =
|
||||
googleApiError?.message ||
|
||||
(error instanceof Error ? error.message : String(error));
|
||||
return new RetryableQuotaError(
|
||||
errorMessage,
|
||||
googleApiError ?? {
|
||||
code: 503,
|
||||
message: errorMessage,
|
||||
details: [],
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
if (
|
||||
!googleApiError ||
|
||||
(googleApiError.code !== 429 && googleApiError.code !== 499) ||
|
||||
(googleApiError.code !== 429 &&
|
||||
googleApiError.code !== 499 &&
|
||||
googleApiError.code !== 503) ||
|
||||
googleApiError.details.length === 0
|
||||
) {
|
||||
// Fallback: try to parse the error message for a retry delay
|
||||
@@ -265,9 +260,9 @@ export function classifyGoogleError(error: unknown): unknown {
|
||||
}
|
||||
return new RetryableQuotaError(errorMessage, cause, retryDelaySeconds);
|
||||
}
|
||||
} else if (status === 429 || status === 499) {
|
||||
// Fallback: If it is a 429 or 499 but doesn't have a specific "retry in" message,
|
||||
// assume it is a temporary rate limit and retry after 5 sec (same as DEFAULT_RETRY_OPTIONS).
|
||||
} else if (status === 429 || status === 499 || status === 503) {
|
||||
// Fallback: If it is a 429, 499, or 503 but doesn't have a specific "retry in" message,
|
||||
// assume it is a temporary rate limit and retry.
|
||||
return new RetryableQuotaError(
|
||||
errorMessage,
|
||||
googleApiError ?? {
|
||||
@@ -282,18 +277,15 @@ export function classifyGoogleError(error: unknown): unknown {
|
||||
}
|
||||
|
||||
const quotaFailure = googleApiError.details.find(
|
||||
(d): d is QuotaFailure =>
|
||||
d['@type'] === 'type.googleapis.com/google.rpc.QuotaFailure',
|
||||
(d): d is QuotaFailure => d['@type'] === GoogleApiType.QUOTA_FAILURE,
|
||||
);
|
||||
|
||||
const errorInfo = googleApiError.details.find(
|
||||
(d): d is ErrorInfo =>
|
||||
d['@type'] === 'type.googleapis.com/google.rpc.ErrorInfo',
|
||||
(d): d is ErrorInfo => d['@type'] === GoogleApiType.ERROR_INFO,
|
||||
);
|
||||
|
||||
const retryInfo = googleApiError.details.find(
|
||||
(d): d is RetryInfo =>
|
||||
d['@type'] === 'type.googleapis.com/google.rpc.RetryInfo',
|
||||
(d): d is RetryInfo => d['@type'] === GoogleApiType.RETRY_INFO,
|
||||
);
|
||||
|
||||
// 1. Check for long-term limits in QuotaFailure or ErrorInfo
|
||||
@@ -321,7 +313,7 @@ export function classifyGoogleError(error: unknown): unknown {
|
||||
// INSUFFICIENT_G1_CREDITS_BALANCE is always terminal, regardless of domain
|
||||
if (errorInfo.reason === 'INSUFFICIENT_G1_CREDITS_BALANCE') {
|
||||
return new TerminalQuotaError(
|
||||
`${googleApiError.message}`,
|
||||
googleApiError.message,
|
||||
googleApiError,
|
||||
delaySeconds,
|
||||
errorInfo.reason,
|
||||
@@ -335,21 +327,21 @@ export function classifyGoogleError(error: unknown): unknown {
|
||||
const effectiveDelay = delaySeconds ?? 10;
|
||||
if (effectiveDelay > MAX_RETRYABLE_DELAY_SECONDS) {
|
||||
return new TerminalQuotaError(
|
||||
`${googleApiError.message}`,
|
||||
googleApiError.message,
|
||||
googleApiError,
|
||||
effectiveDelay,
|
||||
errorInfo.reason,
|
||||
);
|
||||
}
|
||||
return new RetryableQuotaError(
|
||||
`${googleApiError.message}`,
|
||||
googleApiError.message,
|
||||
googleApiError,
|
||||
effectiveDelay,
|
||||
);
|
||||
}
|
||||
if (errorInfo.reason === 'QUOTA_EXHAUSTED') {
|
||||
return new TerminalQuotaError(
|
||||
`${googleApiError.message}`,
|
||||
googleApiError.message,
|
||||
googleApiError,
|
||||
delaySeconds,
|
||||
errorInfo.reason,
|
||||
@@ -400,19 +392,10 @@ export function classifyGoogleError(error: unknown): unknown {
|
||||
}
|
||||
}
|
||||
|
||||
// If we reached this point and the status is still 429 or 499, we return retryable.
|
||||
if (status === 429 || status === 499) {
|
||||
const errorMessage =
|
||||
googleApiError?.message ||
|
||||
(error instanceof Error ? error.message : String(error));
|
||||
return new RetryableQuotaError(
|
||||
errorMessage,
|
||||
googleApiError ?? {
|
||||
code: status,
|
||||
message: errorMessage,
|
||||
details: [],
|
||||
},
|
||||
);
|
||||
}
|
||||
return error; // Fallback to original error if no specific classification fits.
|
||||
// If we reached this point, the status is 429, 499, or 503 and we have details,
|
||||
// but no specific violation was matched. We return a generic retryable error.
|
||||
const errorMessage =
|
||||
googleApiError.message ||
|
||||
(error instanceof Error ? error.message : String(error));
|
||||
return new RetryableQuotaError(errorMessage, googleApiError);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user