mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-07-22 07:41:23 -07:00
feat(core): add LocalSessionInvocation (#26665)
This commit is contained in:
Adam Weidman
@@ -0,0 +1,416 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { type AgentLoopContext } from '../config/agent-loop-context.js';
|
||||
import { MessageBusType } from '../confirmation-bus/types.js';
|
||||
import {
|
||||
BaseToolInvocation,
|
||||
type ToolResult,
|
||||
type ExecuteOptions,
|
||||
} from '../tools/tools.js';
|
||||
import {
|
||||
type LocalAgentDefinition,
|
||||
type AgentInputs,
|
||||
type SubagentActivityEvent,
|
||||
type SubagentProgress,
|
||||
type SubagentActivityItem,
|
||||
AgentTerminateMode,
|
||||
SubagentActivityErrorType,
|
||||
SUBAGENT_REJECTED_ERROR_PREFIX,
|
||||
SUBAGENT_CANCELLED_ERROR_MESSAGE,
|
||||
isToolActivityError,
|
||||
SubagentState,
|
||||
} from './types.js';
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import type { MessageBus } from '../confirmation-bus/message-bus.js';
|
||||
import {
|
||||
sanitizeThoughtContent,
|
||||
sanitizeToolArgs,
|
||||
sanitizeErrorMessage,
|
||||
} from '../utils/agent-sanitization-utils.js';
|
||||
import { checkExhaustive } from '../utils/checks.js';
|
||||
import { LocalSubagentSession } from './local-subagent-protocol.js';
|
||||
import type { AgentEvent } from '../agent/types.js';
|
||||
|
||||
const MAX_RECENT_ACTIVITY = 3;
|
||||
|
||||
/** Optional configuration for subagent invocations. */
|
||||
export interface SubagentInvocationOptions {
|
||||
toolName?: string;
|
||||
toolDisplayName?: string;
|
||||
onAgentEvent?: (event: AgentEvent) => void;
|
||||
}
|
||||
|
||||
/**
|
||||
* Session-based local subagent invocation.
|
||||
*
|
||||
* This class orchestrates the execution of a defined agent by:
|
||||
* 1. Using {@link LocalSubagentSession} as the execution engine.
|
||||
* 2. Bridging the agent's streaming activity (e.g., thoughts) to the tool's
|
||||
* live output stream via the session's rawActivityCallback.
|
||||
* 3. Formatting the final result into a {@link ToolResult}.
|
||||
*/
|
||||
export class LocalSessionInvocation extends BaseToolInvocation<
|
||||
AgentInputs,
|
||||
ToolResult
|
||||
> {
|
||||
private readonly _onAgentEvent?: (event: AgentEvent) => void;
|
||||
|
||||
/**
|
||||
* @param definition The definition object that configures the agent.
|
||||
* @param context The agent loop context.
|
||||
* @param params The validated input parameters for the agent.
|
||||
* @param messageBus Message bus for policy enforcement.
|
||||
* @param options Optional overrides for tool name, display name, and event callback.
|
||||
*/
|
||||
constructor(
|
||||
private readonly definition: LocalAgentDefinition,
|
||||
private readonly context: AgentLoopContext,
|
||||
params: AgentInputs,
|
||||
messageBus: MessageBus,
|
||||
options?: SubagentInvocationOptions,
|
||||
) {
|
||||
super(
|
||||
params,
|
||||
messageBus,
|
||||
options?.toolName ?? definition.name,
|
||||
options?.toolDisplayName ?? definition.displayName,
|
||||
);
|
||||
this._onAgentEvent = options?.onAgentEvent;
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns a concise, human-readable description of the invocation.
|
||||
* Used for logging and display purposes.
|
||||
*/
|
||||
getDescription(): string {
|
||||
const inputSummary = Object.entries(this.params)
|
||||
.map(([key, value]) => `${key}: ${String(value)}`)
|
||||
.join(', ');
|
||||
|
||||
return `Running subagent '${this.definition.name}' with inputs: { ${inputSummary} }`;
|
||||
}
|
||||
|
||||
private publishActivity(activity: SubagentActivityItem): void {
|
||||
void this.messageBus.publish({
|
||||
type: MessageBusType.SUBAGENT_ACTIVITY,
|
||||
subagentName: this.definition.displayName ?? this.definition.name,
|
||||
activity,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Executes the subagent.
|
||||
*
|
||||
* @param options Options for tool execution including signal and output updates.
|
||||
* @returns A `Promise` that resolves with the final `ToolResult`.
|
||||
*/
|
||||
async execute(options: ExecuteOptions): Promise<ToolResult> {
|
||||
const { abortSignal: signal, updateOutput } = options;
|
||||
let recentActivity: SubagentActivityItem[] = [];
|
||||
|
||||
// Raw SubagentActivityEvent handler — preserves all existing progress display logic.
|
||||
// Passed as rawActivityCallback to LocalSubagentSession so the protocol can call it
|
||||
// before translating to AgentEvents.
|
||||
const onActivity = (activity: SubagentActivityEvent): void => {
|
||||
if (!updateOutput) return;
|
||||
|
||||
let updated = false;
|
||||
|
||||
switch (activity.type) {
|
||||
case 'THOUGHT_CHUNK': {
|
||||
const rawText = activity.data['text'];
|
||||
const text = typeof rawText === 'string' ? rawText : '';
|
||||
const lastItem = recentActivity[recentActivity.length - 1];
|
||||
|
||||
if (
|
||||
lastItem &&
|
||||
lastItem.type === 'thought' &&
|
||||
lastItem.status === SubagentState.RUNNING
|
||||
) {
|
||||
lastItem.content = sanitizeThoughtContent(text);
|
||||
} else {
|
||||
recentActivity.push({
|
||||
id: randomUUID(),
|
||||
type: 'thought',
|
||||
content: sanitizeThoughtContent(text),
|
||||
status: SubagentState.RUNNING,
|
||||
});
|
||||
}
|
||||
updated = true;
|
||||
|
||||
const latestThought = recentActivity[recentActivity.length - 1];
|
||||
if (latestThought) {
|
||||
this.publishActivity(latestThought);
|
||||
}
|
||||
break;
|
||||
}
|
||||
case 'TOOL_CALL_START': {
|
||||
const rawName = activity.data['name'];
|
||||
const name = typeof rawName === 'string' ? rawName.trim() : '';
|
||||
const displayName = activity.data['displayName']
|
||||
? sanitizeErrorMessage(String(activity.data['displayName']).trim())
|
||||
: undefined;
|
||||
const description = activity.data['description']
|
||||
? sanitizeErrorMessage(String(activity.data['description']))
|
||||
: undefined;
|
||||
const args = JSON.stringify(sanitizeToolArgs(activity.data['args']));
|
||||
const callId = activity.data['callId']
|
||||
? String(activity.data['callId'])
|
||||
: randomUUID();
|
||||
recentActivity.push({
|
||||
id: callId,
|
||||
type: 'tool_call',
|
||||
content: name,
|
||||
displayName,
|
||||
description,
|
||||
args,
|
||||
status: SubagentState.RUNNING,
|
||||
});
|
||||
updated = true;
|
||||
|
||||
const latestTool = recentActivity[recentActivity.length - 1];
|
||||
if (latestTool) {
|
||||
this.publishActivity(latestTool);
|
||||
}
|
||||
break;
|
||||
}
|
||||
case 'TOOL_CALL_END': {
|
||||
const data = activity.data['data'];
|
||||
const isError = isToolActivityError(data);
|
||||
|
||||
const callId = activity.data['id']
|
||||
? String(activity.data['id'])
|
||||
: undefined;
|
||||
|
||||
if (callId) {
|
||||
for (let i = recentActivity.length - 1; i >= 0; i--) {
|
||||
if (
|
||||
recentActivity[i].type === 'tool_call' &&
|
||||
recentActivity[i].id === callId &&
|
||||
recentActivity[i].status === SubagentState.RUNNING
|
||||
) {
|
||||
recentActivity[i].status = isError
|
||||
? SubagentState.ERROR
|
||||
: SubagentState.COMPLETED;
|
||||
updated = true;
|
||||
|
||||
this.publishActivity(recentActivity[i]);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
case 'ERROR': {
|
||||
const rawError = activity.data['error'];
|
||||
const error = typeof rawError === 'string' ? rawError.trim() : '';
|
||||
const errorType = activity.data['errorType'];
|
||||
const sanitizedError = sanitizeErrorMessage(error);
|
||||
const isCancellation =
|
||||
errorType === SubagentActivityErrorType.CANCELLED ||
|
||||
error === SUBAGENT_CANCELLED_ERROR_MESSAGE;
|
||||
const isRejection =
|
||||
errorType === SubagentActivityErrorType.REJECTED ||
|
||||
error.startsWith(SUBAGENT_REJECTED_ERROR_PREFIX);
|
||||
|
||||
const callId = activity.data['callId']
|
||||
? String(activity.data['callId'])
|
||||
: undefined;
|
||||
|
||||
if (callId) {
|
||||
const targetStatus =
|
||||
isCancellation || isRejection
|
||||
? SubagentState.CANCELLED
|
||||
: SubagentState.ERROR;
|
||||
|
||||
for (let i = recentActivity.length - 1; i >= 0; i--) {
|
||||
if (
|
||||
recentActivity[i].type === 'tool_call' &&
|
||||
recentActivity[i].id === callId &&
|
||||
recentActivity[i].status === SubagentState.RUNNING
|
||||
) {
|
||||
recentActivity[i].status = targetStatus;
|
||||
updated = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
recentActivity.push({
|
||||
id: randomUUID(),
|
||||
type: 'thought',
|
||||
content:
|
||||
isCancellation || isRejection
|
||||
? sanitizedError
|
||||
: `Error: ${sanitizedError}`,
|
||||
status:
|
||||
isCancellation || isRejection
|
||||
? SubagentState.CANCELLED
|
||||
: SubagentState.ERROR,
|
||||
});
|
||||
updated = true;
|
||||
break;
|
||||
}
|
||||
default:
|
||||
checkExhaustive(activity.type);
|
||||
break;
|
||||
}
|
||||
|
||||
if (updated) {
|
||||
// Keep only the last N items
|
||||
if (recentActivity.length > MAX_RECENT_ACTIVITY) {
|
||||
recentActivity = recentActivity.slice(-MAX_RECENT_ACTIVITY);
|
||||
}
|
||||
|
||||
const progress: SubagentProgress = {
|
||||
isSubagentProgress: true,
|
||||
agentName: this.definition.name,
|
||||
recentActivity: [...recentActivity], // Copy to avoid mutation issues
|
||||
state: SubagentState.RUNNING,
|
||||
};
|
||||
|
||||
updateOutput(progress);
|
||||
}
|
||||
};
|
||||
|
||||
// Create session with the raw activity callback for rich progress display
|
||||
const session = new LocalSubagentSession(
|
||||
this.definition,
|
||||
this.context,
|
||||
this.messageBus,
|
||||
onActivity,
|
||||
);
|
||||
|
||||
// Subscribe for parent session observability
|
||||
let unsubscribeParent: (() => void) | undefined;
|
||||
if (this._onAgentEvent) {
|
||||
unsubscribeParent = session.subscribe(this._onAgentEvent);
|
||||
}
|
||||
|
||||
// Wire external abort signal to session abort
|
||||
const abortListener = () => void session.abort();
|
||||
signal.addEventListener('abort', abortListener, { once: true });
|
||||
|
||||
try {
|
||||
if (updateOutput) {
|
||||
const initialProgress: SubagentProgress = {
|
||||
isSubagentProgress: true,
|
||||
agentName: this.definition.name,
|
||||
recentActivity: [],
|
||||
state: SubagentState.RUNNING,
|
||||
};
|
||||
updateOutput(initialProgress);
|
||||
}
|
||||
|
||||
// Buffer non-query params, then send query as message to start execution
|
||||
const query = String(this.params['query'] ?? '');
|
||||
const otherParams = { ...this.params } as Record<string, unknown>;
|
||||
delete otherParams['query'];
|
||||
if (Object.keys(otherParams).length > 0) {
|
||||
await session.send({ update: { config: otherParams } });
|
||||
}
|
||||
await session.send({
|
||||
message: { content: [{ type: 'text', text: query }] },
|
||||
});
|
||||
|
||||
const output = await session.getResult();
|
||||
|
||||
if (output.terminate_reason === AgentTerminateMode.ABORTED) {
|
||||
const progress: SubagentProgress = {
|
||||
isSubagentProgress: true,
|
||||
agentName: this.definition.name,
|
||||
recentActivity: [...recentActivity],
|
||||
state: SubagentState.CANCELLED,
|
||||
};
|
||||
|
||||
if (updateOutput) {
|
||||
updateOutput(progress);
|
||||
}
|
||||
|
||||
const cancelError = new Error('Operation cancelled by user');
|
||||
cancelError.name = 'AbortError';
|
||||
throw cancelError;
|
||||
}
|
||||
|
||||
const progress: SubagentProgress = {
|
||||
isSubagentProgress: true,
|
||||
agentName: this.definition.name,
|
||||
recentActivity: [...recentActivity],
|
||||
state: SubagentState.COMPLETED,
|
||||
result: output.result,
|
||||
terminateReason: output.terminate_reason,
|
||||
};
|
||||
|
||||
if (updateOutput) {
|
||||
updateOutput(progress);
|
||||
}
|
||||
|
||||
const resultContent = `Subagent '${this.definition.name}' finished.
|
||||
Termination Reason: ${output.terminate_reason}
|
||||
Result:
|
||||
${output.result}`;
|
||||
|
||||
return {
|
||||
llmContent: [{ text: resultContent }],
|
||||
returnDisplay: progress,
|
||||
};
|
||||
} catch (error) {
|
||||
const errorMessage =
|
||||
error instanceof Error ? error.message : String(error);
|
||||
|
||||
const isAbort =
|
||||
(error instanceof Error && error.name === 'AbortError') ||
|
||||
errorMessage.includes('Aborted');
|
||||
|
||||
// Mark any running items as error/cancelled
|
||||
for (const item of recentActivity) {
|
||||
if (item.status === SubagentState.RUNNING) {
|
||||
item.status = isAbort ? SubagentState.CANCELLED : SubagentState.ERROR;
|
||||
}
|
||||
}
|
||||
|
||||
// Ensure the error is reflected in the recent activity for display
|
||||
if (!isAbort) {
|
||||
const lastActivity = recentActivity[recentActivity.length - 1];
|
||||
if (!lastActivity || lastActivity.status !== SubagentState.ERROR) {
|
||||
recentActivity.push({
|
||||
id: randomUUID(),
|
||||
type: 'thought',
|
||||
content: `Error: ${errorMessage}`,
|
||||
status: SubagentState.ERROR,
|
||||
});
|
||||
if (recentActivity.length > MAX_RECENT_ACTIVITY) {
|
||||
recentActivity = recentActivity.slice(-MAX_RECENT_ACTIVITY);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const progress: SubagentProgress = {
|
||||
isSubagentProgress: true,
|
||||
agentName: this.definition.name,
|
||||
recentActivity: [...recentActivity],
|
||||
state: isAbort ? SubagentState.CANCELLED : SubagentState.ERROR,
|
||||
};
|
||||
|
||||
if (updateOutput) {
|
||||
updateOutput(progress);
|
||||
}
|
||||
|
||||
if (isAbort) {
|
||||
throw error;
|
||||
}
|
||||
|
||||
return {
|
||||
llmContent: `Subagent '${this.definition.name}' failed. Error: ${errorMessage}`,
|
||||
returnDisplay: progress,
|
||||
};
|
||||
} finally {
|
||||
signal.removeEventListener('abort', abortListener);
|
||||
unsubscribeParent?.();
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user