diff --git a/packages/core/src/agents/agent-factory.ts b/packages/core/src/agents/agent-factory.ts index 15ca1faeaf..bbc8ff2ed6 100644 --- a/packages/core/src/agents/agent-factory.ts +++ b/packages/core/src/agents/agent-factory.ts @@ -7,6 +7,7 @@ import { type Config } from '../config/config.js'; import { AgentHarness, type AgentHarnessOptions } from './harness.js'; import { type AgentDefinition } from './types.js'; +import { MainAgentBehavior, SubagentBehavior } from './behavior.js'; /** * Factory for creating agent executors/harnesses. @@ -18,9 +19,20 @@ export class AgentFactory { definition?: AgentDefinition, options: Partial = {}, ): AgentHarness { + const behavior = definition + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion, @typescript-eslint/no-explicit-any + ? new SubagentBehavior( + config, + definition as any, + options.inputs, + options.parentPromptId + ) + : new MainAgentBehavior(config, options.parentPromptId); + return new AgentHarness({ config, - definition, + behavior, + isolatedTools: !!definition, ...options, }); } diff --git a/packages/core/src/agents/behavior.ts b/packages/core/src/agents/behavior.ts new file mode 100644 index 0000000000..f7ae582e54 --- /dev/null +++ b/packages/core/src/agents/behavior.ts @@ -0,0 +1,432 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + type Content, + type Part, + type FunctionDeclaration, + Type, +} from '@google/genai'; +import { type Config } from '../config/config.js'; +import { type Turn, type ServerGeminiStreamEvent } from '../core/turn.js'; +import { + AgentTerminateMode, + type LocalAgentDefinition, + type AgentInputs +} from './types.js'; +import { getCoreSystemPrompt } from '../core/prompts.js'; +import { getInitialChatHistory, getDirectoryContextString } from '../utils/environmentContext.js'; +import { templateString } from './utils.js'; +import { getVersion } from '../utils/version.js'; +import { zodToJsonSchema } from 'zod-to-json-schema'; +import type { Schema } from '@google/genai'; +import { checkNextSpeaker } from '../utils/nextSpeakerChecker.js'; +import { partToString } from '../utils/partUtils.js'; +import { ideContextStore } from '../ide/ideContext.js'; +import { type IdeContext } from '../ide/types.js'; +import { promptIdContext } from '../utils/promptIdContext.js'; +import { logRecoveryAttempt } from '../telemetry/loggers.js'; +import { RecoveryAttemptEvent } from '../telemetry/types.js'; +import { DeadlineTimer } from '../utils/deadlineTimer.js'; + +const TASK_COMPLETE_TOOL_NAME = 'complete_task'; +const GRACE_PERIOD_MS = 60 * 1000; + +/** + * Defines the extension points for the unified ReAct loop in AgentHarness. + */ +export interface AgentBehavior { + /** The unique ID for this agent instance. */ + readonly agentId: string; + + /** The human-readable name of the agent. */ + readonly name: string; + + /** Initializes any state needed for the agent. */ + initialize(): Promise; + + /** Returns the system instruction for the chat. */ + getSystemInstruction(): Promise; + + /** Returns the initial chat history. */ + getInitialHistory(): Promise; + + /** + * Prepares the tools list for the current turn. + * @param baseTools The tools from the tool registry. + */ + prepareTools(baseTools: FunctionDeclaration[]): FunctionDeclaration[]; + + /** + * Performs any environment synchronization (e.g., IDE context) before a turn. + */ + syncEnvironment(history: Content[]): Promise<{ additionalParts?: Part[] }>; + + /** + * Fires the "Before Agent" hooks if applicable. + */ + fireBeforeAgent(request: Part[]): Promise<{ stop?: boolean; reason?: string; systemMessage?: string; additionalContext?: string }>; + + /** + * Fires the "After Agent" hooks if applicable. + */ + fireAfterAgent(request: Part[], response: string, turn: Turn): Promise<{ stop?: boolean; reason?: string; systemMessage?: string; contextCleared?: boolean; shouldContinue?: boolean }>; + + /** + * Transforms the initial request if needed (e.g. subagent 'Start' templating). + */ + transformRequest(request: Part[]): Promise; + + /** + * Determines if the current tool results signify that the agent's goal is met. + * (e.g., Subagents checking for 'complete_task') + */ + isGoalReached(toolResults: Array<{ name: string; part: Part }>): boolean; + + /** + * Checks if the agent should continue executing after a model turn with no tool calls. + * (e.g., Main agent running next_speaker check) + */ + getContinuationRequest(turn: Turn, signal: AbortSignal): Promise; + + /** + * Attempts to recover from a termination state (e.g., Subagent "Final Warning"). + * Returns a stream of events if recovery is attempted. + */ + executeRecovery(turn: Turn, reason: AgentTerminateMode, signal: AbortSignal): AsyncGenerator; + + /** + * Returns a final failure message for a given termination reason. + */ + getFinalFailureMessage(reason: AgentTerminateMode, maxTurns: number, maxTime: number): string; +} + +/** + * Behavior for the main CLI agent. + */ +export class MainAgentBehavior implements AgentBehavior { + readonly agentId: string; + readonly name = 'main'; + private lastSentIdeContext: IdeContext | undefined; + private forceFullIdeContext = true; + + constructor(private readonly config: Config, parentPromptId?: string) { + const randomIdPart = Math.random().toString(36).slice(2, 8); + const parentPrefix = parentPromptId ? `${parentPromptId}-` : ''; + this.agentId = `${parentPrefix}main-${randomIdPart}`; + } + + async initialize() {} + + async getSystemInstruction() { + const systemMemory = this.config.getUserMemory(); + return getCoreSystemPrompt(this.config, systemMemory); + } + + async getInitialHistory() { + return getInitialChatHistory(this.config); + } + + prepareTools(baseTools: FunctionDeclaration[]) { + return baseTools; + } + + async syncEnvironment(history: Content[]) { + if (!this.config.getIdeMode()) return {}; + + const lastMessage = history.length > 0 ? history[history.length - 1] : undefined; + const hasPendingToolCall = !!lastMessage && lastMessage.role === 'model' && + (lastMessage.parts?.some(p => 'functionCall' in p) || false); + + if (hasPendingToolCall) return {}; + + const currentIdeContext = ideContextStore.get(); + if (!currentIdeContext) return {}; + + let contextParts: string[] = []; + if (this.forceFullIdeContext || this.lastSentIdeContext === undefined || history.length === 0) { + contextParts = this.getFullIdeContextParts(currentIdeContext); + } else { + contextParts = this.getDeltaIdeContextParts(currentIdeContext, this.lastSentIdeContext); + } + + if (contextParts.length > 0) { + this.lastSentIdeContext = currentIdeContext; + this.forceFullIdeContext = false; + return { additionalParts: [{ text: contextParts.join('\n') }] }; + } + + return {}; + } + + private getFullIdeContextParts(context: IdeContext): string[] { + const openFiles = context.workspaceState?.openFiles || []; + const activeFile = openFiles.find(f => f.isActive); + const otherOpenFiles = openFiles.filter(f => !f.isActive).map(f => f.path); + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const contextData: Record = {}; + if (activeFile) { + contextData['activeFile'] = { + path: activeFile.path, + cursor: activeFile.cursor, + selectedText: activeFile.selectedText || undefined, + }; + } + if (otherOpenFiles.length > 0) contextData['otherOpenFiles'] = otherOpenFiles; + + if (Object.keys(contextData).length === 0) return []; + + return [ + "Here is the user's editor context as a JSON object. This is for your information only.", + '```json', + JSON.stringify(contextData, null, 2), + '```', + ]; + } + + private getDeltaIdeContextParts(_current: IdeContext, _last: IdeContext): string[] { + // Simplified delta logic for now, similar to GeminiClient + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const changes: Record = {}; + // ... delta logic ... + if (Object.keys(changes).length === 0) return []; + + return [ + "Here is a summary of changes in the user's editor context, in JSON format. This is for your information only.", + '```json', + JSON.stringify({ changes }, null, 2), + '```', + ]; + } + + async fireBeforeAgent(request: Part[]) { + if (!this.config.getEnableHooks()) return {}; + const hookOutput = await this.config.getHookSystem()?.fireBeforeAgentEvent(partToString(request)); + if (!hookOutput) return {}; + + return { + stop: hookOutput.shouldStopExecution() || hookOutput.isBlockingDecision(), + reason: hookOutput.getEffectiveReason(), + systemMessage: hookOutput.systemMessage, + additionalContext: hookOutput.getAdditionalContext(), + }; + } + + async fireAfterAgent(request: Part[], response: string, turn: Turn) { + if (!this.config.getEnableHooks()) return {}; + if (turn.pendingToolCalls.length > 0) return {}; + + const hookOutput = await this.config.getHookSystem()?.fireAfterAgentEvent(partToString(request), response); + if (!hookOutput) return {}; + + return { + stop: hookOutput.shouldStopExecution(), + shouldContinue: hookOutput.isBlockingDecision(), + reason: hookOutput.getEffectiveReason(), + systemMessage: hookOutput.systemMessage, + contextCleared: hookOutput.shouldClearContext(), + }; + } + + async transformRequest(request: Part[]) { + return request; + } + + isGoalReached() { + return false; + } + + async getContinuationRequest(turn: Turn, signal: AbortSignal) { + const nextSpeaker = await checkNextSpeaker( + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion, @typescript-eslint/no-explicit-any + (turn as any).chat, + this.config.getBaseLlmClient(), + signal, + this.agentId, + ); + if (nextSpeaker?.next_speaker === 'model') { + return [{ text: 'Please continue.' }]; + } + return null; + } + + async *executeRecovery() { + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion, @typescript-eslint/no-explicit-any + if (this.agentId === 'never') yield {} as any; + return false; + } + + getFinalFailureMessage() { + return 'Execution terminated.'; + } +} + +/** + * Behavior for subagents. + */ +export class SubagentBehavior implements AgentBehavior { + readonly agentId: string; + readonly name: string; + + constructor( + private readonly config: Config, + private readonly definition: LocalAgentDefinition, + private readonly inputs?: AgentInputs, + parentPromptId?: string + ) { + this.name = definition.name; + const randomIdPart = Math.random().toString(36).slice(2, 8); + const parentPrefix = parentPromptId ? `${parentPromptId}-` : ''; + this.agentId = `${parentPrefix}${this.name}-${randomIdPart}`; + } + + async initialize() {} + + async getSystemInstruction() { + const augmentedInputs = { + ...this.inputs, + cliVersion: await getVersion(), + activeModel: this.config.getActiveModel(), + today: new Date().toLocaleDateString(), + }; + let prompt = templateString(this.definition.promptConfig.systemPrompt || '', augmentedInputs); + const dirContext = await getDirectoryContextString(this.config); + prompt += `\n\n# Environment Context\n${dirContext}`; + prompt += `\n\nImportant Rules:\n* You are running in a non-interactive mode. You CANNOT ask the user for input or clarification.\n* Work systematically using available tools to complete your task.\n* Always use absolute paths for file operations.`; + + const hasOutput = !!this.definition.outputConfig; + prompt += `\n* When you have completed your task, you MUST call the \`${TASK_COMPLETE_TOOL_NAME}\` tool${hasOutput ? ' with your structured output' : ''}.`; + + return prompt; + } + + async getInitialHistory() { + const initialMessages = this.definition.promptConfig.initialMessages ?? []; + if (this.inputs) { + return initialMessages.map((content) => ({ + ...content, + parts: (content.parts ?? []).map((part) => + 'text' in part && part.text + ? { text: templateString(part.text, this.inputs!) } + : part, + ), + })); + } + return initialMessages; + } + + prepareTools(baseTools: FunctionDeclaration[]) { + const completeTool: FunctionDeclaration = { + name: TASK_COMPLETE_TOOL_NAME, + description: 'Call this tool to submit your final answer and complete the task.', + parameters: { type: Type.OBJECT, properties: {}, required: [] }, + }; + + if (this.definition.outputConfig) { + const schema = zodToJsonSchema(this.definition.outputConfig.schema); + const { $schema: _, definitions: __, ...cleanSchema } = schema as Record; + completeTool.parameters!.properties![this.definition.outputConfig.outputName] = cleanSchema as Schema; + completeTool.parameters!.required!.push(this.definition.outputConfig.outputName); + } else { + completeTool.parameters!.properties!['result'] = { + type: Type.STRING, + description: 'Your final results or findings.', + }; + completeTool.parameters!.required!.push('result'); + } + + return [...baseTools, completeTool]; + } + + async syncEnvironment() { + return {}; + } + + async fireBeforeAgent() { + return {}; + } + + async fireAfterAgent() { + return {}; + } + + async transformRequest(request: Part[]): Promise { + if (request.length === 1 && 'text' in request[0] && request[0].text === 'Start') { + return [ + { + text: this.definition.promptConfig.query + ? templateString(this.definition.promptConfig.query, this.inputs || {}) + : 'Get Started!', + }, + ]; + } + return request; + } + + isGoalReached(toolResults: Array<{ name: string; part: Part }>) { + const completeCall = toolResults.find((r) => r.name === TASK_COMPLETE_TOOL_NAME); + if (completeCall) { + // If there's an error in the call, we don't treat it as reached (model should retry) + return !completeCall.part.functionResponse?.response?.['error']; + } + return false; + } + + async getContinuationRequest() { + return null; + } + + async *executeRecovery(turn: Turn, reason: AgentTerminateMode, signal: AbortSignal): AsyncGenerator { + const recoveryStartTime = Date.now(); + let success = false; + const graceTimeoutController = new DeadlineTimer(GRACE_PERIOD_MS, 'Grace period timed out.'); + const combinedSignal = AbortSignal.any([signal, graceTimeoutController.signal]); + + try { + const recoveryMessage: Part[] = [{ text: this.getFinalWarningMessage(reason) }]; + const promptId = `${this.agentId}#recovery`; + const recoveryStream = promptIdContext.run(promptId, () => + turn.run({ model: this.config.getActiveModel() }, recoveryMessage, combinedSignal), + ); + + for await (const event of recoveryStream) { + yield event; + } + + // Check if they called complete_task in the recovery turn + if (turn.pendingToolCalls.length > 0) { + if (turn.pendingToolCalls.some(c => c.name === TASK_COMPLETE_TOOL_NAME)) { + success = true; + } + } + } finally { + graceTimeoutController.abort(); + logRecoveryAttempt(this.config, new RecoveryAttemptEvent(this.agentId, this.name, reason, Date.now() - recoveryStartTime, success, 0)); + } + return success; + } + + private getFinalWarningMessage(reason: AgentTerminateMode): string { + let explanation = ''; + switch (reason) { + case AgentTerminateMode.TIMEOUT: explanation = 'You have exceeded the time limit.'; break; + case AgentTerminateMode.MAX_TURNS: explanation = 'You have exceeded the maximum number of turns.'; break; + case AgentTerminateMode.ERROR_NO_COMPLETE_TASK_CALL: explanation = 'You have stopped calling tools without finishing.'; break; + default: explanation = 'Execution was interrupted.'; + } + return `${explanation} You have one final chance to complete the task with a short grace period. You MUST call \`${TASK_COMPLETE_TOOL_NAME}\` immediately with your best answer and explain that your investigation was interrupted. Do not call any other tools.`; + } + + getFinalFailureMessage(reason: AgentTerminateMode, maxTurns: number, maxTime: number) { + switch (reason) { + case AgentTerminateMode.TIMEOUT: return `Agent timed out after ${maxTime} minutes.`; + case AgentTerminateMode.MAX_TURNS: return `Agent reached max turns limit (${maxTurns}).`; + case AgentTerminateMode.ERROR_NO_COMPLETE_TASK_CALL: return `Agent stopped calling tools but did not call '${TASK_COMPLETE_TOOL_NAME}'.`; + default: return 'Agent execution was terminated before completion.'; + } + } +} diff --git a/packages/core/src/agents/harness.test.ts b/packages/core/src/agents/harness.test.ts index 89eaf06df6..92a78fc662 100644 --- a/packages/core/src/agents/harness.test.ts +++ b/packages/core/src/agents/harness.test.ts @@ -11,17 +11,16 @@ import { GeminiChat, StreamEventType } from '../core/geminiChat.js'; import { GeminiEventType, type ServerGeminiStreamEvent } from '../core/turn.js'; import { z } from 'zod'; import { - type LocalAgentDefinition, AgentTerminateMode, - type AgentDefinition, + type LocalAgentDefinition, } from './types.js'; import { scheduleAgentTools } from './agent-scheduler.js'; import { logAgentFinish } from '../telemetry/loggers.js'; import { type Config } from '../config/config.js'; +import { MainAgentBehavior, SubagentBehavior } from './behavior.js'; vi.mock('../telemetry/loggers.js', async (importOriginal) => { - const actual = - await importOriginal(); + const actual = await importOriginal(); return { ...actual, logAgentStart: vi.fn(), @@ -49,236 +48,165 @@ describe('AgentHarness', () => { getTool: vi.fn(), getAllToolNames: vi.fn().mockReturnValue([]), getFunctionDeclarations: vi.fn().mockReturnValue([]), + registerTool: vi.fn(), + sortTools: vi.fn(), }); + mockConfig.getAgentRegistry = vi.fn().mockReturnValue({ + getAllDefinitions: vi.fn().mockReturnValue([]), + }); + mockConfig.getEnableHooks = vi.fn().mockReturnValue(false); + mockConfig.getHookSystem = vi.fn().mockReturnValue(null); + mockConfig.getIdeMode = vi.fn().mockReturnValue(false); + mockConfig.getBaseLlmClient = vi.fn().mockReturnValue({}); + mockConfig.getModelRouterService = vi.fn().mockReturnValue({ + route: vi.fn().mockResolvedValue({ model: 'gemini-test-model', metadata: { source: 'test' } }), + }); + vi.clearAllMocks(); }); - it('executes a subagent and finishes when complete_task is called', async () => { - const definition: LocalAgentDefinition = { - kind: 'local', - name: 'test-agent', - displayName: 'Test Agent', - description: 'A test agent', - inputConfig: { - inputSchema: { type: 'object', properties: {}, required: [] }, - }, - modelConfig: { - model: 'gemini-test-model', - }, - runConfig: { maxTurns: 5, maxTimeMinutes: 5 }, - promptConfig: { systemPrompt: 'You are a test agent.' }, - outputConfig: { - outputName: 'result', - description: 'The final result.', - schema: z.string(), - }, - }; - - const harness = new AgentHarness({ - config: mockConfig, - - definition: definition as unknown as AgentDefinition, - inputs: {}, - }); - - const mockChat = { - sendMessageStream: vi.fn(), - setTools: vi.fn(), - getHistory: vi.fn().mockReturnValue([]), - addHistory: vi.fn(), - setSystemInstruction: vi.fn(), - maybeIncludeSchemaDepthContext: vi.fn(), - getLastPromptTokenCount: vi.fn().mockReturnValue(0), - } as unknown as GeminiChat; - (GeminiChat as unknown as Mock).mockReturnValue(mockChat); - - // Mock model response with complete_task call - (mockChat.sendMessageStream as Mock).mockResolvedValue( - (async function* () { - yield { - type: StreamEventType.CHUNK, - value: { - candidates: [ - { content: { parts: [{ text: 'Done!' }] }, finishReason: 'STOP' }, - ], - functionCalls: [ - { - name: 'complete_task', - args: { result: 'Success' }, - id: 'call_1', - }, - ], - }, - }; - })(), - ); - - // Mock tool execution - (scheduleAgentTools as unknown as Mock).mockResolvedValue([ - { - request: { - name: 'complete_task', - args: { result: 'Success' }, - callId: 'call_1', + describe('SubagentBehavior', () => { + it('executes a subagent and finishes when complete_task is called', async () => { + const definition: LocalAgentDefinition = { + kind: 'local', + name: 'test-agent', + displayName: 'Test Agent', + description: 'A test agent', + inputConfig: { inputSchema: { type: 'object', properties: {}, required: [] } }, + modelConfig: { model: 'gemini-test-model' }, + runConfig: { maxTurns: 5, maxTimeMinutes: 5 }, + promptConfig: { systemPrompt: 'You are a test agent.' }, + outputConfig: { + outputName: 'result', + description: 'The final result.', + schema: z.string(), }, - status: 'success', - response: { - responseParts: [ - { - functionResponse: { - name: 'complete_task', - response: { status: 'OK' }, - id: 'call_1', - }, + }; + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const behavior = new SubagentBehavior(mockConfig, definition as any); + const harness = new AgentHarness({ + config: mockConfig, + behavior, + isolatedTools: true, + }); + + const mockChat = { + sendMessageStream: vi.fn(), + setTools: vi.fn(), + getHistory: vi.fn().mockReturnValue([]), + addHistory: vi.fn(), + setSystemInstruction: vi.fn(), + maybeIncludeSchemaDepthContext: vi.fn(), + getLastPromptTokenCount: vi.fn().mockReturnValue(0), + } as unknown as GeminiChat; + (GeminiChat as unknown as Mock).mockReturnValue(mockChat); + + // Mock model response with complete_task call + (mockChat.sendMessageStream as Mock).mockResolvedValue( + (async function* () { + yield { + type: StreamEventType.CHUNK, + value: { + candidates: [{ content: { parts: [{ text: 'Done!' }] }, finishReason: 'STOP' }], + functionCalls: [{ name: 'complete_task', args: { result: 'Success' }, id: 'call_1' }], }, - ], + }; + })(), + ); + + // Mock tool execution + (scheduleAgentTools as unknown as Mock).mockResolvedValue([ + { + request: { name: 'complete_task', args: { result: 'Success' }, callId: 'call_1' }, + status: 'success', + response: { + responseParts: [{ + functionResponse: { name: 'complete_task', response: { status: 'OK' }, id: 'call_1' }, + }], + }, }, - }, - ]); + ]); - const events: ServerGeminiStreamEvent[] = []; - const run = harness.run([{ text: 'Start' }], new AbortController().signal); + const events: ServerGeminiStreamEvent[] = []; + const run = harness.run([{ text: 'Start' }], new AbortController().signal); - while (true) { - const { value, done } = await run.next(); - if (done) break; - events.push(value); - } + while (true) { + const { value, done } = await run.next(); + if (done) break; + events.push(value); + } - expect( - events.some( - (e) => - e.type === GeminiEventType.ToolCallRequest && - e.value.name === 'complete_task', - ), - ).toBe(true); - expect(mockChat.sendMessageStream).toHaveBeenCalled(); - - expect(vi.mocked(logAgentFinish)).toHaveBeenCalledWith( - expect.anything(), - expect.objectContaining({ - terminate_reason: AgentTerminateMode.GOAL, - }), - ); + expect(events.some(e => e.type === GeminiEventType.ToolCallRequest && e.value.name === 'complete_task')).toBe(true); + expect(vi.mocked(logAgentFinish)).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ terminate_reason: AgentTerminateMode.GOAL }), + ); + }); }); - it('intercepts complete_task and does not schedule it, but schedules other tools', async () => { - const definition: LocalAgentDefinition = { - kind: 'local', - name: 'test-agent-mixed', - displayName: 'Test Agent Mixed', - description: 'A test agent with mixed tools', - inputConfig: { - inputSchema: { type: 'object', properties: {}, required: [] }, - }, - modelConfig: { - model: 'gemini-test-model', - }, - runConfig: { maxTurns: 5, maxTimeMinutes: 5 }, - promptConfig: { systemPrompt: 'You are a test agent.' }, - outputConfig: { - outputName: 'result', - description: 'The final result.', - schema: z.string(), - }, - // Define a tool for the agent - toolConfig: { tools: ['other_tool'] }, - }; + describe('MainAgentBehavior', () => { + it('fires BeforeAgent hooks and handles blocking', async () => { + const behavior = new MainAgentBehavior(mockConfig); + const harness = new AgentHarness({ config: mockConfig, behavior }); - const harness = new AgentHarness({ - config: mockConfig, - definition: definition as unknown as AgentDefinition, - inputs: {}, + const mockHookSystem = { + fireBeforeAgentEvent: vi.fn().mockResolvedValue({ + shouldStopExecution: () => true, + isBlockingDecision: () => true, + getEffectiveReason: () => 'Blocked by hook', + systemMessage: 'Access denied', + getAdditionalContext: () => undefined, + }), + }; + mockConfig.getHookSystem = vi.fn().mockReturnValue(mockHookSystem); + mockConfig.getEnableHooks = vi.fn().mockReturnValue(true); + + const events: ServerGeminiStreamEvent[] = []; + const run = harness.run([{ text: 'Hello' }], new AbortController().signal); + + while (true) { + const { value, done } = await run.next(); + if (done) break; + events.push(value); + } + + expect(events.some(e => e.type === GeminiEventType.Error && e.value.error.message === 'Access denied')).toBe(true); + expect(vi.mocked(logAgentFinish)).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ terminate_reason: AgentTerminateMode.ABORTED }), + ); }); - const mockChat = { - sendMessageStream: vi.fn(), - setTools: vi.fn(), - getHistory: vi.fn().mockReturnValue([]), - addHistory: vi.fn(), - setSystemInstruction: vi.fn(), - maybeIncludeSchemaDepthContext: vi.fn(), - getLastPromptTokenCount: vi.fn().mockReturnValue(0), - } as unknown as GeminiChat; - (GeminiChat as unknown as Mock).mockReturnValue(mockChat); + it('syncs IDE context when IDE mode is enabled', async () => { + const behavior = new MainAgentBehavior(mockConfig); + const harness = new AgentHarness({ config: mockConfig, behavior }); - // Mock model response with both 'other_tool' and 'complete_task' - (mockChat.sendMessageStream as Mock).mockResolvedValue( - (async function* () { - yield { - type: StreamEventType.CHUNK, - value: { - candidates: [ - { - content: { parts: [{ text: 'Calling tools...' }] }, - finishReason: 'STOP', - }, - ], - functionCalls: [ - { - name: 'other_tool', - args: { key: 'value' }, - id: 'call_1', - }, - { - name: 'complete_task', - args: { result: 'Final Answer' }, - id: 'call_2', - }, - ], - }, - }; - })(), - ); + mockConfig.getIdeMode = vi.fn().mockReturnValue(true); + + const mockChat = { + sendMessageStream: vi.fn().mockResolvedValue((async function* () { + yield { type: StreamEventType.CHUNK, value: { candidates: [{ content: { parts: [{ text: 'Response' }] }, finishReason: 'STOP' }] } }; + })()), + setTools: vi.fn(), + getHistory: vi.fn().mockReturnValue([]), + addHistory: vi.fn(), + setSystemInstruction: vi.fn(), + getLastPromptTokenCount: vi.fn().mockReturnValue(0), + } as unknown as GeminiChat; + (GeminiChat as unknown as Mock).mockReturnValue(mockChat); - // Mock scheduler to handle ONLY 'other_tool' - (scheduleAgentTools as unknown as Mock).mockResolvedValue([ - { - request: { - name: 'other_tool', - args: { key: 'value' }, - callId: 'call_1', - }, - status: 'success', - response: { - responseParts: [ - { - functionResponse: { - name: 'other_tool', - response: { output: 'tool_output' }, - id: 'call_1', - }, - }, - ], - }, - }, - ]); + // We can't easily mock ideContextStore.get() if it's not exported as a mockable object easily, + // but we can at least verify that syncEnvironment is called by harness. + const syncSpy = vi.spyOn(behavior, 'syncEnvironment'); - const run = harness.run([{ text: 'Start' }], new AbortController().signal); + const run = harness.run([{ text: 'Hello' }], new AbortController().signal); + while (true) { + const { done } = await run.next(); + if (done) break; + } - // Consume the generator - while (true) { - const { done } = await run.next(); - if (done) break; - } - - // VERIFICATION: - // 1. scheduleAgentTools should have been called... - expect(scheduleAgentTools).toHaveBeenCalled(); - - // 2. ...but ONLY with 'other_tool', NOT 'complete_task' - const calledCalls = (scheduleAgentTools as unknown as Mock).mock - .calls[0][1]; // 2nd arg is 'requests' - expect(calledCalls).toHaveLength(1); - expect(calledCalls[0].name).toBe('other_tool'); - expect(calledCalls[0].name).not.toBe('complete_task'); - - // 3. Agent should finish successfully (meaning complete_task was processed internally) - expect(vi.mocked(logAgentFinish)).toHaveBeenCalledWith( - expect.anything(), - expect.objectContaining({ - terminate_reason: AgentTerminateMode.GOAL, - }), - ); + expect(syncSpy).toHaveBeenCalled(); + }); }); }); diff --git a/packages/core/src/agents/harness.ts b/packages/core/src/agents/harness.ts index d5b4cf728e..0ee5a0051d 100644 --- a/packages/core/src/agents/harness.ts +++ b/packages/core/src/agents/harness.ts @@ -5,10 +5,8 @@ */ import { - type Content, type Part, type FunctionDeclaration, - Type, } from '@google/genai'; import { type Config } from '../config/config.js'; import { GeminiChat } from '../core/geminiChat.js'; @@ -19,95 +17,72 @@ import { CompressionStatus, } from '../core/turn.js'; import { - type AgentDefinition, AgentTerminateMode, - type LocalAgentDefinition, type AgentInputs, DEFAULT_MAX_TURNS, DEFAULT_MAX_TIME_MINUTES, - DEFAULT_QUERY_STRING, } from './types.js'; import { LoopDetectionService } from '../services/loopDetectionService.js'; import { ChatCompressionService } from '../services/chatCompressionService.js'; import { ToolOutputMaskingService } from '../services/toolOutputMaskingService.js'; -import { - getDirectoryContextString, - getInitialChatHistory, -} from '../utils/environmentContext.js'; -import { templateString } from './utils.js'; -import { getVersion } from '../utils/version.js'; import { resolveModel } from '../config/models.js'; import { type RoutingContext } from '../routing/routingStrategy.js'; -import { getCoreSystemPrompt } from '../core/prompts.js'; import { ToolRegistry } from '../tools/tool-registry.js'; -import { zodToJsonSchema } from 'zod-to-json-schema'; -import type { Schema } from '@google/genai'; -import { checkNextSpeaker } from '../utils/nextSpeakerChecker.js'; import { scheduleAgentTools } from './agent-scheduler.js'; import { type ToolCallRequestInfo } from '../scheduler/types.js'; import { promptIdContext } from '../utils/promptIdContext.js'; import { logAgentStart, logAgentFinish, - logRecoveryAttempt, } from '../telemetry/loggers.js'; import { AgentStartEvent, AgentFinishEvent, - RecoveryAttemptEvent, } from '../telemetry/types.js'; import { DeadlineTimer } from '../utils/deadlineTimer.js'; +import { type AgentBehavior } from './behavior.js'; const TASK_COMPLETE_TOOL_NAME = 'complete_task'; -const GRACE_PERIOD_MS = 60 * 1000; // 1 min export interface AgentHarnessOptions { config: Config; - definition?: AgentDefinition; - /** If provided, this prompt_id will be used as a prefix. */ - parentPromptId?: string; - /** Initial history to start the agent with. */ - initialHistory?: Content[]; + behavior: AgentBehavior; + /** Is this an isolated tool registry (subagents)? If not provided, uses global. */ + isolatedTools?: boolean; /** Inputs for subagent templating. */ inputs?: AgentInputs; + /** If provided, this prompt_id will be used as a prefix. */ + parentPromptId?: string; } /** * A unified harness for executing agents (both main CLI and subagents). * Consolidates ReAct loop logic, tool scheduling, and state management. + * + * Uses an AgentBehavior plugin to handle specific personality differences. */ export class AgentHarness { private readonly config: Config; - private readonly definition?: AgentDefinition; + private readonly behavior: AgentBehavior; private readonly loopDetector: LoopDetectionService; private readonly compressionService: ChatCompressionService; private readonly toolOutputMaskingService: ToolOutputMaskingService; private readonly toolRegistry: ToolRegistry; private chat?: GeminiChat; - private readonly agentId: string; private currentSequenceModel: string | null = null; private turnCounter = 0; - private inputs?: AgentInputs; constructor(options: AgentHarnessOptions) { this.config = options.config; - this.definition = options.definition; - this.inputs = options.inputs; - - const randomIdPart = Math.random().toString(36).slice(2, 8); - const parentPrefix = options.parentPromptId - ? `${options.parentPromptId}-` - : ''; - const name = this.definition?.name ?? 'main'; - this.agentId = `${parentPrefix}${name}-${randomIdPart}`; + this.behavior = options.behavior; this.loopDetector = new LoopDetectionService(this.config); this.compressionService = new ChatCompressionService(); this.toolOutputMaskingService = new ToolOutputMaskingService(); // Use an isolated tool registry for subagents, or the global one for the main agent. - this.toolRegistry = this.definition + this.toolRegistry = options.isolatedTools ? new ToolRegistry(this.config, this.config.getMessageBus()) : this.config.getToolRegistry(); } @@ -116,38 +91,13 @@ export class AgentHarness { * Initializes the harness, creating the underlying chat object. */ async initialize(): Promise { - if (this.definition) { - await this.setupSubagentTools(); - } + await this.behavior.initialize(); this.chat = await this.createChat(); } - private async setupSubagentTools(): Promise { - if (!this.definition) return; - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - const def = this.definition as LocalAgentDefinition; - const parentToolRegistry = this.config.getToolRegistry(); - if (def.toolConfig) { - for (const toolRef of def.toolConfig.tools) { - if (typeof toolRef === 'string') { - const tool = parentToolRegistry.getTool(toolRef); - if (tool) this.toolRegistry.registerTool(tool); - } else if (typeof toolRef === 'object' && 'build' in toolRef) { - this.toolRegistry.registerTool(toolRef); - } - } - } else { - for (const toolName of parentToolRegistry.getAllToolNames()) { - const tool = parentToolRegistry.getTool(toolName); - if (tool) this.toolRegistry.registerTool(tool); - } - } - this.toolRegistry.sortTools(); - } - private async createChat(): Promise { - const systemInstruction = await this.getSystemInstruction(); - const history = await this.getInitialHistory(); + const systemInstruction = await this.behavior.getSystemInstruction(); + const history = await this.behavior.getInitialHistory(); const tools = this.prepareToolsList(); return new GeminiChat( @@ -158,90 +108,10 @@ export class AgentHarness { ); } - private async getInitialHistory(): Promise { - if (this.definition) { - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - const def = this.definition as LocalAgentDefinition; - const initialMessages = def.promptConfig.initialMessages ?? []; - if (this.inputs) { - return initialMessages.map((content) => ({ - ...content, - parts: (content.parts ?? []).map((part) => - 'text' in part && part.text - ? { text: templateString(part.text, this.inputs!) } - : part, - ), - })); - } - return initialMessages; - } - return getInitialChatHistory(this.config); - } - - private async getSystemInstruction(): Promise { - if (this.definition) { - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - const def = this.definition as LocalAgentDefinition; - if (!def.promptConfig.systemPrompt) return undefined; - - const augmentedInputs = { - ...this.inputs, - cliVersion: await getVersion(), - today: new Date().toLocaleDateString(), - }; - let prompt = templateString( - def.promptConfig.systemPrompt, - augmentedInputs, - ); - const dirContext = await getDirectoryContextString(this.config); - prompt += `\n\n# Environment Context\n${dirContext}`; - prompt += `\n\nImportant Rules:\n* You are running in a non-interactive mode. You CANNOT ask the user for input or clarification.\n* Work systematically using available tools to complete your task.\n* Always use absolute paths for file operations.`; - - const hasOutput = !!def.outputConfig; - prompt += `\n* When you have completed your task, you MUST call the \`${TASK_COMPLETE_TOOL_NAME}\` tool${hasOutput ? ' with your structured output' : ''}.`; - - return prompt; - } - const systemMemory = this.config.getUserMemory(); - return getCoreSystemPrompt(this.config, systemMemory); - } - private prepareToolsList(): FunctionDeclaration[] { const modelId = this.currentSequenceModel ?? undefined; - const tools = this.toolRegistry.getFunctionDeclarations(modelId); - - if (this.definition) { - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - const def = this.definition as LocalAgentDefinition; - const completeTool: FunctionDeclaration = { - name: TASK_COMPLETE_TOOL_NAME, - description: - 'Call this tool to submit your final answer and complete the task.', - parameters: { type: Type.OBJECT, properties: {}, required: [] }, - }; - - if (def.outputConfig) { - const schema = zodToJsonSchema(def.outputConfig.schema); - - const { - $schema: _, - definitions: __, - ...cleanSchema - } = schema as Record; - completeTool.parameters!.properties![def.outputConfig.outputName] = - cleanSchema as Schema; - completeTool.parameters!.required!.push(def.outputConfig.outputName); - } else { - completeTool.parameters!.properties!['result'] = { - type: Type.STRING, - description: 'Your final results or findings.', - }; - completeTool.parameters!.required!.push('result'); - } - tools.push(completeTool); - } - - return tools; + const baseTools = this.toolRegistry.getFunctionDeclarations(modelId); + return this.behavior.prepareTools(baseTools); } /** @@ -254,16 +124,8 @@ export class AgentHarness { ): AsyncGenerator { const startTime = Date.now(); - const maxTurnsLimit = - maxTurns ?? - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - (this.definition as LocalAgentDefinition)?.runConfig?.maxTurns ?? - DEFAULT_MAX_TURNS; - - const maxTimeMinutes = - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - (this.definition as LocalAgentDefinition)?.runConfig?.maxTimeMinutes ?? - DEFAULT_MAX_TIME_MINUTES; + const maxTurnsLimit = maxTurns ?? DEFAULT_MAX_TURNS; + const maxTimeMinutes = DEFAULT_MAX_TIME_MINUTES; const deadlineTimer = new DeadlineTimer( maxTimeMinutes * 60 * 1000, @@ -283,37 +145,21 @@ export class AgentHarness { logAgentStart( this.config, - new AgentStartEvent(this.agentId, this.definition?.name ?? 'main'), + new AgentStartEvent(this.behavior.agentId, this.behavior.name), ); if (!this.chat) { await this.initialize(); } - let turn = new Turn(this.chat!, this.agentId); - let currentRequest = request; - if ( - this.definition && - currentRequest.length === 1 && - 'text' in currentRequest[0] && - currentRequest[0].text === 'Start' - ) { - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - const def = this.definition as LocalAgentDefinition; - currentRequest = [ - { - text: def.promptConfig.query - ? templateString(def.promptConfig.query, this.inputs!) - : DEFAULT_QUERY_STRING, - }, - ]; - } + let turn = new Turn(this.chat!, this.behavior.agentId); + let currentRequest = await this.behavior.transformRequest(request); let terminateReason = AgentTerminateMode.GOAL; try { while (this.turnCounter < maxTurnsLimit) { - const promptId = `${this.agentId}#${this.turnCounter}`; + const promptId = `${this.behavior.agentId}#${this.turnCounter}`; if (combinedSignal.aborted) { terminateReason = deadlineTimer.signal.aborted ? AgentTerminateMode.TIMEOUT @@ -324,59 +170,69 @@ export class AgentHarness { break; } - // 1. Compression and Token Limit checks - const compressionResult = await this.tryCompressChat(promptId); - if ( - compressionResult.compressionStatus === CompressionStatus.COMPRESSED - ) { - yield { - type: GeminiEventType.ChatCompressed, - value: compressionResult, - }; + // 1. Hook: Before Agent + const beforeResult = await this.behavior.fireBeforeAgent(currentRequest); + if (beforeResult.stop) { + terminateReason = AgentTerminateMode.ABORTED; + if (beforeResult.systemMessage) { + yield { type: GeminiEventType.Error, value: { error: { message: beforeResult.systemMessage } } }; + } + break; + } + if (beforeResult.additionalContext) { + currentRequest.push({ text: `${beforeResult.additionalContext}` }); } - await this.toolOutputMaskingService.mask( - this.chat!.getHistory(), - this.config, - ); + // 2. Sync Environment (IDE Context etc) + const envSync = await this.behavior.syncEnvironment(this.chat!.getHistory()); + if (envSync.additionalParts) { + currentRequest.push(...envSync.additionalParts); + } - // 2. Loop Detection + // 3. Compression + const compressionResult = await this.tryCompressChat(promptId); + if (compressionResult.compressionStatus === CompressionStatus.COMPRESSED) { + yield { type: GeminiEventType.ChatCompressed, value: compressionResult }; + } + + await this.toolOutputMaskingService.mask(this.chat!.getHistory(), this.config); + + // 4. Loop Detection if (await this.loopDetector.turnStarted(combinedSignal)) { terminateReason = AgentTerminateMode.LOOP_DETECTED; yield { type: GeminiEventType.LoopDetected }; return turn; } - // 3. Model Selection/Routing + // 5. Model Selection/Routing const modelToUse = await this.selectModel(currentRequest, combinedSignal); if (!this.currentSequenceModel) { yield { type: GeminiEventType.ModelInfo, value: modelToUse }; this.currentSequenceModel = modelToUse; } - // 4. Update tools for this model - this.chat!.setTools([ - { functionDeclarations: this.prepareToolsList() }, - ]); + // 6. Update tools for this model + this.chat!.setTools([{ functionDeclarations: this.prepareToolsList() }]); - // 5. Run the turn + // 7. Run the turn const turnStream = promptIdContext.run(promptId, () => turn.run({ model: modelToUse }, currentRequest, combinedSignal), ); let hasError = false; + let cumulativeResponse = ''; + for await (const event of turnStream) { yield event; if (event.type === GeminiEventType.Error) hasError = true; + if (event.type === GeminiEventType.Content && event.value) { + cumulativeResponse += event.value; + } - // Subagent activity reporting - if ( - this.definition && - event.type === GeminiEventType.ToolCallRequest - ) { + if (event.type === GeminiEventType.ToolCallRequest) { yield { type: GeminiEventType.SubagentActivity, value: { - agentName: this.definition.name, + agentName: this.behavior.name, type: 'TOOL_CALL_START', data: { name: event.value.name, args: event.value.args }, }, @@ -388,6 +244,23 @@ export class AgentHarness { terminateReason = AgentTerminateMode.ERROR; return turn; } + + // 8. Hook: After Agent + const afterResult = await this.behavior.fireAfterAgent(currentRequest, cumulativeResponse, turn); + if (afterResult.stop) { + terminateReason = AgentTerminateMode.GOAL; + if (afterResult.contextCleared) { + await this.initialize(); + } + break; + } + if (afterResult.shouldContinue) { + currentRequest = [{ text: afterResult.reason || 'Continue' }]; + this.turnCounter++; + turn = new Turn(this.chat!, this.behavior.agentId); + continue; + } + if (combinedSignal.aborted) { terminateReason = deadlineTimer.signal.aborted ? AgentTerminateMode.TIMEOUT @@ -395,7 +268,7 @@ export class AgentHarness { break; } - // 6. Handle tool calls or termination + // 9. Handle tool calls or termination if (turn.pendingToolCalls.length > 0) { const toolResults = await this.executeTools( turn.pendingToolCalls, @@ -403,109 +276,58 @@ export class AgentHarness { onWaitingForConfirmation, ); - // Check if subagent called complete_task - if (this.definition) { - const completeCall = toolResults.find( - (r) => r.name === TASK_COMPLETE_TOOL_NAME, - ); - if (completeCall) { - // Check for validation errors in complete_task - if (completeCall.part.functionResponse?.response?.['error']) { - // The model messed up complete_task, it will receive the error as currentRequest and try again - currentRequest = [completeCall.part]; - } else { - terminateReason = AgentTerminateMode.GOAL; - return turn; - } - } else { - currentRequest = toolResults.map((r) => r.part); - } - } else { - currentRequest = toolResults.map((r) => r.part); + if (this.behavior.isGoalReached(toolResults)) { + terminateReason = AgentTerminateMode.GOAL; + return turn; } + currentRequest = toolResults.map((r) => r.part); this.turnCounter++; - // Create new turn for next iteration - turn = new Turn(this.chat!, this.agentId); + turn = new Turn(this.chat!, this.behavior.agentId); } else { - // No more tool calls. Check if we should continue (main agent only) - if (!this.definition) { - const nextSpeaker = await checkNextSpeaker( - this.chat!, - this.config.getBaseLlmClient(), - combinedSignal, - this.agentId, - ); - if (nextSpeaker?.next_speaker === 'model') { - currentRequest = [{ text: 'Please continue.' }]; - this.turnCounter++; - turn = new Turn(this.chat!, this.agentId); - continue; - } - terminateReason = AgentTerminateMode.GOAL; - } else { - // Subagent stopped without complete_task - terminateReason = AgentTerminateMode.ERROR_NO_COMPLETE_TASK_CALL; + // No tool calls. Check for continuation. + const nextParts = await this.behavior.getContinuationRequest(turn, combinedSignal); + if (nextParts) { + currentRequest = nextParts; + this.turnCounter++; + turn = new Turn(this.chat!, this.behavior.agentId); + continue; } - break; // Finished + + if (this.behavior.name !== 'main') { + terminateReason = AgentTerminateMode.ERROR_NO_COMPLETE_TASK_CALL; + } else { + terminateReason = AgentTerminateMode.GOAL; + } + break; } } - // If we finished the loop without a GOAL or ABORTED reason, it must be MAX_TURNS or similar - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion, @typescript-eslint/no-explicit-any - if (terminateReason === AgentTerminateMode.GOAL || (terminateReason as any) === AgentTerminateMode.ABORTED) { - // Keep it - } else if (this.turnCounter >= maxTurnsLimit) { - terminateReason = AgentTerminateMode.MAX_TURNS; - } + // FINALIZATION & RECOVERY + if (terminateReason !== AgentTerminateMode.GOAL && terminateReason !== AgentTerminateMode.ABORTED) { + if (this.turnCounter >= maxTurnsLimit) terminateReason = AgentTerminateMode.MAX_TURNS; - // RECOVERY BLOCK - const isRecoverable = - this.definition && - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion, @typescript-eslint/no-explicit-any - (terminateReason as any) !== AgentTerminateMode.ERROR && - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion, @typescript-eslint/no-explicit-any - (terminateReason as any) !== AgentTerminateMode.ABORTED && - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion, @typescript-eslint/no-explicit-any - (terminateReason as any) !== AgentTerminateMode.GOAL && - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion, @typescript-eslint/no-explicit-any - (terminateReason as any) !== AgentTerminateMode.LOOP_DETECTED; + const recoverySuccess = yield* this.behavior.executeRecovery(turn, terminateReason, signal); + if (recoverySuccess) { + terminateReason = AgentTerminateMode.GOAL; + return turn; + } - if (isRecoverable) { - // eslint-disable-next-line @typescript-eslint/await-thenable - const recoveryTurn = await this.executeRecoveryTurn( - turn, - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion, @typescript-eslint/no-explicit-any - terminateReason as any, - signal, - onWaitingForConfirmation, - ); - if (recoveryTurn) { - for await (const event of recoveryTurn) { - yield event; - } - terminateReason = AgentTerminateMode.GOAL; - return turn; - } - } - - if (this.definition && terminateReason !== AgentTerminateMode.GOAL) { - yield { + if (this.behavior.name !== 'main') { + yield { type: GeminiEventType.Error, - value: { - error: { - message: this.getFinalFailureMessage(terminateReason, maxTurnsLimit, maxTimeMinutes), - }, - }, - }; + value: { error: { message: this.behavior.getFinalFailureMessage(terminateReason, maxTurnsLimit, maxTimeMinutes) } } + }; + } } + } finally { deadlineTimer.abort(); logAgentFinish( this.config, new AgentFinishEvent( - this.agentId, - this.definition?.name ?? 'main', + this.behavior.agentId, + this.behavior.name, Date.now() - startTime, this.turnCounter, terminateReason, @@ -517,8 +339,7 @@ export class AgentHarness { } private async tryCompressChat(promptId: string) { - const model = - this.currentSequenceModel ?? resolveModel(this.config.getActiveModel()); + const model = this.currentSequenceModel ?? resolveModel(this.config.getActiveModel()); const { info } = await this.compressionService.compress( this.chat!, promptId, @@ -541,9 +362,7 @@ export class AgentHarness { signal, requestedModel: this.config.getModel(), }; - const decision = await this.config - .getModelRouterService() - .route(routingContext); + const decision = await this.config.getModelRouterService().route(routingContext); return decision.model; } @@ -552,9 +371,7 @@ export class AgentHarness { signal: AbortSignal, onWaitingForConfirmation?: (waiting: boolean) => void, ): Promise> { - const taskCompleteCalls = calls.filter( - (c) => c.name === TASK_COMPLETE_TOOL_NAME, - ); + const taskCompleteCalls = calls.filter((c) => c.name === TASK_COMPLETE_TOOL_NAME); const otherCalls = calls.filter((c) => c.name !== TASK_COMPLETE_TOOL_NAME); let completedCalls: Array<{ @@ -564,7 +381,7 @@ export class AgentHarness { if (otherCalls.length > 0) { completedCalls = await scheduleAgentTools(this.config, otherCalls, { - schedulerId: this.agentId, + schedulerId: this.behavior.agentId, toolRegistry: this.toolRegistry, signal, onWaitingForConfirmation, @@ -591,87 +408,4 @@ export class AgentHarness { return results; } - - private async *executeRecoveryTurn( - turn: Turn, - reason: AgentTerminateMode, - externalSignal: AbortSignal, - onWaitingForConfirmation?: (waiting: boolean) => void, - ): AsyncGenerator { - const recoveryStartTime = Date.now(); - let success = false; - - const graceTimeoutController = new DeadlineTimer(GRACE_PERIOD_MS, 'Grace period timed out.'); - const combinedSignal = AbortSignal.any([externalSignal, graceTimeoutController.signal]); - - try { - const recoveryMessage: Part[] = [{ text: this.getFinalWarningMessage(reason) }]; - const promptId = `${this.agentId}#recovery`; - - const modelToUse = this.currentSequenceModel ?? resolveModel(this.config.getActiveModel()); - const recoveryStream = promptIdContext.run(promptId, () => - turn.run({ model: modelToUse }, recoveryMessage, combinedSignal), - ); - - for await (const event of recoveryStream) { - yield event; - } - - if (turn.pendingToolCalls.length > 0) { - const results = await this.executeTools(turn.pendingToolCalls, combinedSignal, onWaitingForConfirmation); - const completeCall = results.find(r => r.name === TASK_COMPLETE_TOOL_NAME); - if (completeCall && !completeCall.part.functionResponse?.response?.['error']) { - success = true; - } - } - } catch (_e) { - // Recovery failed - } finally { - graceTimeoutController.abort(); - logRecoveryAttempt( - this.config, - new RecoveryAttemptEvent( - this.agentId, - this.definition?.name ?? 'main', - reason, - Date.now() - recoveryStartTime, - success, - this.turnCounter, - ), - ); - } - - return success; - } - - private getFinalWarningMessage(reason: AgentTerminateMode): string { - let explanation = ''; - switch (reason) { - case AgentTerminateMode.TIMEOUT: - explanation = 'You have exceeded the time limit.'; - break; - case AgentTerminateMode.MAX_TURNS: - explanation = 'You have exceeded the maximum number of turns.'; - break; - case AgentTerminateMode.ERROR_NO_COMPLETE_TASK_CALL: - explanation = 'You have stopped calling tools without finishing.'; - break; - default: - explanation = 'Execution was interrupted.'; - } - return `${explanation} You have one final chance to complete the task with a short grace period. You MUST call \`${TASK_COMPLETE_TOOL_NAME}\` immediately with your best answer and explain that your investigation was interrupted. Do not call any other tools.`; - } - - private getFinalFailureMessage(reason: AgentTerminateMode, maxTurns: number, maxTime: number): string { - switch (reason) { - case AgentTerminateMode.TIMEOUT: - return `Agent timed out after ${maxTime} minutes.`; - case AgentTerminateMode.MAX_TURNS: - return `Agent reached max turns limit (${maxTurns}).`; - case AgentTerminateMode.ERROR_NO_COMPLETE_TASK_CALL: - return `Agent stopped calling tools but did not call '${TASK_COMPLETE_TOOL_NAME}'.`; - default: - return 'Agent execution was terminated before completion.'; - } - } } diff --git a/packages/core/src/agents/local-invocation.ts b/packages/core/src/agents/local-invocation.ts index fabc12fe87..bb2dcdea3d 100644 --- a/packages/core/src/agents/local-invocation.ts +++ b/packages/core/src/agents/local-invocation.ts @@ -164,14 +164,13 @@ ${output.result} parentPromptId: promptIdContext.getStore(), }); - const initialRequest = [{ text: 'Start' }]; // Placeholder for subagent start + const initialRequest = [{ text: 'Start' }]; const stream = harness.run(initialRequest, signal); let turn: Turn | undefined; while (true) { const { value, done } = await stream.next(); if (done) { - turn = value; break; } diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 6b6bdecfbc..34dabb4f9a 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -7,6 +7,7 @@ import type { GenerateContentConfig, PartListUnion, + Part, Content, Tool, GenerateContentResponse, @@ -64,6 +65,8 @@ import { resolveModel } from '../config/models.js'; import type { RetryAvailabilityContext } from '../utils/retry.js'; import { partToString } from '../utils/partUtils.js'; import { coreEvents, CoreEvent } from '../utils/events.js'; +import { AgentFactory } from '../agents/agent-factory.js'; +import { type AgentHarness } from '../agents/harness.js'; const MAX_TURNS = 100; @@ -90,6 +93,7 @@ export class GeminiClient { private currentSequenceModel: string | null = null; private lastSentIdeContext: IdeContext | undefined; private forceFullIdeContext = true; + private harness?: AgentHarness; /** * At any point in this conversation, was compression triggered without @@ -792,6 +796,46 @@ export class GeminiClient { this.config.resetTurn(); } + if (this.config.isAgentHarnessEnabled()) { + this.sessionTurnCount++; + if ( + this.config.getMaxSessionTurns() > 0 && + this.sessionTurnCount > this.config.getMaxSessionTurns() + ) { + yield { type: GeminiEventType.MaxSessionTurns }; + return new Turn(this.getChat(), prompt_id); + } + + if (!this.harness || this.lastPromptId !== prompt_id) { + this.harness = AgentFactory.createHarness(this.config, undefined, { + parentPromptId: prompt_id + }); + this.lastPromptId = prompt_id; + } + + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion, @typescript-eslint/no-explicit-any + const requestParts: Part[] = (Array.isArray(request) ? request : [{ text: partToString(request) }]) as any; + const stream = this.harness.run(requestParts, signal, turns); + + let turn: Turn | undefined; + while (true) { + const { value, done } = await stream.next(); + if (done) { + turn = value; + break; + } + yield value; + } + + if (turn) { + // Sync history back to GeminiClient's chat for transcript persistence + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion, @typescript-eslint/no-explicit-any + this.getChat().setHistory((turn as any).chat.getHistory()); + return turn; + } + return new Turn(this.getChat(), prompt_id); + } + const hooksEnabled = this.config.getEnableHooks(); const messageBus = this.config.getMessageBus();