mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-07-12 19:10:55 -07:00
feat(core): add LocalSubagentProtocol behind AgentProtocol (#25302)
This commit is contained in:
Adam Weidman
@@ -0,0 +1,438 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
/**
|
||||
* @fileoverview LocalSubagentProtocol — wraps LocalAgentExecutor behind the
|
||||
* AgentProtocol interface, translating SubagentActivityEvent callbacks into
|
||||
* AgentEvents and exposing the executor result via getResult().
|
||||
*
|
||||
* Pattern mirrors LegacyAgentProtocol, but the loop body runs
|
||||
* LocalAgentExecutor instead of GeminiClient.sendMessageStream().
|
||||
*/
|
||||
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import type { AgentLoopContext } from '../config/agent-loop-context.js';
|
||||
import { AgentSession } from '../agent/agent-session.js';
|
||||
import type {
|
||||
AgentProtocol,
|
||||
AgentSend,
|
||||
AgentEvent,
|
||||
StreamEndReason,
|
||||
Unsubscribe,
|
||||
ContentPart,
|
||||
} from '../agent/types.js';
|
||||
import { LocalAgentExecutor } from './local-executor.js';
|
||||
import {
|
||||
AgentTerminateMode,
|
||||
type LocalAgentDefinition,
|
||||
type AgentInputs,
|
||||
type OutputObject,
|
||||
type SubagentActivityEvent,
|
||||
} from './types.js';
|
||||
import type { MessageBus } from '../confirmation-bus/message-bus.js';
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function isAbortLikeError(err: unknown): boolean {
|
||||
return err instanceof Error && err.name === 'AbortError';
|
||||
}
|
||||
|
||||
function mapTerminateMode(mode: AgentTerminateMode): StreamEndReason {
|
||||
switch (mode) {
|
||||
case AgentTerminateMode.GOAL:
|
||||
return 'completed';
|
||||
case AgentTerminateMode.TIMEOUT:
|
||||
return 'max_time';
|
||||
case AgentTerminateMode.MAX_TURNS:
|
||||
return 'max_turns';
|
||||
case AgentTerminateMode.ABORTED:
|
||||
return 'aborted';
|
||||
case AgentTerminateMode.ERROR:
|
||||
case AgentTerminateMode.ERROR_NO_COMPLETE_TASK_CALL:
|
||||
return 'failed';
|
||||
default: {
|
||||
void (mode satisfies never);
|
||||
return 'failed';
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// LocalSubagentProtocol
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
class LocalSubagentProtocol implements AgentProtocol {
|
||||
private _events: AgentEvent[] = [];
|
||||
private _subscribers = new Set<(event: AgentEvent) => void>();
|
||||
private _streamId: string = randomUUID();
|
||||
private _eventCounter = 0;
|
||||
private _agentStartEmitted = false;
|
||||
private _agentEndEmitted = false;
|
||||
private _activeStreamId: string | undefined;
|
||||
private _abortController = new AbortController();
|
||||
|
||||
// Result promise wiring — re-created per stream in _beginNewStream()
|
||||
private _resultResolve!: (output: OutputObject) => void;
|
||||
private _resultReject!: (err: unknown) => void;
|
||||
private _resultPromise: Promise<OutputObject> | undefined;
|
||||
|
||||
// Buffered config from send({update})
|
||||
private _bufferedConfig: Record<string, unknown> = {};
|
||||
|
||||
constructor(
|
||||
private readonly definition: LocalAgentDefinition,
|
||||
private readonly context: AgentLoopContext,
|
||||
// Required for API parity across protocol constructors (local, remote, legacy)
|
||||
_messageBus: MessageBus,
|
||||
private readonly _rawActivityCallback?: (
|
||||
activity: SubagentActivityEvent,
|
||||
) => void,
|
||||
) {}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// AgentProtocol interface
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
get events(): readonly AgentEvent[] {
|
||||
return this._events;
|
||||
}
|
||||
|
||||
subscribe(callback: (event: AgentEvent) => void): Unsubscribe {
|
||||
this._subscribers.add(callback);
|
||||
return () => {
|
||||
this._subscribers.delete(callback);
|
||||
};
|
||||
}
|
||||
|
||||
async send(payload: AgentSend): Promise<{ streamId: string | null }> {
|
||||
if ('update' in payload && payload.update) {
|
||||
// Buffer config for use when message send arrives
|
||||
if (payload.update.config) {
|
||||
this._bufferedConfig = {
|
||||
...this._bufferedConfig,
|
||||
...payload.update.config,
|
||||
};
|
||||
}
|
||||
return { streamId: null };
|
||||
}
|
||||
|
||||
if ('message' in payload && payload.message) {
|
||||
if (this._activeStreamId) {
|
||||
throw new Error(
|
||||
'LocalSubagentProtocol.send() cannot be called while a stream is active.',
|
||||
);
|
||||
}
|
||||
|
||||
// Extract query text from the message ContentParts
|
||||
const queryText = payload.message.content
|
||||
.filter((p): p is ContentPart & { type: 'text' } => p.type === 'text')
|
||||
.map((p) => p.text)
|
||||
.join('');
|
||||
|
||||
// Only include 'query' in params when the message text is non-empty,
|
||||
// so that callers that pass all fields via update.config are not affected.
|
||||
const params: AgentInputs = {
|
||||
...this._bufferedConfig,
|
||||
...(queryText.length > 0 ? { query: queryText } : {}),
|
||||
};
|
||||
this._bufferedConfig = {};
|
||||
|
||||
this._beginNewStream();
|
||||
const streamId = this._streamId;
|
||||
|
||||
// Schedule execution in a macrotask so send() resolves before agent_start
|
||||
setTimeout(() => {
|
||||
void this._runExecutionInBackground(params);
|
||||
}, 0);
|
||||
|
||||
return { streamId };
|
||||
}
|
||||
|
||||
// action and elicitations are not supported
|
||||
return { streamId: null };
|
||||
}
|
||||
|
||||
async abort(): Promise<void> {
|
||||
this._abortController.abort();
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Protocol-specific: result access
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/**
|
||||
* Resolves when the executor completes, with the raw OutputObject.
|
||||
* Used by LocalSubagentInvocation to build the ToolResult.
|
||||
*/
|
||||
getResult(): Promise<OutputObject> {
|
||||
if (!this._resultPromise) {
|
||||
return Promise.reject(new Error('No active or completed stream'));
|
||||
}
|
||||
return this._resultPromise;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Core: execution
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
private _beginNewStream(): void {
|
||||
this._streamId = randomUUID();
|
||||
this._eventCounter = 0;
|
||||
this._abortController = new AbortController();
|
||||
this._agentStartEmitted = false;
|
||||
this._agentEndEmitted = false;
|
||||
this._activeStreamId = this._streamId;
|
||||
this._resultPromise = new Promise<OutputObject>((resolve, reject) => {
|
||||
this._resultResolve = resolve;
|
||||
this._resultReject = reject;
|
||||
});
|
||||
}
|
||||
|
||||
private async _runExecutionInBackground(params: AgentInputs): Promise<void> {
|
||||
this._ensureAgentStart();
|
||||
try {
|
||||
await this._runExecution(params);
|
||||
} catch (err: unknown) {
|
||||
if (this._abortController.signal.aborted || isAbortLikeError(err)) {
|
||||
this._ensureAgentEnd('aborted');
|
||||
// Abort resolves with an empty result — partial output is intentionally
|
||||
// dropped since the caller requested cancellation.
|
||||
this._resultResolve({
|
||||
result: '',
|
||||
terminate_reason: AgentTerminateMode.ABORTED,
|
||||
});
|
||||
} else {
|
||||
this._emitErrorAndAgentEnd(err);
|
||||
this._resultReject(err);
|
||||
}
|
||||
} finally {
|
||||
this._clearActiveStream();
|
||||
}
|
||||
}
|
||||
|
||||
private async _runExecution(params: AgentInputs): Promise<void> {
|
||||
const signal = this._abortController.signal;
|
||||
|
||||
const onActivity = (activity: SubagentActivityEvent): void => {
|
||||
// Forward raw activity to invocation-level callback (for rich SubagentProgress display)
|
||||
this._rawActivityCallback?.(activity);
|
||||
this._emit(this._translateActivity(activity));
|
||||
};
|
||||
|
||||
const executor = await LocalAgentExecutor.create(
|
||||
this.definition,
|
||||
this.context,
|
||||
onActivity,
|
||||
);
|
||||
|
||||
const output = await executor.run(params, signal);
|
||||
|
||||
if (
|
||||
output.terminate_reason === AgentTerminateMode.ABORTED ||
|
||||
signal.aborted
|
||||
) {
|
||||
this._finishStream('aborted');
|
||||
} else {
|
||||
this._finishStream(mapTerminateMode(output.terminate_reason));
|
||||
}
|
||||
|
||||
this._resultResolve(output);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Activity → AgentEvent translation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
private _translateActivity(activity: SubagentActivityEvent): AgentEvent[] {
|
||||
switch (activity.type) {
|
||||
case 'THOUGHT_CHUNK': {
|
||||
const rawText = activity.data['text'];
|
||||
const text = String(rawText ?? '');
|
||||
return [
|
||||
this._makeEvent('message', {
|
||||
role: 'agent',
|
||||
content: [{ type: 'thought', thought: text }],
|
||||
}),
|
||||
];
|
||||
}
|
||||
case 'TOOL_CALL_START': {
|
||||
const rawCallId = activity.data['callId'];
|
||||
const callId = String(rawCallId ?? randomUUID());
|
||||
const rawName = activity.data['name'];
|
||||
const name = String(rawName ?? 'unknown');
|
||||
const rawArgs = activity.data['args'];
|
||||
const args: Record<string, unknown> =
|
||||
rawArgs !== null &&
|
||||
typeof rawArgs === 'object' &&
|
||||
!Array.isArray(rawArgs)
|
||||
? // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
(rawArgs as Record<string, unknown>)
|
||||
: {};
|
||||
return [
|
||||
this._makeEvent('tool_request', {
|
||||
requestId: callId,
|
||||
name,
|
||||
args,
|
||||
}),
|
||||
];
|
||||
}
|
||||
case 'TOOL_CALL_END': {
|
||||
const rawId = activity.data['id'];
|
||||
const requestId = String(rawId ?? randomUUID());
|
||||
const rawName = activity.data['name'];
|
||||
const name = String(rawName ?? 'unknown');
|
||||
const rawOutput = activity.data['output'];
|
||||
const output = String(rawOutput ?? '');
|
||||
return [
|
||||
this._makeEvent('tool_response', {
|
||||
requestId,
|
||||
name,
|
||||
content: [{ type: 'text', text: output }],
|
||||
}),
|
||||
];
|
||||
}
|
||||
case 'ERROR': {
|
||||
const rawError = activity.data['error'];
|
||||
const errorMsg = String(rawError ?? 'Unknown error');
|
||||
return [
|
||||
this._makeEvent('error', {
|
||||
status: 'INTERNAL',
|
||||
message: errorMsg,
|
||||
fatal: false,
|
||||
}),
|
||||
];
|
||||
}
|
||||
default: {
|
||||
void (activity.type satisfies never);
|
||||
return [];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Internal helpers (mirrors LegacyAgentProtocol)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
private _emit(events: AgentEvent[]): void {
|
||||
if (events.length === 0) return;
|
||||
const subscribers = [...this._subscribers];
|
||||
for (const event of events) {
|
||||
this._events.push(event);
|
||||
if (event.type === 'agent_end') {
|
||||
this._agentEndEmitted = true;
|
||||
}
|
||||
for (const sub of subscribers) {
|
||||
sub(event);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private _clearActiveStream(): void {
|
||||
this._activeStreamId = undefined;
|
||||
}
|
||||
|
||||
private _ensureAgentStart(): void {
|
||||
if (!this._agentStartEmitted) {
|
||||
this._agentStartEmitted = true;
|
||||
this._emit([this._makeEvent('agent_start', {})]);
|
||||
}
|
||||
}
|
||||
|
||||
private _ensureAgentEnd(reason: StreamEndReason = 'completed'): void {
|
||||
if (!this._agentEndEmitted && this._agentStartEmitted) {
|
||||
this._emit([this._makeEvent('agent_end', { reason })]);
|
||||
}
|
||||
}
|
||||
|
||||
private _finishStream(reason: StreamEndReason): void {
|
||||
this._ensureAgentEnd(reason);
|
||||
this._clearActiveStream();
|
||||
}
|
||||
|
||||
private _emitErrorAndAgentEnd(err: unknown): void {
|
||||
const message = err instanceof Error ? err.message : String(err);
|
||||
this._ensureAgentStart();
|
||||
|
||||
const meta: Record<string, unknown> = {};
|
||||
if (err instanceof Error) {
|
||||
meta['errorName'] = err.constructor.name;
|
||||
meta['stack'] = err.stack;
|
||||
if ('exitCode' in err && typeof err.exitCode === 'number') {
|
||||
meta['exitCode'] = err.exitCode;
|
||||
}
|
||||
if ('code' in err) {
|
||||
meta['code'] = err.code;
|
||||
}
|
||||
if ('status' in err) {
|
||||
meta['status'] = err.status;
|
||||
}
|
||||
}
|
||||
|
||||
this._emit([
|
||||
this._makeEvent('error', {
|
||||
status: 'INTERNAL',
|
||||
message,
|
||||
fatal: true,
|
||||
...(Object.keys(meta).length > 0 ? { _meta: meta } : {}),
|
||||
}),
|
||||
]);
|
||||
this._ensureAgentEnd('failed');
|
||||
}
|
||||
|
||||
private _nextEventFields() {
|
||||
return {
|
||||
id: `${this._streamId}-${this._eventCounter++}`,
|
||||
timestamp: new Date().toISOString(),
|
||||
streamId: this._streamId,
|
||||
};
|
||||
}
|
||||
|
||||
private _makeEvent<T extends AgentEvent['type']>(
|
||||
type: T,
|
||||
payload: Omit<AgentEvent<T>, 'id' | 'timestamp' | 'streamId' | 'type'>,
|
||||
): AgentEvent {
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
return {
|
||||
...this._nextEventFields(),
|
||||
type,
|
||||
...payload,
|
||||
} as AgentEvent;
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Public export
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export class LocalSubagentSession extends AgentSession {
|
||||
private readonly _localProtocol: LocalSubagentProtocol;
|
||||
|
||||
constructor(
|
||||
definition: LocalAgentDefinition,
|
||||
context: AgentLoopContext,
|
||||
messageBus: MessageBus,
|
||||
rawActivityCallback?: (activity: SubagentActivityEvent) => void,
|
||||
) {
|
||||
const protocol = new LocalSubagentProtocol(
|
||||
definition,
|
||||
context,
|
||||
messageBus,
|
||||
rawActivityCallback,
|
||||
);
|
||||
super(protocol);
|
||||
this._localProtocol = protocol;
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns the raw executor OutputObject once execution completes.
|
||||
* Used by LocalSubagentInvocation to build the ToolResult.
|
||||
*/
|
||||
getResult(): Promise<OutputObject> {
|
||||
return this._localProtocol.getResult();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user