diff --git a/packages/cli/src/nonInteractiveCliAgentSession.ts b/packages/cli/src/nonInteractiveCliAgentSession.ts index e0a532becf5..e64b4a5686f 100644 --- a/packages/cli/src/nonInteractiveCliAgentSession.ts +++ b/packages/cli/src/nonInteractiveCliAgentSession.ts @@ -193,6 +193,7 @@ export async function runNonInteractive({ let errorToHandle: unknown | undefined; let scheduler: Scheduler | undefined; + let session: LegacyAgentSession | undefined; let abortSession = () => {}; try { consolePatcher.patch(); @@ -296,7 +297,7 @@ export async function runNonInteractive({ } // Create LegacyAgentSession — owns the agentic loop - const session = new LegacyAgentSession({ + session = new LegacyAgentSession({ client: geminiClient, scheduler, config, @@ -305,7 +306,7 @@ export async function runNonInteractive({ // Wire Ctrl+C to session abort abortSession = () => { - void session.abort(); + void session?.abort(); }; abortController.signal.addEventListener('abort', abortSession); if (abortController.signal.aborted) { @@ -640,6 +641,7 @@ export async function runNonInteractive({ cleanupStdinCancellation(); abortController.signal.removeEventListener('abort', abortSession); + session?.dispose(); scheduler?.dispose(); consolePatcher.cleanup(); coreEvents.off(CoreEvent.UserFeedback, handleUserFeedback); diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 313f377b02b..4c1fc2fae72 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -1180,6 +1180,13 @@ Logging in with Google... Restarting Gemini CLI to continue. [config, getPreferredEditor], ); + useEffect( + () => () => { + streamAgent?.dispose?.(); + }, + [streamAgent], + ); + const activeStream = streamAgent ? // eslint-disable-next-line react-hooks/rules-of-hooks useAgentStream({ diff --git a/packages/cli/src/ui/hooks/useAgentStream.ts b/packages/cli/src/ui/hooks/useAgentStream.ts index 982391a4370..6d46d2f8e60 100644 --- a/packages/cli/src/ui/hooks/useAgentStream.ts +++ b/packages/cli/src/ui/hooks/useAgentStream.ts @@ -224,14 +224,17 @@ export const useAgentStream = ({ if (tc.callId !== event.requestId) return tc; const legacyState = event._meta?.legacyState; - const evtStatus = legacyState?.status; - let status = tc.status; - if (evtStatus === 'executing') + if (event.status === 'executing') status = CoreToolCallStatus.Executing; - else if (evtStatus === 'error') status = CoreToolCallStatus.Error; - else if (evtStatus === 'success') + else if (event.status === 'pending_input') + status = CoreToolCallStatus.AwaitingApproval; + else if (event.status === 'errored') + status = CoreToolCallStatus.Error; + else if (event.status === 'succeeded') status = CoreToolCallStatus.Success; + else if (event.status === 'aborted') + status = CoreToolCallStatus.Cancelled; const display = event.display?.result; const liveOutput = @@ -272,11 +275,16 @@ export const useAgentStream = ({ const resultDisplay = displayContentToString(display) ?? tc.resultDisplay; + let status = CoreToolCallStatus.Success; + if (event.status === 'errored') status = CoreToolCallStatus.Error; + else if (event.status === 'aborted') + status = CoreToolCallStatus.Cancelled; + else if (event.status === 'succeeded') + status = CoreToolCallStatus.Success; + return { ...tc, - status: event.isError - ? CoreToolCallStatus.Error - : CoreToolCallStatus.Success, + status, display: event.display ? { ...tc.display, ...event.display } : tc.display, diff --git a/packages/core/src/agent/agent-session.ts b/packages/core/src/agent/agent-session.ts index 6a4c295fc86..e07695532db 100644 --- a/packages/core/src/agent/agent-session.ts +++ b/packages/core/src/agent/agent-session.ts @@ -34,6 +34,10 @@ export class AgentSession implements AgentProtocol { return this._protocol.abort(); } + dispose(): void { + this._protocol.dispose?.(); + } + get events(): readonly AgentEvent[] { return this._protocol.events; } diff --git a/packages/core/src/agent/event-translator.ts b/packages/core/src/agent/event-translator.ts index f60822a8e6f..f91b3e6ddd7 100644 --- a/packages/core/src/agent/event-translator.ts +++ b/packages/core/src/agent/event-translator.ts @@ -235,6 +235,7 @@ export function translateEvent( makeEvent('tool_request', state, { requestId: event.value.callId, name: event.value.name, + status: 'pending', args: event.value.args, display: event.value.display, }), @@ -257,6 +258,7 @@ export function translateEvent( makeEvent('tool_response', state, { requestId: event.value.callId, name: state.pendingToolNames.get(event.value.callId) ?? 'unknown', + status: event.value.error ? 'errored' : 'succeeded', content: event.value.error ? [{ type: 'text', text: event.value.error.message }] : geminiPartsToContentParts(event.value.responseParts), diff --git a/packages/core/src/agent/legacy-agent-session.ts b/packages/core/src/agent/legacy-agent-session.ts index e8d5e56ef5c..587c107c696 100644 --- a/packages/core/src/agent/legacy-agent-session.ts +++ b/packages/core/src/agent/legacy-agent-session.ts @@ -13,7 +13,10 @@ import { GeminiEventType } from '../core/turn.js'; import type { Part } from '@google/genai'; import type { GeminiClient } from '../core/client.js'; import type { Config } from '../config/config.js'; -import type { ToolCallRequestInfo } from '../scheduler/types.js'; +import { + type ToolCallRequestInfo, + CoreToolCallStatus, +} from '../scheduler/types.js'; import { Scheduler } from '../scheduler/scheduler.js'; import { recordToolCallInteractions } from '../code_assist/telemetry.js'; import { ToolErrorType, isFatalToolError } from '../tools/tool-error.js'; @@ -39,7 +42,12 @@ import type { ContentPart, StreamEndReason, Unsubscribe, + ToolEventStatus, } from './types.js'; +import { + MessageBusType, + type ToolCallsUpdateMessage, +} from '../confirmation-bus/types.js'; function isAbortLikeError(err: unknown): boolean { return err instanceof Error && err.name === 'AbortError'; @@ -63,7 +71,9 @@ export class LegacyAgentProtocol implements AgentProtocol { private _agentEndEmitted = false; private _activeStreamId?: string; private _abortController = new AbortController(); + private _disposalController = new AbortController(); private _nextStreamIdOverride?: string; + private _lastToolStatuses = new Map(); private readonly _client: GeminiClient; private readonly _scheduler: Scheduler; @@ -92,6 +102,19 @@ export class LegacyAgentProtocol implements AgentProtocol { } this._scheduler = scheduler; } + + if (this._config.messageBus) { + this._config.messageBus.subscribe( + MessageBusType.TOOL_CALLS_UPDATE, + this._handleToolCallsUpdate.bind(this), + { signal: this._disposalController.signal }, + ); + } + } + + dispose(): void { + this._disposalController.abort(); + void this.abort(); } get events(): readonly AgentEvent[] { @@ -275,6 +298,11 @@ export class LegacyAgentProtocol implements AgentProtocol { this._makeToolResponseEvent({ requestId: request.callId, name: request.name, + status: response.error + ? 'errored' + : tc.status === CoreToolCallStatus.Cancelled + ? 'aborted' + : 'succeeded', content, isError: response.error !== undefined, ...(display ? { display } : {}), @@ -489,6 +517,87 @@ export class LegacyAgentProtocol implements AgentProtocol { } satisfies AgentEvent<'error'>; return event; } + + private _handleToolCallsUpdate(msg: ToolCallsUpdateMessage): void { + if (!this._activeStreamId) { + return; + } + + const eventsToEmit: AgentEvent[] = []; + + for (const tc of msg.toolCalls) { + const callId = tc.request.callId; + + if (!this._translationState.pendingToolNames.has(callId)) { + continue; + } + + let status: ToolEventStatus = 'pending'; + switch (tc.status) { + case CoreToolCallStatus.Validating: + case CoreToolCallStatus.Scheduled: + status = 'pending'; + break; + case CoreToolCallStatus.AwaitingApproval: + status = 'pending_input'; + break; + case CoreToolCallStatus.Executing: + status = 'executing'; + break; + case CoreToolCallStatus.Success: + status = 'succeeded'; + break; + case CoreToolCallStatus.Error: + status = 'errored'; + break; + case CoreToolCallStatus.Cancelled: + status = 'aborted'; + break; + default: + status = 'pending'; + break; + } + + const lastStatus = this._lastToolStatuses.get(callId); + + if (lastStatus !== status) { + this._lastToolStatuses.set(callId, status); + + const display = populateToolDisplay({ + name: tc.request.name, + invocation: 'invocation' in tc ? tc.invocation : undefined, + displayName: 'tool' in tc ? tc.tool?.displayName : undefined, + display: 'response' in tc ? tc.response?.display : undefined, + }); + + eventsToEmit.push( + this._makeToolUpdateEvent({ + requestId: callId, + status, + ...(display ? { display } : {}), + }), + ); + } + } + + if (eventsToEmit.length > 0) { + this._emit(eventsToEmit); + } + } + + private _makeToolUpdateEvent( + payload: Omit< + AgentEvent<'tool_update'>, + 'id' | 'timestamp' | 'streamId' | 'type' + >, + ): AgentEvent<'tool_update'> { + const event = { + ...this._nextEventFields(), + type: 'tool_update', + ...payload, + } satisfies AgentEvent<'tool_update'>; + return event; + } } export class LegacyAgentSession extends AgentSession { diff --git a/packages/core/src/agent/types.ts b/packages/core/src/agent/types.ts index 0d41c466024..9595135da38 100644 --- a/packages/core/src/agent/types.ts +++ b/packages/core/src/agent/types.ts @@ -37,6 +37,11 @@ export interface AgentProtocol extends Trajectory { */ abort(): Promise; + /** + * Disposes of the protocol, cleaning up any long-lived resources. + */ + dispose?(): void; + /** * AgentProtocol implements the Trajectory interface and can retrieve existing events. */ @@ -227,11 +232,21 @@ export interface ToolDisplay { format?: ToolDisplayFormat; } +export type ToolEventStatus = + | 'pending' + | 'pending_input' + | 'executing' + | 'succeeded' + | 'errored' + | 'aborted'; + export interface ToolRequest { /** A unique identifier for this tool request to be correlated by the response. */ requestId: string; /** The name of the tool being requested. */ name: string; + /** The status of the tool execution. */ + status: ToolEventStatus; /** The arguments for the tool. */ /** Tool-controlled display information. */ display?: ToolDisplay; @@ -255,6 +270,8 @@ export interface ToolRequest { */ export interface ToolUpdate { requestId: string; + /** The status of the tool execution. */ + status: ToolEventStatus; /** Tool-controlled display information. */ display?: ToolDisplay; content?: ContentPart[]; @@ -276,6 +293,8 @@ export interface ToolUpdate { export interface ToolResponse { requestId: string; name: string; + /** The status of the tool execution. */ + status: ToolEventStatus; /** Tool-controlled display information. */ display?: ToolDisplay; /** Multi-part content to be sent to the model. */ diff --git a/packages/core/src/agents/local-subagent-protocol.ts b/packages/core/src/agents/local-subagent-protocol.ts index 57c6121ea34..1328a072f38 100644 --- a/packages/core/src/agents/local-subagent-protocol.ts +++ b/packages/core/src/agents/local-subagent-protocol.ts @@ -277,6 +277,7 @@ class LocalSubagentProtocol implements AgentProtocol { this._makeEvent('tool_request', { requestId: callId, name, + status: 'executing', args, }), ]; @@ -292,6 +293,7 @@ class LocalSubagentProtocol implements AgentProtocol { this._makeEvent('tool_response', { requestId, name, + status: 'succeeded', content: [{ type: 'text', text: output }], }), ]; diff --git a/packages/core/src/tools/shell.ts b/packages/core/src/tools/shell.ts index 13965d94f6e..cce401cdf5d 100644 --- a/packages/core/src/tools/shell.ts +++ b/packages/core/src/tools/shell.ts @@ -1036,7 +1036,7 @@ export class ShellToolInvocation extends BaseToolInvocation< return { llmContent, display: { - name: 'Shell', + name: this._toolDisplayName, description: this.getDescription(), resultSummary: displayResultSummary, result: