From 9d8351c5c96fc13da441f506e07ec31d44d1cbd4 Mon Sep 17 00:00:00 2001 From: mkorwel Date: Thu, 12 Feb 2026 13:10:40 -0600 Subject: [PATCH] wip: fork agent harness UI and fix engine delegation loop --- packages/cli/src/ui/AppContainer.tsx | 55 ++-- .../cli/src/ui/components/Notifications.tsx | 3 +- .../cli/src/ui/contexts/UIStateContext.tsx | 2 +- packages/cli/src/ui/hooks/useAgentHarness.ts | 247 ++++++++++++++++++ packages/cli/src/ui/hooks/useGeminiStream.ts | 5 +- packages/cli/src/ui/types.ts | 1 + packages/core/src/agents/agent-scheduler.ts | 14 +- packages/core/src/agents/behavior.ts | 17 +- packages/core/src/agents/harness.ts | 126 +++++++-- packages/core/src/agents/local-invocation.ts | 21 +- packages/core/src/core/turn.ts | 8 +- packages/core/src/scheduler/scheduler.test.ts | 6 +- packages/core/src/scheduler/scheduler.ts | 22 +- packages/core/src/scheduler/state-manager.ts | 6 +- 14 files changed, 459 insertions(+), 74 deletions(-) create mode 100644 packages/cli/src/ui/hooks/useAgentHarness.ts diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 72fdb0ce48..21901596e0 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -95,6 +95,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 { useAgentHarness } from './hooks/useAgentHarness.js'; import { type BackgroundShell } from './hooks/shellCommandProcessor.js'; import { useVim } from './hooks/vim.js'; import { type LoadableSettingScope, SettingScope } from '../config/settings.js'; @@ -966,26 +967,9 @@ Logging in with Google... Restarting Gemini CLI to continue. } }, [pendingRestorePrompt, inputHistory, historyManager.history]); - const { - streamingState, - submitQuery, - initError, - pendingHistoryItems: pendingGeminiHistoryItems, - thought, - cancelOngoingRequest, - pendingToolCalls, - handleApprovalModeChange, - activePtyId, - loopDetectionConfirmationRequest, - lastOutputTime, - backgroundShellCount, - isBackgroundShellVisible, - toggleBackgroundShell, - backgroundCurrentShell, - backgroundShells, - dismissBackgroundShell, - retryStatus, - } = useGeminiStream( + const isAgentHarnessEnabled = config.isAgentHarnessEnabled(); + + const legacyStream = useGeminiStream( config.getGeminiClient(), historyManager.history, historyManager.addItem, @@ -1006,6 +990,37 @@ Logging in with Google... Restarting Gemini CLI to continue. embeddedShellFocused, ); + const harnessStream = useAgentHarness( + historyManager.addItem, + config, + onCancelSubmit, + ); + + const activeStream = isAgentHarnessEnabled ? harnessStream : legacyStream; + + const { + streamingState, + submitQuery, + initError, + pendingHistoryItems: pendingGeminiHistoryItems, + thought, + cancelOngoingRequest, + toolCalls: pendingToolCalls, + handleApprovalModeChange, + activePtyId: rawActivePtyId, + loopDetectionConfirmationRequest, + lastOutputTime, + backgroundShellCount, + isBackgroundShellVisible, + toggleBackgroundShell, + backgroundCurrentShell, + backgroundShells, + dismissBackgroundShell, + retryStatus, + } = activeStream; + + const activePtyId = rawActivePtyId ?? undefined; + toggleBackgroundShellRef.current = toggleBackgroundShell; isBackgroundShellVisibleRef.current = isBackgroundShellVisible; backgroundShellsRef.current = backgroundShells; diff --git a/packages/cli/src/ui/components/Notifications.tsx b/packages/cli/src/ui/components/Notifications.tsx index c252dd12de..5b80ad88e3 100644 --- a/packages/cli/src/ui/components/Notifications.tsx +++ b/packages/cli/src/ui/components/Notifications.tsx @@ -110,7 +110,8 @@ export const Notifications = () => { marginBottom={1} > - Initialization Error: {initError} + Initialization Error:{' '} + {initError instanceof Error ? initError.message : initError} {' '} diff --git a/packages/cli/src/ui/contexts/UIStateContext.tsx b/packages/cli/src/ui/contexts/UIStateContext.tsx index cf41896232..4888d0294b 100644 --- a/packages/cli/src/ui/contexts/UIStateContext.tsx +++ b/packages/cli/src/ui/contexts/UIStateContext.tsx @@ -98,7 +98,7 @@ export interface UIState { permissionConfirmationRequest: PermissionConfirmationRequest | null; geminiMdFileCount: number; streamingState: StreamingState; - initError: string | null; + initError: string | Error | null; pendingGeminiHistoryItems: HistoryItemWithoutId[]; thought: ThoughtSummary | null; shellModeActive: boolean; diff --git a/packages/cli/src/ui/hooks/useAgentHarness.ts b/packages/cli/src/ui/hooks/useAgentHarness.ts new file mode 100644 index 0000000000..042e1b1551 --- /dev/null +++ b/packages/cli/src/ui/hooks/useAgentHarness.ts @@ -0,0 +1,247 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { useState, useRef, useCallback, useEffect } from 'react'; +import { + GeminiEventType as ServerGeminiEventType, + debugLogger, + AgentFactory, +} from '@google/gemini-cli-core'; +import type { + Config, + ServerGeminiStreamEvent as GeminiEvent, + ThoughtSummary, +} from '@google/gemini-cli-core'; +import { type PartListUnion, type Part } from '@google/genai'; +import { + StreamingState, + MessageType, + type HistoryItemWithoutId, +} from '../types.js'; +import { useStateAndRef } from './useStateAndRef.js'; +import type { UseHistoryManagerReturn } from './useHistoryManager.js'; +import { mapToDisplay as mapTrackedToolCallsToDisplay } from './toolMapping.js'; +import { ROOT_SCHEDULER_ID } from '@google/gemini-cli-core'; +import type { TrackedToolCall } from './useToolScheduler.js'; + +export interface UseAgentHarnessReturn { + streamingState: StreamingState; + isResponding: boolean; + thought: ThoughtSummary | null; + streamingContent: string; + toolCalls: TrackedToolCall[]; + submitQuery: (query: PartListUnion) => Promise; + cancelOngoingRequest: () => void; + reset: () => void; + // Legacy compatibility properties + initError: Error | null; + pendingHistoryItems: HistoryItemWithoutId[]; + handleApprovalModeChange: (mode: string) => void; + activePtyId: number | null; + loopDetectionConfirmationRequest: unknown | null; + lastOutputTime: number; + backgroundShellCount: number; + isBackgroundShellVisible: boolean; + toggleBackgroundShell: () => void; + backgroundCurrentShell: unknown | null; + backgroundShells: Map; + dismissBackgroundShell: (pid: number) => void; + retryStatus: unknown | null; +} + +/** + * A specialized hook for processing streams from the AgentHarness. + * COMPLETELY FORKED from useGeminiStream to ensure zero regressions in legacy mode. + */ +export const useAgentHarness = ( + addItem: UseHistoryManagerReturn['addItem'], + config: Config, + onCancelSubmit: (fullReset: boolean) => void, +): UseAgentHarnessReturn => { + const [streamingState, setStreamingState] = useState( + StreamingState.Idle, + ); + const [streamingContent, setStreamingContent] = useState(''); + const [thought, thoughtRef, setThought] = useStateAndRef(null); + + // Tools for the CURRENT turn of the main agent + const [toolCalls, setToolCalls] = useState([]); + const toolCallsRef = useRef([]); + + // Sync ref with state + useEffect(() => { + toolCallsRef.current = toolCalls; + }, [toolCalls]); + + const pushedToolCallIdsRef = useRef>(new Set()); + + // Track subagent activities for hierarchical display + // For now, we just log them. In future, we can add a specialized "SubagentBlock" history item. + + const abortControllerRef = useRef(null); + + const reset = useCallback(() => { + setStreamingState(StreamingState.Idle); + setStreamingContent(''); + setThought(null); + setToolCalls([]); + pushedToolCallIdsRef.current.clear(); + }, [setThought]); + + const cancelOngoingRequest = useCallback(() => { + if (abortControllerRef.current) { + abortControllerRef.current.abort(); + } + onCancelSubmit(true); + reset(); + }, [onCancelSubmit, reset]); + + const processEvent = useCallback( + (event: GeminiEvent) => { + switch (event.type) { + case ServerGeminiEventType.Content: + setStreamingState(StreamingState.Responding); + setStreamingContent((prev) => prev + (event.value || '')); + break; + + case ServerGeminiEventType.Thought: + setThought(event.value); + break; + + case ServerGeminiEventType.ToolCallRequest: + setToolCalls((prev) => [ + ...prev, + { + request: event.value, + status: 'validating', + schedulerId: event.value.schedulerId || ROOT_SCHEDULER_ID, + } as TrackedToolCall, + ]); + break; + + case ServerGeminiEventType.ToolCallResponse: + { + const response = event.value; + setToolCalls((prev) => + prev.map((tc) => + tc.request.callId === response.callId + ? ({ + ...tc, + status: 'success', + result: response, + } as TrackedToolCall) + : tc, + ), + ); + } + break; + + case ServerGeminiEventType.TurnFinished: + // MAIN AGENT turn finished. Flush current state to history. + if (thoughtRef.current) { + addItem({ + type: MessageType.THINKING, + thought: thoughtRef.current, + } as HistoryItemWithoutId); + setThought(null); + } + + if (toolCallsRef.current.length > 0) { + const unpushed = toolCallsRef.current.filter( + (tc) => !pushedToolCallIdsRef.current.has(tc.request.callId), + ); + if (unpushed.length > 0) { + addItem( + mapTrackedToolCallsToDisplay(unpushed as TrackedToolCall[], { + borderBottom: true, + }), + ); + unpushed.forEach((tc) => + pushedToolCallIdsRef.current.add(tc.request.callId), + ); + } + } + + if (streamingContent) { + addItem({ type: 'gemini', text: streamingContent }); + setStreamingContent(''); + } + + setToolCalls([]); + break; + + case ServerGeminiEventType.SubagentActivity: + debugLogger.debug(`[HarnessHook] Subagent activity: ${event.value.type} for ${event.value.agentName}`); + break; + + case ServerGeminiEventType.Finished: + setStreamingState(StreamingState.Idle); + break; + + default: + break; + } + }, + [addItem, setThought, thoughtRef, streamingContent], + ); + + const submitQuery = useCallback( + async (parts: PartListUnion) => { + reset(); + setStreamingState(StreamingState.Responding); + + abortControllerRef.current = new AbortController(); + const harness = AgentFactory.createHarness(config); + + // Convert parts to Part[] array for harness + const requestParts: Part[] = Array.isArray(parts) + ? (parts as Part[]) + : [{ text: String(parts) }]; + + const stream = harness.run( + requestParts, + abortControllerRef.current.signal, + ); + + try { + for await (const event of stream) { + processEvent(event); + } + } catch (err: unknown) { + if (err instanceof Error && err.name === 'AbortError') return; + const msg = err instanceof Error ? err.message : String(err); + addItem({ type: MessageType.ERROR, text: msg }); + } finally { + setStreamingState(StreamingState.Idle); + } + }, + [config, reset, processEvent, addItem], + ); + + return { + streamingState, + isResponding: streamingState !== StreamingState.Idle, + thought, + streamingContent, + toolCalls, + submitQuery, + cancelOngoingRequest, + reset, + initError: null, + pendingHistoryItems: [], + handleApprovalModeChange: () => {}, + activePtyId: null, + loopDetectionConfirmationRequest: null, + lastOutputTime: 0, + backgroundShellCount: 0, + isBackgroundShellVisible: false, + toggleBackgroundShell: () => {}, + backgroundCurrentShell: null, + backgroundShells: new Map(), + dismissBackgroundShell: () => {}, + retryStatus: null, + }; +}; diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 74e6e635b9..c853fdd7ea 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -1171,6 +1171,9 @@ export const useGeminiStream = ( case ServerGeminiEventType.SubagentActivity: // TODO: UI implementation for subagent activity break; + case ServerGeminiEventType.TurnFinished: + // No-op for now to satisfy exhaustive switch + break; default: { // enforces exhaustive switch-case const unreachable: never = event; @@ -1680,7 +1683,7 @@ export const useGeminiStream = ( pendingHistoryItems, thought, cancelOngoingRequest, - pendingToolCalls: toolCalls, + toolCalls, handleApprovalModeChange, activePtyId, loopDetectionConfirmationRequest, diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index c48b81bf9c..90acf77688 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -373,6 +373,7 @@ export enum MessageType { AGENTS_LIST = 'agents_list', MCP_STATUS = 'mcp_status', CHAT_LIST = 'chat_list', + THINKING = 'thinking', HOOKS_LIST = 'hooks_list', } diff --git a/packages/core/src/agents/agent-scheduler.ts b/packages/core/src/agents/agent-scheduler.ts index 4b2e0fa587..19a3cad92f 100644 --- a/packages/core/src/agents/agent-scheduler.ts +++ b/packages/core/src/agents/agent-scheduler.ts @@ -12,6 +12,7 @@ import type { } from '../scheduler/types.js'; import type { ToolRegistry } from '../tools/tool-registry.js'; import type { EditorType } from '../utils/editor.js'; +import type { MessageBus } from '../confirmation-bus/message-bus.js'; /** * Options for scheduling agent tools. @@ -29,6 +30,13 @@ export interface AgentSchedulingOptions { getPreferredEditor?: () => EditorType | undefined; /** Optional function to be notified when the scheduler is waiting for user confirmation. */ onWaitingForConfirmation?: (waiting: boolean) => void; + /** + * Optional message bus override. + * If provided, the scheduler will broadcast to this bus. + * If explicitly null, broadcasting is disabled. + * If omitted, the global config message bus is used. + */ + messageBus?: MessageBus | null; } /** @@ -51,6 +59,7 @@ export async function scheduleAgentTools( signal, getPreferredEditor, onWaitingForConfirmation, + messageBus, } = options; // Create a proxy/override of the config to provide the agent-specific tool registry. @@ -59,7 +68,10 @@ export async function scheduleAgentTools( const scheduler = new Scheduler({ config: agentConfig, - messageBus: config.getMessageBus(), + messageBus: + messageBus === undefined + ? config.getMessageBus() + : (messageBus ?? undefined), getPreferredEditor: getPreferredEditor ?? (() => undefined), schedulerId, parentCallId, diff --git a/packages/core/src/agents/behavior.ts b/packages/core/src/agents/behavior.ts index 1ee995014d..add317db99 100644 --- a/packages/core/src/agents/behavior.ts +++ b/packages/core/src/agents/behavior.ts @@ -41,6 +41,7 @@ import { DeadlineTimer } from '../utils/deadlineTimer.js'; import { debugLogger } from '../utils/debugLogger.js'; import type { ToolRegistry } from '../tools/tool-registry.js'; +import type { ToolCallResponseInfo } from '../scheduler/types.js'; const TASK_COMPLETE_TOOL_NAME = 'complete_task'; const GRACE_PERIOD_MS = 60 * 1000; @@ -110,7 +111,13 @@ export interface AgentBehavior { * 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; + isGoalReached( + toolResults: Array<{ + name: string; + part: Part; + result: ToolCallResponseInfo; + }>, + ): boolean; /** * Checks if the agent should continue executing after a model turn with no tool calls. @@ -467,7 +474,13 @@ export class SubagentBehavior implements AgentBehavior { return request; } - isGoalReached(toolResults: Array<{ name: string; part: Part }>) { + isGoalReached( + toolResults: Array<{ + name: string; + part: Part; + result: ToolCallResponseInfo; + }>, + ) { const completeCall = toolResults.find( (r) => r.name === TASK_COMPLETE_TOOL_NAME, ); diff --git a/packages/core/src/agents/harness.ts b/packages/core/src/agents/harness.ts index 334d183b6c..7ff8a782ec 100644 --- a/packages/core/src/agents/harness.ts +++ b/packages/core/src/agents/harness.ts @@ -29,8 +29,12 @@ import { ToolOutputMaskingService } from '../services/toolOutputMaskingService.j import { resolveModel } from '../config/models.js'; import { type RoutingContext } from '../routing/routingStrategy.js'; import { ToolRegistry } from '../tools/tool-registry.js'; +import { SubagentTool } from './subagent-tool.js'; import { scheduleAgentTools } from './agent-scheduler.js'; -import { type ToolCallRequestInfo } from '../scheduler/types.js'; +import { + type ToolCallRequestInfo, + type ToolCallResponseInfo, +} from '../scheduler/types.js'; import { promptIdContext } from '../utils/promptIdContext.js'; import { logAgentStart, logAgentFinish } from '../telemetry/loggers.js'; import { AgentStartEvent, AgentFinishEvent } from '../telemetry/types.js'; @@ -258,14 +262,17 @@ export class AgentHarness { } if (event.type === GeminiEventType.ToolCallRequest) { - yield { - type: GeminiEventType.SubagentActivity, - value: { - agentName: this.behavior.name, - type: 'TOOL_CALL_START', - data: { name: event.value.name, args: event.value.args }, - }, - }; + const tool = this.toolRegistry.getTool(event.value.name); + if (tool instanceof SubagentTool) { + yield { + type: GeminiEventType.SubagentActivity, + value: { + agentName: this.behavior.name, + type: 'TOOL_CALL_START', + data: { name: event.value.name, args: event.value.args }, + }, + }; + } } } @@ -290,7 +297,9 @@ export class AgentHarness { if (afterResult.shouldContinue) { currentRequest = [{ text: afterResult.reason || 'Continue' }]; this.turnCounter++; - turn = new Turn(this.chat!, this.behavior.agentId); + if (this.behavior.name === 'main') { + yield { type: GeminiEventType.TurnFinished }; + } continue; } @@ -309,14 +318,71 @@ export class AgentHarness { onWaitingForConfirmation, ); + // Yield responses so UI knows they are done + for (const result of toolResults) { + if (result.result) { + yield { + type: GeminiEventType.ToolCallResponse, + value: result.result, + }; + + const tool = this.toolRegistry.getTool(result.name); + if (tool instanceof SubagentTool) { + yield { + type: GeminiEventType.SubagentActivity, + value: { + agentName: this.behavior.name, + type: 'TOOL_CALL_END', + data: { + name: result.name, + output: result.result.resultDisplay, + }, + }, + }; + } + } + } + if (this.behavior.isGoalReached(toolResults)) { terminateReason = AgentTerminateMode.GOAL; + + // If it's a subagent, find the complete_task call and extract the result string + const goalCall = toolResults.find((r) => r.name === 'complete_task'); + if (goalCall?.result?.resultDisplay && this.chat) { + // Ensure the chat session records the final text result so future turns or getResponseText() can see it + this.chat.addHistory({ + role: 'model', + parts: [{ text: String(goalCall.result.resultDisplay) }], + }); + } + return turn; } - currentRequest = toolResults.map((r) => r.part); + currentRequest = toolResults.map((r) => { + // Ensure the LLM "sees" the rich result display if it's available. + // We use the resultDisplay text as the definitive function response. + if ( + r.result?.resultDisplay && + 'functionResponse' in r.part && + r.part.functionResponse + ) { + return { + functionResponse: { + ...r.part.functionResponse, + response: { result: String(r.result.resultDisplay) }, + }, + }; + } + return r.part; + }); this.turnCounter++; - turn = new Turn(this.chat!, this.behavior.agentId); + + // Only yield TurnFinished if we are the main agent. + // Nested subagent turns should be internal and not trigger UI flushes in the parent. + if (this.behavior.name === 'main') { + yield { type: GeminiEventType.TurnFinished }; + } } else { // No tool calls. Check for continuation. const nextParts = await this.behavior.getContinuationRequest( @@ -326,7 +392,9 @@ export class AgentHarness { if (nextParts) { currentRequest = nextParts; this.turnCounter++; - turn = new Turn(this.chat!, this.behavior.agentId); + if (this.behavior.name === 'main') { + yield { type: GeminiEventType.TurnFinished }; + } continue; } @@ -428,7 +496,7 @@ export class AgentHarness { calls: ToolCallRequestInfo[], signal: AbortSignal, onWaitingForConfirmation?: (waiting: boolean) => void, - ): Promise> { + ): Promise> { const taskCompleteCalls = calls.filter( (c) => c.name === TASK_COMPLETE_TOOL_NAME, ); @@ -440,7 +508,7 @@ export class AgentHarness { let completedCalls: Array<{ request: ToolCallRequestInfo; - response: { responseParts: Part[] }; + response: ToolCallResponseInfo; }> = []; if (otherCalls.length > 0) { @@ -449,24 +517,38 @@ export class AgentHarness { toolRegistry: this.toolRegistry, signal, onWaitingForConfirmation, + // Only broadcast to global UI if we are the main top-level agent + messageBus: this.behavior.name === 'main' ? undefined : null, }); } const results = completedCalls.map((call) => ({ name: call.request.name, part: call.response.responseParts[0], + result: call.response, })); for (const call of taskCompleteCalls) { + const response: ToolCallResponseInfo = { + callId: call.callId, + responseParts: [ + { + functionResponse: { + name: TASK_COMPLETE_TOOL_NAME, + response: { result: 'Task completed locally' }, + id: call.callId, + }, + }, + ], + resultDisplay: 'Task completed locally', + error: undefined, + errorType: undefined, + contentLength: 'Task completed locally'.length, + }; results.push({ name: TASK_COMPLETE_TOOL_NAME, - part: { - functionResponse: { - name: TASK_COMPLETE_TOOL_NAME, - response: { result: 'Task completed locally' }, - id: call.callId, - }, - }, + part: response.responseParts[0], + result: response, }); } diff --git a/packages/core/src/agents/local-invocation.ts b/packages/core/src/agents/local-invocation.ts index bb2dcdea3d..7cbdad4168 100644 --- a/packages/core/src/agents/local-invocation.ts +++ b/packages/core/src/agents/local-invocation.ts @@ -7,7 +7,10 @@ import type { Config } from '../config/config.js'; import { LocalAgentExecutor } from './local-executor.js'; import type { AnsiOutput } from '../utils/terminalSerializer.js'; -import { BaseToolInvocation, type ToolResult } from '../tools/tools.js'; +import { + BaseToolInvocation, + type ToolResult, +} from '../tools/tools.js'; import { ToolErrorType } from '../tools/tool-error.js'; import type { LocalAgentDefinition, @@ -17,7 +20,6 @@ import type { import type { MessageBus } from '../confirmation-bus/message-bus.js'; import { AgentFactory } from './agent-factory.js'; import type { Turn } from '../core/turn.js'; -import { GeminiEventType } from '../core/turn.js'; import { promptIdContext } from '../utils/promptIdContext.js'; const INPUT_PREVIEW_MAX_LENGTH = 50; @@ -168,24 +170,15 @@ ${output.result} const stream = harness.run(initialRequest, signal); let turn: Turn | undefined; + while (true) { const { value, done } = await stream.next(); if (done) { turn = value; break; } - const event = value; - if (updateOutput) { - if (event.type === GeminiEventType.Thought) { - updateOutput(`🤖💭 ${event.value.subject}`); - } else if (event.type === GeminiEventType.SubagentActivity) { - if (event.value.type === 'TOOL_CALL_START') { - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - const toolName = event.value.data['name'] as string; - updateOutput(`🛠️ Calling tool: ${toolName}...`); - } - } - } + // For the subagent box, we don't want to stream internal thoughts or tool calls + // to the persistent history. We just wait for the final result. } if (!turn) { diff --git a/packages/core/src/core/turn.ts b/packages/core/src/core/turn.ts index baef5596e4..147fc9f771 100644 --- a/packages/core/src/core/turn.ts +++ b/packages/core/src/core/turn.ts @@ -69,6 +69,7 @@ export enum GeminiEventType { AgentExecutionStopped = 'agent_execution_stopped', AgentExecutionBlocked = 'agent_execution_blocked', SubagentActivity = 'subagent_activity', + TurnFinished = 'turn_finished', } export type ServerGeminiSubagentActivityEvent = { @@ -80,6 +81,10 @@ export type ServerGeminiSubagentActivityEvent = { }; }; +export type ServerGeminiTurnFinishedEvent = { + type: GeminiEventType.TurnFinished; +}; + export type ServerGeminiRetryEvent = { type: GeminiEventType.Retry; }; @@ -240,7 +245,8 @@ export type ServerGeminiStreamEvent = | ServerGeminiModelInfoEvent | ServerGeminiAgentExecutionStoppedEvent | ServerGeminiAgentExecutionBlockedEvent - | ServerGeminiSubagentActivityEvent; + | ServerGeminiSubagentActivityEvent + | ServerGeminiTurnFinishedEvent; // A turn manages the agentic loop turn within the server context. export class Turn { diff --git a/packages/core/src/scheduler/scheduler.test.ts b/packages/core/src/scheduler/scheduler.test.ts index a3979f43a6..56d911bfc4 100644 --- a/packages/core/src/scheduler/scheduler.test.ts +++ b/packages/core/src/scheduler/scheduler.test.ts @@ -219,7 +219,11 @@ describe('Scheduler (Orchestrator)', () => { let capturedTerminalHandler: TerminalCallHandler | undefined; vi.mocked(SchedulerStateManager).mockImplementation( - (_messageBus, _schedulerId, onTerminalCall) => { + ( + _messageBus: MessageBus | undefined, + _schedulerId: string | undefined, + onTerminalCall: TerminalCallHandler | undefined, + ) => { capturedTerminalHandler = onTerminalCall; return mockStateManager as unknown as SchedulerStateManager; }, diff --git a/packages/core/src/scheduler/scheduler.ts b/packages/core/src/scheduler/scheduler.ts index 1cd8dc3317..c3152b0acb 100644 --- a/packages/core/src/scheduler/scheduler.ts +++ b/packages/core/src/scheduler/scheduler.ts @@ -47,7 +47,7 @@ interface SchedulerQueueItem { export interface SchedulerOptions { config: Config; - messageBus: MessageBus; + messageBus?: MessageBus; getPreferredEditor: () => EditorType | undefined; schedulerId: string; parentCallId?: string; @@ -87,7 +87,7 @@ export class Scheduler { private readonly executor: ToolExecutor; private readonly modifier: ToolModificationHandler; private readonly config: Config; - private readonly messageBus: MessageBus; + private readonly messageBus?: MessageBus; private readonly getPreferredEditor: () => EditorType | undefined; private readonly schedulerId: string; private readonly parentCallId?: string; @@ -112,11 +112,13 @@ export class Scheduler { this.executor = new ToolExecutor(this.config); this.modifier = new ToolModificationHandler(); - this.setupMessageBusListener(this.messageBus); + if (this.messageBus) { + this.setupMessageBusListener(this.messageBus); + } } private setupMessageBusListener(messageBus: MessageBus): void { - if (Scheduler.subscribedMessageBuses.has(messageBus)) { + if (!messageBus || Scheduler.subscribedMessageBuses.has(messageBus)) { return; } @@ -432,7 +434,7 @@ export class Scheduler { let outcome = ToolConfirmationOutcome.ProceedOnce; let lastDetails: SerializableConfirmationDetails | undefined; - if (decision === PolicyDecision.ASK_USER) { + if (decision === PolicyDecision.ASK_USER && this.messageBus) { const result = await resolveConfirmation(toolCall, signal, { config: this.config, messageBus: this.messageBus, @@ -449,10 +451,12 @@ export class Scheduler { } // Handle Policy Updates - await updatePolicy(toolCall.tool, outcome, lastDetails, { - config: this.config, - messageBus: this.messageBus, - }); + if (this.messageBus) { + await updatePolicy(toolCall.tool, outcome, lastDetails, { + config: this.config, + messageBus: this.messageBus, + }); + } // Handle cancellation (cascades to entire batch) if (outcome === ToolConfirmationOutcome.Cancel) { diff --git a/packages/core/src/scheduler/state-manager.ts b/packages/core/src/scheduler/state-manager.ts index 21e931a18a..bf6fc33502 100644 --- a/packages/core/src/scheduler/state-manager.ts +++ b/packages/core/src/scheduler/state-manager.ts @@ -46,7 +46,7 @@ export class SchedulerStateManager { private _completedBatch: CompletedToolCall[] = []; constructor( - private readonly messageBus: MessageBus, + private readonly messageBus: MessageBus | undefined, private readonly schedulerId: string = ROOT_SCHEDULER_ID, private readonly onTerminalCall?: TerminalCallHandler, ) {} @@ -210,6 +210,10 @@ export class SchedulerStateManager { } private emitUpdate() { + if (!this.messageBus) { + return; + } + const snapshot = this.getSnapshot(); // Fire and forget - The message bus handles the publish and error handling.