bug(core): Avoid stateful tool use in executor. (#14305)

This commit is contained in:
joshualitt
2025-12-01 10:54:28 -08:00
committed by GitHub
parent db027dd95b
commit 62f890b5aa
2 changed files with 14 additions and 25 deletions

View File

@@ -65,13 +65,11 @@ const {
mockExecuteToolCall,
mockSetSystemInstruction,
mockCompress,
mockSetTools,
} = vi.hoisted(() => ({
mockSendMessageStream: vi.fn(),
mockExecuteToolCall: vi.fn(),
mockSetSystemInstruction: vi.fn(),
mockCompress: vi.fn(),
mockSetTools: vi.fn(),
}));
let mockChatHistory: Content[] = [];
@@ -94,7 +92,6 @@ vi.mock('../core/geminiChat.js', async (importOriginal) => {
getHistory: vi.fn((_curated?: boolean) => [...mockChatHistory]),
setHistory: mockSetHistory,
setSystemInstruction: mockSetSystemInstruction,
setTools: mockSetTools,
})),
};
});
@@ -238,7 +235,6 @@ describe('AgentExecutor', () => {
mockSetHistory.mockClear();
mockSendMessageStream.mockReset();
mockSetSystemInstruction.mockReset();
mockSetTools.mockReset();
mockExecuteToolCall.mockReset();
mockedLogAgentStart.mockReset();
mockedLogAgentFinish.mockReset();
@@ -258,7 +254,6 @@ describe('AgentExecutor', () => {
({
sendMessageStream: mockSendMessageStream,
setSystemInstruction: mockSetSystemInstruction,
setTools: mockSetTools,
getHistory: vi.fn((_curated?: boolean) => [...mockChatHistory]),
getLastPromptTokenCount: vi.fn(() => 100),
setHistory: mockSetHistory,
@@ -490,8 +485,10 @@ describe('AgentExecutor', () => {
const { modelConfigKey } = getMockMessageParams(0);
expect(modelConfigKey.model).toBe(getModelConfigAlias(definition));
const call = mockSetTools.mock.calls[0];
const sentTools = (call[0] as Tool[])[0].functionDeclarations;
const chatConstructorArgs = MockedGeminiChat.mock.calls[0];
// tools are the 3rd argument (index 2), passed as [{ functionDeclarations: [...] }]
const passedToolsArg = chatConstructorArgs[2] as Tool[];
const sentTools = passedToolsArg[0].functionDeclarations;
expect(sentTools).toBeDefined();
expect(sentTools).toEqual(
@@ -615,8 +612,9 @@ describe('AgentExecutor', () => {
const { modelConfigKey } = getMockMessageParams(0);
expect(modelConfigKey.model).toBe(getModelConfigAlias(definition));
const call = mockSetTools.mock.calls[0];
const sentTools = (call[0] as Tool[])[0].functionDeclarations;
const chatConstructorArgs = MockedGeminiChat.mock.calls[0];
const passedToolsArg = chatConstructorArgs[2] as Tool[];
const sentTools = passedToolsArg[0].functionDeclarations;
expect(sentTools).toBeDefined();
const completeToolDef = sentTools!.find(

View File

@@ -182,7 +182,6 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
private async executeTurn(
chat: GeminiChat,
currentMessage: Content,
tools: FunctionDeclaration[],
turnCounter: number,
combinedSignal: AbortSignal,
timeoutSignal: AbortSignal, // Pass the timeout controller's signal
@@ -192,7 +191,7 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
await this.tryCompressChat(chat, promptId);
const { functionCalls } = await promptIdContext.run(promptId, async () =>
this.callModel(chat, currentMessage, tools, combinedSignal, promptId),
this.callModel(chat, currentMessage, combinedSignal, promptId),
);
if (combinedSignal.aborted) {
@@ -272,7 +271,6 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
*/
private async executeFinalWarningTurn(
chat: GeminiChat,
tools: FunctionDeclaration[],
turnCounter: number,
reason:
| AgentTerminateMode.TIMEOUT
@@ -309,7 +307,6 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
const turnResult = await this.executeTurn(
chat,
recoveryMessage,
tools,
turnCounter, // This will be the "last" turn number
combinedSignal,
graceTimeoutController.signal, // Pass grace signal to identify a *grace* timeout
@@ -387,8 +384,8 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
let chat: GeminiChat | undefined;
let tools: FunctionDeclaration[] | undefined;
try {
chat = await this.createChatObject(inputs);
tools = this.prepareToolsList();
chat = await this.createChatObject(inputs, tools);
const query = this.definition.promptConfig.query
? templateString(this.definition.promptConfig.query, inputs)
: 'Get Started!';
@@ -414,7 +411,6 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
const turnResult = await this.executeTurn(
chat,
currentMessage,
tools,
turnCounter++,
combinedSignal,
timeoutController.signal,
@@ -443,7 +439,6 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
) {
const recoveryResult = await this.executeFinalWarningTurn(
chat,
tools,
turnCounter, // Use current turnCounter for the recovery attempt
terminateReason,
signal, // Pass the external signal
@@ -509,7 +504,6 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
if (chat && tools) {
const recoveryResult = await this.executeFinalWarningTurn(
chat,
tools,
turnCounter, // Use current turnCounter
AgentTerminateMode.TIMEOUT,
signal,
@@ -591,15 +585,9 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
private async callModel(
chat: GeminiChat,
message: Content,
tools: FunctionDeclaration[],
signal: AbortSignal,
promptId: string,
): Promise<{ functionCalls: FunctionCall[]; textResponse: string }> {
if (tools.length > 0) {
// TODO(12622): Move tools back to config.
chat.setTools([{ functionDeclarations: tools }]);
}
const responseStream = await chat.sendMessageStream(
{
model: getModelConfigAlias(this.definition),
@@ -650,7 +638,10 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
}
/** Initializes a `GeminiChat` instance for the agent run. */
private async createChatObject(inputs: AgentInputs): Promise<GeminiChat> {
private async createChatObject(
inputs: AgentInputs,
tools: FunctionDeclaration[],
): Promise<GeminiChat> {
const { promptConfig } = this.definition;
if (!promptConfig.systemPrompt && !promptConfig.initialMessages) {
@@ -673,7 +664,7 @@ export class AgentExecutor<TOutput extends z.ZodTypeAny> {
return new GeminiChat(
this.runtimeContext,
systemInstruction,
[], // set in `callModel`,
[{ functionDeclarations: tools }],
startHistory,
);
} catch (error) {